Star378
269
代码
代码
Issues14
Pull Requests50
Actions
项目讨论
Wiki
分析
Star378
269
已合并
megamoe优化，改成主E模式 #1466
zy_tt创建于 25 天前
megamoe优化，改成主E模式 #1466
已合并
zy_tt创建于 25 天前
从zy_tt/pto-isa_cann_zymaster1合入到cann/pto-isamaster
共 30 个文件变更+2726-1759
Rdocs/figures/megamoe/permute_all2all_count.png→docs/figures/megamoe/permuate_all2all_count.png+0-0
文件重命名但无更改。
Mkernels/manual/a2a3/dispatch_mega_combine/CMakeLists.txt+66-50
@@ -11,13 +11,17 @@
11 11 
12cmake_minimum_required(VERSION 3.16)12cmake_minimum_required(VERSION 3.16)
13 13 
14-set(CMAKE_COMPILER "bisheng" CACHE STRING "Compiler used for host and CCE sources")14+if(NOT DEFINED ENV{ASCEND_HOME_PATH})
15-if(NOT DEFINED CMAKE_C_COMPILER)15+ message(FATAL_ERROR "Cannot find ASCEND_HOME_PATH, please run set_env.sh.")
16- set(CMAKE_C_COMPILER "${CMAKE_COMPILER}")
17endif()16endif()
18-if(NOT DEFINED CMAKE_CXX_COMPILER)17+set(ASCEND_HOME_PATH $ENV{ASCEND_HOME_PATH})
19- set(CMAKE_CXX_COMPILER "${CMAKE_COMPILER}")18+ 
19+get_filename_component(ASCEND_BISHENG_COMPILER "${ASCEND_HOME_PATH}/bin/bisheng" REALPATH)
20+if(NOT EXISTS "${ASCEND_BISHENG_COMPILER}")
21+ message(FATAL_ERROR "Cannot find ${ASCEND_HOME_PATH}/bin/bisheng, please run the CANN set_env.sh.")
20endif()22endif()
23+set(CMAKE_C_COMPILER "${ASCEND_BISHENG_COMPILER}" CACHE FILEPATH "C compiler" FORCE)
24+set(CMAKE_CXX_COMPILER "${ASCEND_BISHENG_COMPILER}" CACHE FILEPATH "CXX compiler" FORCE)
21 25 
22project(pto_example LANGUAGES C CXX)26project(pto_example LANGUAGES C CXX)
23 27 
@@ -31,14 +35,63 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
31set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})35set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
32set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)36set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
33 37 
34-if(NOT DEFINED ENV{ASCEND_HOME_PATH})
35- message(FATAL_ERROR "Cannot find ASCEND_HOME_PATH, please run set_env.sh.")
36-else()
37- set(ASCEND_HOME_PATH $ENV{ASCEND_HOME_PATH})
38-endif()
39- 
40set(ASCEND_DRIVER_PATH /usr/local/Ascend/driver)38set(ASCEND_DRIVER_PATH /usr/local/Ascend/driver)
41 39 
40+# The CANN environment exposes include/bin links for its active ARM or x86 package.
41+# Resolve that link instead of duplicating CANN's platform selection logic here.
42+get_filename_component(ASCEND_KERNEL_PLATFORM_INCLUDE "${ASCEND_HOME_PATH}/include" REALPATH)
43+get_filename_component(ASCEND_KERNEL_PLATFORM_ROOT "${ASCEND_KERNEL_PLATFORM_INCLUDE}" DIRECTORY)
44+if(NOT IS_DIRECTORY "${ASCEND_KERNEL_PLATFORM_INCLUDE}")
45+ message(FATAL_ERROR "Cannot resolve the active CANN platform directory from ${ASCEND_HOME_PATH}/include.")
46+endif()
47+ 
48+# CANN 9.x uses asc/, while older compatibility packages use include/ascendc/.
49+set(KERNEL_INCLUDE_CANDIDATES
50+ ${ASCEND_KERNEL_PLATFORM_ROOT}/asc
51+ ${ASCEND_KERNEL_PLATFORM_ROOT}/asc/include
52+ ${ASCEND_KERNEL_PLATFORM_ROOT}/asc/include/basic_api
53+ ${ASCEND_KERNEL_PLATFORM_ROOT}/asc/include/interface
54+ ${ASCEND_KERNEL_PLATFORM_ROOT}/asc/include/utils
55+ ${ASCEND_KERNEL_PLATFORM_ROOT}/asc/impl/basic_api
56+ ${ASCEND_KERNEL_PLATFORM_ROOT}/include/ascendc/basic_api
57+ ${ASCEND_KERNEL_PLATFORM_ROOT}/include/ascendc/basic_api/interface
58+ ${ASCEND_KERNEL_PLATFORM_ROOT}/include/ascendc/basic_api/impl
59+)
60+ 
61+set(KERNEL_INCLUDE_DIRS)
62+set(KERNEL_OPERATOR_HEADER_FOUND FALSE)
63+set(KERNEL_TPIPE_HEADER_FOUND FALSE)
64+set(KERNEL_TPIPE_IMPL_HEADER_FOUND FALSE)
65+foreach(KERNEL_INCLUDE_CANDIDATE IN LISTS KERNEL_INCLUDE_CANDIDATES)
66+ if(IS_DIRECTORY "${KERNEL_INCLUDE_CANDIDATE}")
67+ list(APPEND KERNEL_INCLUDE_DIRS "${KERNEL_INCLUDE_CANDIDATE}")
68+ if(EXISTS "${KERNEL_INCLUDE_CANDIDATE}/kernel_operator.h")
69+ set(KERNEL_OPERATOR_HEADER_FOUND TRUE)
70+ endif()
71+ if(EXISTS "${KERNEL_INCLUDE_CANDIDATE}/kernel_tpipe.h")
72+ set(KERNEL_TPIPE_HEADER_FOUND TRUE)
73+ endif()
74+ if(EXISTS "${KERNEL_INCLUDE_CANDIDATE}/kernel_tpipe_impl.h")
75+ set(KERNEL_TPIPE_IMPL_HEADER_FOUND TRUE)
76+ endif()
77+ endif()
78+endforeach()
79+ 
80+if(NOT KERNEL_OPERATOR_HEADER_FOUND
81+ OR NOT KERNEL_TPIPE_HEADER_FOUND
82+ OR NOT KERNEL_TPIPE_IMPL_HEADER_FOUND)
83+ message(FATAL_ERROR
84+ "Incomplete AscendC headers under ASCEND_HOME_PATH=${ASCEND_HOME_PATH}: expected kernel_operator.h, "
85+ "kernel_tpipe.h, and kernel_tpipe_impl.h. Check the active CANN environment.")
86+endif()
87+ 
88+if(EXISTS ${ASCEND_KERNEL_PLATFORM_ROOT}/include/ascendc/highlevel_api/lib/std/tuple.h)
89+ list(APPEND KERNEL_INCLUDE_DIRS ${ASCEND_KERNEL_PLATFORM_ROOT}/include/ascendc/highlevel_api)
90+endif()
91+if(EXISTS ${ASCEND_KERNEL_PLATFORM_ROOT}/pkg_inc/runtime/runtime)
92+ list(APPEND KERNEL_INCLUDE_DIRS ${ASCEND_KERNEL_PLATFORM_ROOT}/pkg_inc/runtime/runtime)
93+endif()
94+list(REMOVE_DUPLICATES KERNEL_INCLUDE_DIRS)
42add_compile_options(95add_compile_options(
43 -D_FORTIFY_SOURCE=296 -D_FORTIFY_SOURCE=2
44 -O2 -std=c++1797 -O2 -std=c++17
@@ -62,39 +115,7 @@ set(CMAKE_CCE_COMPILE_OPTIONS
62 "SHELL:-mllvm -cce-aicore-dcci-insert-for-scalar=false"115 "SHELL:-mllvm -cce-aicore-dcci-insert-for-scalar=false"
63)116)
64 117 
65-if(DEBUG_MODE)118+set(CCE_AICORE_ARCH "dav-c220")
66- message(STATUS "Debug Mode Enabled, Add Debug Options")
67- add_compile_definitions(_DEBUG)
68- set(CMAKE_CCE_COMPILE_OPTIONS "${CMAKE_CCE_COMPILE_OPTIONS} --cce-enable-print")
69-endif()
70- 
71-set(A2A3_SOC_VERSIONS
72- Ascend910B1
73- Ascend910B
74- Ascend910C
75- Ascend910_93
76- Ascend910_9391
77- Ascend910_9381
78- Ascend910_9372
79- Ascend910_9392
80- Ascend910_9382
81- Ascend910_9362
82- ascend910_93
83- ascend910_9391
84- ascend910_9381
85- ascend910_9372
86- ascend910_9392
87- ascend910_9382
88- ascend910_9362
89-)
90- 
91-if("${SOC_VERSION}" IN_LIST A2A3_SOC_VERSIONS)
92- set(CCE_AICORE_ARCH "dav-c220")
93-elseif("${SOC_VERSION}" STREQUAL "Ascend910_9599")
94- set(CCE_AICORE_ARCH "dav-c310")
95-else()
96- message(FATAL_ERROR "Unsupported SOC_VERSION: ${SOC_VERSION}")
97-endif()
98 119 
99set(CMAKE_CPP_COMPILE_OPTIONS120set(CMAKE_CPP_COMPILE_OPTIONS
100 -xc++121 -xc++
@@ -124,12 +145,7 @@ function(pto_example_mixed NAME)
124 ${PROJECT_SOURCE_DIR}/op_kernel145 ${PROJECT_SOURCE_DIR}/op_kernel
125 ${PROJECT_SOURCE_DIR}/op_kernel/utils146 ${PROJECT_SOURCE_DIR}/op_kernel/utils
126 ${PROJECT_BINARY_DIR}147 ${PROJECT_BINARY_DIR}
127- ${ASCEND_HOME_PATH}/aarch64-linux/pkg_inc/runtime/runtime148+ ${KERNEL_INCLUDE_DIRS}
128- ${ASCEND_HOME_PATH}/aarch64-linux/include/ascendc/basic_api
129- ${ASCEND_HOME_PATH}/aarch64-linux/asc
130- ${ASCEND_HOME_PATH}/aarch64-linux/asc/include/basic_api
131- ${ASCEND_HOME_PATH}/aarch64-linux/asc/include/interface
132- ${ASCEND_HOME_PATH}/aarch64-linux/ascendc/include/basic_api/impl
133 )149 )
134 target_link_options(${NAME}_kernel PRIVATE --cce-fatobj-link)150 target_link_options(${NAME}_kernel PRIVATE --cce-fatobj-link)
135 151 
Mkernels/manual/a2a3/dispatch_mega_combine/README.md+32-39
@@ -48,14 +48,10 @@ kernels/manual/a2a3/dispatch_mega_combine/
48│ ├── combine.h # Remote writeback of GMM2 output to offsetD48│ ├── combine.h # Remote writeback of GMM2 output to offsetD
49│ ├── unpermute.h # TopK weighted reduction and original token order restoration49│ ├── unpermute.h # TopK weighted reduction and original token order restoration
50│ └── utils/ # PTO vector, sync, HCCL window, and GMM pipeline helpers50│ └── utils/ # PTO vector, sync, HCCL window, and GMM pipeline helpers
51-├── overview.md # Design overview, performance comparison, and stage pseudocode51+├── overview.md # Design overview and performance comparison
52-├── front_reorder.md # Front reorder / sort / count-as-flag details52+├── overview_v1.md # Fixed-group scheduling and overlap design
53-├── dispatch.md # Dispatch contract and data movement strategy53+├── pseudocode.md # Current seven-stage data flow and pseudocode
54-├── gmm1.md / gmm2.md # GMM tile scheduling, swizzle, sync, and pipeline details54+└── README_zh.md # Chinese README
55-├── swiglu.md # SwiGLU segmentation and quantization strategy
56-├── combine.md # Combine large/small paths and remote writeback protocol
57-├── unpermute.md # Unpermute restoration and accumulation strategy
58-└── glden.md # Python batch golden rewrite design
59```55```
60 56 
61## Operator Description57## Operator Description
@@ -97,12 +93,12 @@ for each rank, token:
97 93 
98## Optimization Notes94## Optimization Notes
99 95 
100-- **Expert-level overlap**: AIC-side GMM1/GMM2 and AIV-side Dispatch/SwiGLU/Combine progress by local expert group, connected by hard flags.96+- **Expert-level overlap**: AIC-side GMM1/GMM2 and AIV-side Dispatch/SwiGLU/Combine progress by local expert, connected by per-expert arrival/ready notifications.
101- **Three front reorder paths**: FullLoad, OneCore, and MultiCore are selected by UB working-set size. Small-route cases keep sort/count/inverse/quant work in UB as much as possible.97- **Three front reorder paths**: FullLoad, OneCore, and MultiCore are selected by UB working-set size. Small-route cases keep sort/count/inverse/quant work in UB as much as possible.
102- **Count-as-flag**: FrontReorder publishes count rows with a marker value so peer ranks can wait on data arrival directly instead of adding a full count-exchange barrier.98- **Count-as-flag**: FrontReorder publishes count rows with a marker value so peer ranks can wait on data arrival directly instead of adding a full count-exchange barrier.
103- **PTO tile GMM optimization**: GMM1/GMM2 use output-tile swizzle, L1-to-L0 multi-level reuse, double buffering, and fixpipe quant/cast.99- **PTO tile GMM optimization**: GMM1/GMM2 use output-tile swizzle, L1-to-L0 multi-level reuse, double buffering, and fixpipe quant/cast.
104-- **Segmented SwiGLU overlap**: SwiGLU is split into segments so the first segment can finish before GMM2 starts.100+- **Dynamic AIC grouping**: GMM1 starts with all 24 AICs and then releases 8; GMM2 starts with 8 AICs and expands to 24 after GMM1 completes.
105-- **Dual combine paths**: large-token cases use full-row writeback, while small-token cases split GMM2 tiles into subtiles to improve AIV occupancy.101+- **Two-stage Unpermute**: 32 AIVs first process tokens whose routes are ready, then all 48 AIVs process the remaining tokens after Combine completes.
106 102 
107## Tiling Parameters103## Tiling Parameters
108 104 
@@ -115,11 +111,11 @@ for each rank, token:
115| `expertPerRank` | Number of local experts per rank |111| `expertPerRank` | Number of local experts per rank |
116| `worldSize` | MPI/HCCL rank count |112| `worldSize` | MPI/HCCL rank count |
117| `maxOutputSize` | Per-rank routed-row workspace limit; typical performance cases pass a fixed workspace limit explicitly |113| `maxOutputSize` | Per-rank routed-row workspace limit; typical performance cases pass a fixed workspace limit explicitly |
118-| `aicNum` | Logical AIC count; default script value is 24 |114+| `aicNum` | Fixed physical AIC count for the production path: 24 |
119-| `aivNum` | Logical AIV count; default script value is 48 |115+| `aivNum` | Fixed physical AIV count for the production path: 48 |
120| `GMM baseM/baseN` | Main output tile shape is `128 x 256` |116| `GMM baseM/baseN` | Main output tile shape is `128 x 256` |
121| `Front FullLoad` | Selected by `routeElems`, `K`, `expertNum`, and the 192 KiB UB budget |117| `Front FullLoad` | Selected by `routeElems`, `K`, `expertNum`, and the 192 KiB UB budget |
122-| `Combine small path` | Preferred when `problemM * topK <= 4096` |118+| Fixed groups | GMM1/GMM2 use `16 + 8` AICs; Dispatch/SwiGLU/Combine have `16/16/8` AIV slots, with SwiGLU activating 8 or 16 by M |
123 119 
124## Supported Cases120## Supported Cases
125 121 
@@ -152,19 +148,19 @@ aivNum=48
152│ FrontReorder (AIV) │148│ FrontReorder (AIV) │
153│ sort expertId routes -> offsetA + count/prefix metadata │149│ sort expertId routes -> offsetA + count/prefix metadata │
154└──────────────────────────────┬───────────────────────────────────────────────┘150└──────────────────────────────┬───────────────────────────────────────────────┘
155- │ D2C ready / count-as-flag151+ │ count metadata ready / count-as-flag
156┌──────────────────────────────▼───────────────────────────────────────────────┐152┌──────────────────────────────▼───────────────────────────────────────────────┐
157│ Expert-level overlapped pipeline │153│ Expert-level overlapped pipeline │
158│ │154│ │
159-│ AIV: Dispatch(group i) -> SwiGLU(segment/group i) -> Combine(group i) │155+│ AIV: Dispatch(expert i) -> SwiGLU(expert i) -> Combine(expert i) │
160-│ AIC: GMM1(group i) -> GMM2(group i) │156+│ AIC: GMM1 starts 24 -> 16; released AICs join GMM2, which grows 8 -> 24 │
161│ │157│ │
162-│ Stages communicate through hard flags: D2C, C2V, V2C, G2C/Combine ready │158+│ Stages communicate through per-expert GM arrival/ready progress │
163└──────────────────────────────┬───────────────────────────────────────────────┘159└──────────────────────────────┬───────────────────────────────────────────────┘
164- │ final boundary160+ │ per-rank expert progress / DataReady
165┌──────────────────────────────▼───────────────────────────────────────────────┐161┌──────────────────────────────▼───────────────────────────────────────────────┐
166-│ Unpermute (AIV) │162+│ Two-stage Unpermute (AIV) │
167-│ offsetD + probs + expandedRowIdx -> TopK weighted reduce -> out[M, K] │163+│ 32 AIVs process ready tokens -> 48 AIVs process remaining tokens -> out │
168└──────────────────────────────────────────────────────────────────────────────┘164└──────────────────────────────────────────────────────────────────────────────┘
169```165```
170 166 
@@ -211,7 +207,7 @@ srcRowBase = preSumBeforeRank[srcRank, localExpert]
211dstRowBase = groupBase + (srcRank == 0 ? 0 : cumsumMM[srcRank - 1, localExpert])207dstRowBase = groupBase + (srcRank == 0 ? 0 : cumsumMM[srcRank - 1, localExpert])
212```208```
213 209 
214-After each local expert group is gathered, Dispatch sets the GMM1-ready flag for that group.210+After each local expert is gathered, Dispatch coordinates its active workers and publishes GMM1 ready progress.
215 211 
216## GMM1 / SwiGLU / GMM2 Stages212## GMM1 / SwiGLU / GMM2 Stages
217 213 
@@ -226,7 +222,7 @@ gmA[int8] x weight1[int8]
226 -> gmC[half]222 -> gmC[half]
227```223```
228 224 
229-Each local expert is split into `128 x 256` output tiles. Linear tile ids are mapped to `(blockM, blockN)` with swizzle to improve L1 reuse of the B-side weights.225+Each local expert is split into `128 x 256` output tiles. Linear tile ids are mapped to `(blockM, blockN)` with swizzle to improve L1 reuse of the B-side weights. The first two experts use all 24 AICs; subsequent experts use the 16-AIC GMM1 group.
230 226 
231### SwiGLU227### SwiGLU
232 228 
@@ -239,7 +235,7 @@ gmC * perTokenScale1
239 -> gmPermutedToken[int8] + perTokenScale2[float]235 -> gmPermutedToken[int8] + perTokenScale2[float]
240```236```
241 237 
242-SwiGLU is split into segments. Core0 writes segment metadata and the other AIVs split rows from that metadata.238+SwiGLU advances one expert at a time and splits its rows across the active fixed group. After all active AIVs finish an expert, its coordinator publishes GMM2 ready progress. M=16 uses 8 AIVs; the other listed cases use 16.
243 239 
244### GMM2240### GMM2
245 241 
@@ -252,7 +248,7 @@ gmPermutedToken[int8] x weight2[int8]
252 -> gmm2Output[half]248 -> gmm2Output[half]
253```249```
254 250 
255-After GMM2 finishes a local expert group, it sets the combine-ready flag for that group.251+GMM2 starts on the 8-AIC group. Once GMM1 is done, its released 16 AICs join GMM2 at an expert boundary; each completed expert publishes arrival progress for Combine.
256 252 
257## Combine / Unpermute Stages253## Combine / Unpermute Stages
258 254 
@@ -266,16 +262,14 @@ gmm2Output[srcRow, 0:K] half
266 -> srcRank.remoteWindow.offsetD[dstRow, 0:K]262 -> srcRank.remoteWindow.offsetD[dstRow, 0:K]
267```263```
268 264 
269-Path selection:265+The 8-AIV Combine group assigns full-row writeback by source rank. It publishes per-rank expert progress while running and publishes final DataReady after all experts are visible.
270- 
271-- **DirectLarge**: full-row writeback for large-token cases.
272-- **DirectSmall**: subtile writeback for small-token cases to improve AIV occupancy.
273 266 
274Unpermute restores the source-rank token order:267Unpermute restores the source-rank token order:
275 268 
276```text269```text
277offsetD + probs + expandedRowIdx270offsetD + probs + expandedRowIdx
278- -> TopK weighted accumulation271+ -> phase 1: 32 AIVs reduce tokens whose routes are ready
272+ -> phase 2: 48 AIVs reduce the remaining tokens
279 -> out[M, K]273 -> out[M, K]
280```274```
281 275 
@@ -288,6 +282,7 @@ The HCCL remote window carries cross-rank visible data:
288| `offsetA` | HCCL window | FrontReorder writes packed int8 token rows; Dispatch pulls from peer ranks |282| `offsetA` | HCCL window | FrontReorder writes packed int8 token rows; Dispatch pulls from peer ranks |
289| `offsetD` | HCCL window | Combine writes back to source ranks; Unpermute consumes locally |283| `offsetD` | HCCL window | Combine writes back to source ranks; Unpermute consumes locally |
290| `tokenPerExpert` | HCCL window | count-as-flag cross-rank count rows |284| `tokenPerExpert` | HCCL window | count-as-flag cross-rank count rows |
285+| `ExpertProgress / DataReady` | HCCL window | Per-rank Combine progress and launch-scoped completion notification |
291| `gmA` | workspace GM | GMM1 input generated by Dispatch |286| `gmA` | workspace GM | GMM1 input generated by Dispatch |
292| `gmC` | workspace GM | GMM1 output and SwiGLU input |287| `gmC` | workspace GM | GMM1 output and SwiGLU input |
293| `gmPermutedToken` | workspace GM | SwiGLU dynamic-quant output and GMM2 input |288| `gmPermutedToken` | workspace GM | SwiGLU dynamic-quant output and GMM2 input |
@@ -321,7 +316,7 @@ Small route counts should prefer FullLoad, which reduces GM intermediate traffic
321 316 
322### 2. Keep Expert-Level Stage Boundaries Clear317### 2. Keep Expert-Level Stage Boundaries Clear
323 318 
324-Dispatch, GMM1, SwiGLU, GMM2, and Combine depend on hard flags. Before removing any synchronization, verify the exact per-group set/wait relationship.319+Dispatch, GMM1, SwiGLU, GMM2, and Combine use producer arrival and consumer ready slots. Before changing synchronization, verify the producer count, consumer range, and expected expert progress at each boundary.
325 320 
326### 3. Prioritize GMM Tile Efficiency321### 3. Prioritize GMM Tile Efficiency
327 322 
@@ -332,9 +327,9 @@ GMM1/GMM2 dominate runtime. Check:
332- whether small `currentM` causes AIC imbalance;327- whether small `currentM` causes AIC imbalance;
333- whether AIV communication/writeback competes with GMM HBM traffic.328- whether AIV communication/writeback competes with GMM HBM traffic.
334 329 
335-### 4. Use Subtile Combine for Small Token Counts330+### 4. Balance Combine Against GMM2
336 331 
337-For small M, the direct row path may underuse AIVs. DirectSmall splits GMM2 tiles into subtiles, but requires the GMM2 tiling column width to align with the small-path subtile width.332+Combine uses 8 AIVs and a shape-specific delayed start to avoid competing with GMM2 for HBM bandwidth. When tuning small M, check the GMM2 join point, Combine start expert, and SwiGLU active-worker count together.
338 333 
339### 5. Use Batch Golden Generation334### 5. Use Batch Golden Generation
340 335 
@@ -345,9 +340,7 @@ Large synthetic cases use the `python-batch` golden backend by default. Use `pyt
345Configure the Ascend CANN environment:340Configure the Ascend CANN environment:
346 341 
347```bash342```bash
348-export ASCEND_CANN_PATH=/usr/local/Ascend/cann/set_env.sh343+source <cann-install>/set_env.sh
349-export ASCEND_HOME_PATH=/usr/local/Ascend/cann/cann
350-source /usr/local/Ascend/cann/cann/set_env.sh
351```344```
352 345 
353Run the default 2048 case:346Run the default 2048 case:
@@ -369,13 +362,13 @@ bash run.sh --world-size 8 --m 512 --k 7168 --n 4096 --topk 8 --experts 16 --max
369 362 
370| Environment Variable | Purpose | Default Behavior |363| Environment Variable | Purpose | Default Behavior |
371| --- | --- | --- |364| --- | --- | --- |
372-| `ASCEND_HOME_PATH` | CANN installation path | Must be set before running |365+| `ASCEND_HOME_PATH` | Active CANN installation path | Set by the CANN `set_env.sh` |
373-| `CMAKE_COMPILER` | Compiler used by CMake | `bisheng` |
374| `MPI_ENV_BIN` | MPI/conda `bin` path | `/home/ntlab/miniconda3/envs/ltr_pto/bin` |366| `MPI_ENV_BIN` | MPI/conda `bin` path | `/home/ntlab/miniconda3/envs/ltr_pto/bin` |
375| `MPI_ENV_LIB` | MPI/conda `lib` path | `/home/ntlab/miniconda3/envs/ltr_pto/lib` |367| `MPI_ENV_LIB` | MPI/conda `lib` path | `/home/ntlab/miniconda3/envs/ltr_pto/lib` |
376| `MPI_LIB_PATH` | Absolute path to `libmpi.so` | `${MPI_ENV_LIB}/libmpi.so` |368| `MPI_LIB_PATH` | Absolute path to `libmpi.so` | `${MPI_ENV_LIB}/libmpi.so` |
377| `MPI_RUNNER` | MPI launch command | `mpirun` |369| `MPI_RUNNER` | MPI launch command | `mpirun` |
378| `HCCL_BUFFSIZE` | HCCL RDMA window size | Raised automatically by `run.sh` when needed |370| `HCCL_BUFFSIZE` | HCCL RDMA window size | Raised automatically by `run.sh` when needed |
371+| `DISPATCH_MEGA_COMBINE_START_SYNC_DEBUG` | Synchronize kernel entry for cross-core timing comparison | Disabled |
379 372 
380## Changing Case Parameters373## Changing Case Parameters
381 374 
@@ -400,7 +393,7 @@ Common constraints:
400| HCCL window too small | The manually set `HCCL_BUFFSIZE` is below the case requirement; unset it or increase it |393| HCCL window too small | The manually set `HCCL_BUFFSIZE` is below the case requirement; unset it or increase it |
401| MPI launch fails | Check that `MPI_ENV_BIN`, `MPI_ENV_LIB`, and `MPI_LIB_PATH` point to the same conda/MPI environment |394| MPI launch fails | Check that `MPI_ENV_BIN`, `MPI_ENV_LIB`, and `MPI_LIB_PATH` point to the same conda/MPI environment |
402| Golden generation is slow | Use the default `python-batch` backend; use `python-naive` only for debug comparison |395| Golden generation is slow | Use the default `python-batch` backend; use `python-naive` only for debug comparison |
403-| Small-M performance is unstable | Check whether FullLoad is selected, whether Combine uses DirectSmall, and whether AIV concurrency is hurting GMM |396+| Small-M performance is unstable | Check whether FullLoad is selected and whether the configured SwiGLU workers, GMM2 join point, and Combine start point increase HBM contention |
404| Result diff is abnormal | Check whether old generated data was reused; do not reuse stale `out/` after changing expert distribution or key case parameters |397| Result diff is abnormal | Check whether old generated data was reused; do not reuse stale `out/` after changing expert distribution or key case parameters |
405 398 
406## Build System399## Build System
Mkernels/manual/a2a3/dispatch_mega_combine/README_zh.md+35-43
@@ -10,7 +10,7 @@
10- Ascend910B / Ascend910C10- Ascend910B / Ascend910C
11- Ascend910_93 / Ascend910_9391 / Ascend910_9381 / Ascend910_9372 / Ascend910_9392 / Ascend910_9382 / Ascend910_936211- Ascend910_93 / Ascend910_9391 / Ascend910_9381 / Ascend910_9372 / Ascend910_9392 / Ascend910_9382 / Ascend910_9362
12 12 
13-> 当前目录位于 `a2a3` 手写 kernel 下，性能和运行验证主要面向 A2/A3 形态。`CMakeLists.txt` 中也保留了 `Ascend910_9599` 的 `dav-c310` 编译分支，但使用前需要按目标环境重新验证。13+当前目录位于 `a2a3` 手写 kernel 下。典型性能 case 使用脚本默认的 A2/A3 配置，命令中无需显式指定 SoC。
14 14 
15## 目录结构15## 目录结构
16 16 
@@ -40,14 +40,10 @@ kernels/manual/a2a3/dispatch_mega_combine/
40│ ├── combine.h # GMM2 输出远端写回 offsetD40│ ├── combine.h # GMM2 输出远端写回 offsetD
41│ ├── unpermute.h # topK weighted reduce 和原 token 顺序还原41│ ├── unpermute.h # topK weighted reduce 和原 token 顺序还原
42│ └── utils/ # PTO vector、sync、HCCL window、GMM pipeline helper42│ └── utils/ # PTO vector、sync、HCCL window、GMM pipeline helper
43-├── overview.md # 总体设计、性能对比和阶段伪码43+├── overview.md # 总体设计和性能对比
44-├── front_reorder.md # front reorder / sort / count-as-flag 细节44+├── overview_v1.md # 固定分组调度和 overlap 设计
45-├── dispatch.md # dispatch 阶段契约和搬运策略45+├── pseudocode.md # 当前七阶段数据流和伪码
46-├── gmm1.md / gmm2.md # GMM tile 调度、swizzle、同步和 pipeline46+└── README.md # 英文 README
47-├── swiglu.md # SwiGLU 分段和量化策略
48-├── combine.md # combine large/small path 和远端写回协议
49-├── unpermute.md # unpermute 还原与累加策略
50-└── glden.md # Python golden batch rewrite 设计
51```47```
52 48 
53## 算子说明49## 算子说明
@@ -89,12 +85,12 @@ for each rank, token:
89 85 
90## 优化说明86## 优化说明
91 87 
92-- **expert 级流水重叠**：AIC 侧 GMM1/GMM2 和 AIV 侧 dispatch/SwiGLU/combine 按 local expert group 轮转推进，通过 hard flag 串接阶段边界。88+- **expert 级流水重叠**：AIC 侧 GMM1/GMM2 和 AIV 侧 Dispatch/SwiGLU/Combine 按 local expert 推进，通过逐 expert 的 arrival/ready 通知串接阶段边界。
93- **front reorder 三路径**：按 UB 工作集大小选择 FullLoad、OneCore、MultiCore。小 route 尽量留在 UB 内完成排序、count、反排和 quant，避免不必要的 GM 中间态。89- **front reorder 三路径**：按 UB 工作集大小选择 FullLoad、OneCore、MultiCore。小 route 尽量留在 UB 内完成排序、count、反排和 quant，避免不必要的 GM 中间态。
94- **count-as-flag**：front 发布 `tokenPerExpert` 时给 count row 加 marker，peer 通过数据值等待到达，减少 AlltoAll count 后的整机同步。90- **count-as-flag**：front 发布 `tokenPerExpert` 时给 count row 加 marker，peer 通过数据值等待到达，减少 AlltoAll count 后的整机同步。
95- **GMM PTO tile 优化**：GMM1/GMM2 使用 PTO tile 编程，包含 output tile swizzle、L1 -> L0 多级复用、双缓冲和 fixpipe quant/cast。91- **GMM PTO tile 优化**：GMM1/GMM2 使用 PTO tile 编程，包含 output tile swizzle、L1 -> L0 多级复用、双缓冲和 fixpipe quant/cast。
96-- **SwiGLU 分段 overlap**：SwiGLU 按 segment 切分，第一段尽量压到 GMM2 启动前完成，降低 GMM1 -> GMM2 中间等待。92+- **AIC 动态分组**：GMM1 先使用全部 24 个 AIC，随后释放 8 个；GMM2 从 8 个 AIC 起步，在 GMM1 完成后扩展到 24 个。
97-- **combine 双路径**：large token path 按完整 row 写回，small token path 按 GMM2 tile 拆 subtile，提高小 M 场景 AIV 利用率。93+- **两阶段 Unpermute**：前 32 个 AIV 先处理 route 已就绪的 token，Combine 全部完成后由 48 个 AIV 处理剩余 token。
98 94 
99## Tiling 参数95## Tiling 参数
100 96 
@@ -107,11 +103,11 @@ for each rank, token:
107| `expertPerRank` | 每 rank 本地 expert 数 |103| `expertPerRank` | 每 rank 本地 expert 数 |
108| `worldSize` | MPI/HCCL rank 数 |104| `worldSize` | MPI/HCCL rank 数 |
109| `maxOutputSize` | 每 rank routed row workspace 上限；典型性能 case 显式使用固定 workspace 上限 |105| `maxOutputSize` | 每 rank routed row workspace 上限；典型性能 case 显式使用固定 workspace 上限 |
110-| `aicNum` | AIC 逻辑核数，默认脚本为 24 |106+| `aicNum` | 生产路径固定物理 AIC 数：24 |
111-| `aivNum` | AIV 逻辑核数，默认脚本为 48 |107+| `aivNum` | 生产路径固定物理 AIV 数：48 |
112| `GMM baseM/baseN` | 主要 tile 口径为 `128 x 256` output tile |108| `GMM baseM/baseN` | 主要 tile 口径为 `128 x 256` output tile |
113| `Front FullLoad` | 由 `routeElems`、`K`、`expertNum` 和 UB 192 KiB 预算共同决定 |109| `Front FullLoad` | 由 `routeElems`、`K`、`expertNum` 和 UB 192 KiB 预算共同决定 |
114-| `Combine small path` | `problemM * topK <= 4096` 时倾向使用 subtile path |110+| 固定分组 | GMM1/GMM2 使用 `16 + 8` 个 AIC；Dispatch/SwiGLU/Combine 预留 `16/16/8` 个 AIV slot，SwiGLU 按 M 激活 8 或 16 个 |
115 111 
116## 支持 Case112## 支持 Case
117 113 
@@ -146,23 +142,23 @@ aivNum=48
146│ FrontReorder (AIV) │142│ FrontReorder (AIV) │
147│ sort expertId route -> offsetA + count/prefix metadata │143│ sort expertId route -> offsetA + count/prefix metadata │
148└──────────────────────────────┬───────────────────────────────────────────────┘144└──────────────────────────────┬───────────────────────────────────────────────┘
149- │ D2C ready / count-as-flag145+ │ count metadata ready / count-as-flag
150┌──────────────────────────────▼───────────────────────────────────────────────┐146┌──────────────────────────────▼───────────────────────────────────────────────┐
151-│ Expert-level overlapped pipeline │147+│ Expert 级重叠流水 │
152│ │148│ │
153-│ AIV: Dispatch(group i) -> SwiGLU(segment/group i) -> Combine(group i) │149+│ AIV: Dispatch(expert i) -> SwiGLU(expert i) -> Combine(expert i) │
154-│ AIC: GMM1(group i) -> GMM2(group i) │150+│ AIC: GMM1 从 24 缩为 16；释放的 AIC 加入 GMM2，使其从 8 扩为 24 │
155│ │151│ │
156-│ Stages communicate with hard flags: D2C, C2V, V2C, G2C/Combine ready │152+│ 各阶段通过逐 expert 的 GM arrival/ready 进度衔接 │
157└──────────────────────────────┬───────────────────────────────────────────────┘153└──────────────────────────────┬───────────────────────────────────────────────┘
158- │ final boundary154+ │ 各 rank expert 进度 / DataReady
159┌──────────────────────────────▼───────────────────────────────────────────────┐155┌──────────────────────────────▼───────────────────────────────────────────────┐
160-│ Unpermute (AIV) │156+│ 两阶段 Unpermute (AIV) │
161-│ offsetD + probs + expandedRowIdx -> topK weighted reduce -> out[M, K] │157+│ 32 个 AIV 处理已就绪 token -> 48 个 AIV 处理剩余 token -> out │
162└──────────────────────────────────────────────────────────────────────────────┘158└──────────────────────────────────────────────────────────────────────────────┘
163```159```
164 160 
165-完整执行顺序在 `MegaMoe::Process()` 中串联，阶段顺序为：161+`MegaMoe::Process()` 先完成 FrontReorder，再由固定分组按以下依赖关系并发推进：
166 162 
167```text163```text
168FrontReorder -> Dispatch -> GMM1 -> SwiGLU -> GMM2 -> Combine -> Unpermute164FrontReorder -> Dispatch -> GMM1 -> SwiGLU -> GMM2 -> Combine -> Unpermute
@@ -211,7 +207,7 @@ srcRowBase = preSumBeforeRank[srcRank, localExpert]
211dstRowBase = groupBase + (srcRank == 0 ? 0 : cumsumMM[srcRank - 1, localExpert])207dstRowBase = groupBase + (srcRank == 0 ? 0 : cumsumMM[srcRank - 1, localExpert])
212```208```
213 209 
214-每个 local expert group 搬运完成后，dispatch 设置 GMM1 ready flag，允许 AIC 开始消费该 group。210+每个 local expert 搬运完成后，Dispatch 汇聚活跃 worker 的到达状态并发布 GMM1 ready 进度。
215 211 
216## GMM1 / SwiGLU / GMM2 阶段212## GMM1 / SwiGLU / GMM2 阶段
217 213 
@@ -226,7 +222,7 @@ gmA[int8] x weight1[int8]
226 -> gmC[half]222 -> gmC[half]
227```223```
228 224 
229-每个 local expert 的输出 tile 网格按 `128 x 256` output tile 切分。线性 tile id 会通过 swizzle 映射到 `(blockM, blockN)`，让相邻 tile 更容易复用 L1 中的 B 侧权重。225+每个 local expert 的输出 tile 网格按 `128 x 256` output tile 切分。线性 tile id 会通过 swizzle 映射到 `(blockM, blockN)`，让相邻 tile 更容易复用 L1 中的 B 侧权重。前两个 expert 使用全部 24 个 AIC，后续 expert 使用 16-AIC GMM1 组。
230 226 
231### SwiGLU227### SwiGLU
232 228 
@@ -239,7 +235,7 @@ gmC * perTokenScale1
239 -> gmPermutedToken[int8] + perTokenScale2[float]235 -> gmPermutedToken[int8] + perTokenScale2[float]
240```236```
241 237 
242-SwiGLU 按 segment 切分；core0 写 segment metadata，其它 AIV 读取后按 row 分担计算。238+SwiGLU 按 expert 推进，并在固定活跃组内按 row 分工。全部活跃 AIV 完成当前 expert 后，由 coordinator 发布 GMM2 ready 进度；M=16 使用 8 个 AIV，其余表中 case 使用 16 个。
243 239 
244### GMM2240### GMM2
245 241 
@@ -252,7 +248,7 @@ gmPermutedToken[int8] x weight2[int8]
252 -> gmm2Output[half]248 -> gmm2Output[half]
253```249```
254 250 
255-GMM2 完成每个 local expert group 后设置 combine ready flag，AIV combine 才能写回该 group。251+GMM2 从 8-AIC 组开始执行。GMM1 完成后，其释放的 16 个 AIC 在 expert 边界加入 GMM2；每个 expert 完成后发布 arrival 进度供 Combine 消费。
256 252 
257## Combine / Unpermute 阶段253## Combine / Unpermute 阶段
258 254 
@@ -266,17 +262,14 @@ gmm2Output[srcRow, 0:K] half
266 -> srcRank.remoteWindow.offsetD[dstRow, 0:K]262 -> srcRank.remoteWindow.offsetD[dstRow, 0:K]
267```263```
268 264 
269-路径选择：265+8-AIV Combine 组按 source rank 分配完整 row 写回。执行过程中发布各 rank 的 expert 进度，全部 expert 可见后发布最终 DataReady。
270- 
271-- **DirectLarge**：大 token 量场景，按完整 row 写回。
272-- **DirectSmall**：小 token 量场景，按 subtile 拆分，提升 AIV 并行度。
273- 
274 266 
275Unpermute 是最后的源 rank 还原阶段：267Unpermute 是最后的源 rank 还原阶段：
276 268 
277```text269```text
278offsetD + probs + expandedRowIdx270offsetD + probs + expandedRowIdx
279- -> 按原 token/topK 加权累加271+ -> 第一阶段：32 个 AIV 处理 route 已就绪的 token
272+ -> 第二阶段：48 个 AIV 处理剩余 token
280 -> out[M, K]273 -> out[M, K]
281```274```
282 275 
@@ -289,6 +282,7 @@ HCCL remote window 主要承载跨 rank 可见的数据：
289| `offsetA` | HCCL window | FrontReorder 写入 packed int8 token row，Dispatch 从 peer 拉取 |282| `offsetA` | HCCL window | FrontReorder 写入 packed int8 token row，Dispatch 从 peer 拉取 |
290| `offsetD` | HCCL window | Combine 写回源 rank，Unpermute 在源 rank 消费 |283| `offsetD` | HCCL window | Combine 写回源 rank，Unpermute 在源 rank 消费 |
291| `tokenPerExpert` | HCCL window | count-as-flag 的跨 rank count row |284| `tokenPerExpert` | HCCL window | count-as-flag 的跨 rank count row |
285+| `ExpertProgress / DataReady` | HCCL window | 各 rank 的 Combine 进度和当前 launch 完成通知 |
292| `gmA` | workspace GM | Dispatch 生成的 GMM1 输入 |286| `gmA` | workspace GM | Dispatch 生成的 GMM1 输入 |
293| `gmC` | workspace GM | GMM1 输出，SwiGLU 输入 |287| `gmC` | workspace GM | GMM1 输出，SwiGLU 输入 |
294| `gmPermutedToken` | workspace GM | SwiGLU dynamic quant 后的 GMM2 输入 |288| `gmPermutedToken` | workspace GM | SwiGLU dynamic quant 后的 GMM2 输入 |
@@ -322,7 +316,7 @@ overview.md
322 316 
323### 2. 保持 expert 级阶段边界清晰317### 2. 保持 expert 级阶段边界清晰
324 318 
325-Dispatch、GMM1、SwiGLU、GMM2、Combine 之间依赖 hard flag。优化时优先确认每个 group 的 set/wait 是否一一匹配，避免为了减少 `SYNCALL` 破坏跨 AIC/AIV 的真实数据依赖。319+Dispatch、GMM1、SwiGLU、GMM2、Combine 之间使用 producer arrival 和 consumer ready slot。修改同步前需要确认每个边界的 producer 数、consumer 范围和期望 expert 进度。
326 320 
327### 3. 优先优化 GMM tile 效率321### 3. 优先优化 GMM tile 效率
328 322 
@@ -333,9 +327,9 @@ GMM1/GMM2 是主耗时阶段。重点关注：
333- `currentM` 较小时 AIC 是否负载不均；327- `currentM` 较小时 AIC 是否负载不均；
334- AIV 通信/写回是否与 GMM HBM 访问冲突。328- AIV 通信/写回是否与 GMM HBM 访问冲突。
335 329 
336-### 4. 小 token 场景使用 subtile combine330+### 4. 平衡 Combine 与 GMM2
337 331 
338-小 M 下 direct row path 容易让 AIV 并行度不足。DirectSmall 通过 subtile 切分提高核利用率，但需要保持 `gmm2Tiling.l1TileN` 与 small path subtile 列宽对齐。332+Combine 固定使用 8 个 AIV，并按 shape 延后启动，避免与 GMM2 竞争 HBM 带宽。调优小 M 时需要一起检查 GMM2 扩组点、Combine 启动 expert 和 SwiGLU 活跃 worker 数。
339 333 
340### 5. Golden 生成使用 batch backend334### 5. Golden 生成使用 batch backend
341 335 
@@ -346,9 +340,7 @@ large synthetic case 默认使用 `python-batch` golden backend，避免逐 toke
346配置 Ascend CANN 环境：340配置 Ascend CANN 环境：
347 341 
348```bash342```bash
349-export ASCEND_CANN_PATH=/usr/local/Ascend/cann/set_env.sh343+source <cann-install>/set_env.sh
350-export ASCEND_HOME_PATH=/usr/local/Ascend/cann/cann
351-source /usr/local/Ascend/cann/cann/set_env.sh
352```344```
353 345 
354运行默认 2048 case：346运行默认 2048 case：
@@ -370,13 +362,13 @@ bash run.sh --world-size 8 --m 512 --k 7168 --n 4096 --topk 8 --experts 16 --max
370 362 
371| 环境变量 | 用途 | 默认行为 |363| 环境变量 | 用途 | 默认行为 |
372| --- | --- | --- |364| --- | --- | --- |
373-| `ASCEND_HOME_PATH` | CANN 安装目录 | 必须提前设置 |365+| `ASCEND_HOME_PATH` | 当前 CANN 安装目录 | 由 CANN `set_env.sh` 设置 |
374-| `CMAKE_COMPILER` | CMake 使用的编译器 | `bisheng` |
375| `MPI_ENV_BIN` | MPI/conda bin 路径 | `/home/ntlab/miniconda3/envs/ltr_pto/bin` |366| `MPI_ENV_BIN` | MPI/conda bin 路径 | `/home/ntlab/miniconda3/envs/ltr_pto/bin` |
376| `MPI_ENV_LIB` | MPI/conda lib 路径 | `/home/ntlab/miniconda3/envs/ltr_pto/lib` |367| `MPI_ENV_LIB` | MPI/conda lib 路径 | `/home/ntlab/miniconda3/envs/ltr_pto/lib` |
377| `MPI_LIB_PATH` | `libmpi.so` 绝对路径 | `${MPI_ENV_LIB}/libmpi.so` |368| `MPI_LIB_PATH` | `libmpi.so` 绝对路径 | `${MPI_ENV_LIB}/libmpi.so` |
378| `MPI_RUNNER` | MPI 启动命令 | `mpirun` |369| `MPI_RUNNER` | MPI 启动命令 | `mpirun` |
379| `HCCL_BUFFSIZE` | HCCL RDMA window 大小 | `run.sh` 按 case 自动抬高到安全值 |370| `HCCL_BUFFSIZE` | HCCL RDMA window 大小 | `run.sh` 按 case 自动抬高到安全值 |
371+| `DISPATCH_MEGA_COMBINE_START_SYNC_DEBUG` | 同步 kernel 入口，用于比较跨核计时 | 默认关闭 |
380 372 
381## 修改 Case 参数373## 修改 Case 参数
382 374 
@@ -401,7 +393,7 @@ bash run.sh --world-size 8 --m 512 --k 7168 --n 4096 --topk 8 --experts 16 --max
401| HCCL window too small | 手动设置的 `HCCL_BUFFSIZE` 低于 case 需求；取消覆盖或调大该变量 |393| HCCL window too small | 手动设置的 `HCCL_BUFFSIZE` 低于 case 需求；取消覆盖或调大该变量 |
402| MPI 启动失败 | 检查 `MPI_ENV_BIN`、`MPI_ENV_LIB`、`MPI_LIB_PATH` 是否指向同一个 conda/MPI 环境 |394| MPI 启动失败 | 检查 `MPI_ENV_BIN`、`MPI_ENV_LIB`、`MPI_LIB_PATH` 是否指向同一个 conda/MPI 环境 |
403| golden 生成很慢 | 使用默认 `python-batch`，必要时调大 `--golden-chunk-rows`；只有调试对照才使用 `python-naive` |395| golden 生成很慢 | 使用默认 `python-batch`，必要时调大 `--golden-chunk-rows`；只有调试对照才使用 `python-naive` |
404-| 小 M 性能不稳定 | 优先检查 FullLoad case 是否命中、combine 是否走 DirectSmall、AIV 并发是否过高影响 GMM |396+| 小 M 性能不稳定 | 检查 FullLoad 是否命中，以及 SwiGLU worker 数、GMM2 扩组点和 Combine 启动点是否加重 HBM 竞争 |
405| 结果 diff 异常 | 先检查 data cache 是否复用旧分布；改变 expert 分布或 case 关键参数后不要使用旧 `out/` |397| 结果 diff 异常 | 先检查 data cache 是否复用旧分布；改变 expert 分布或 case 关键参数后不要使用旧 `out/` |
406 398 
407## 构建系统399## 构建系统
Mkernels/manual/a2a3/dispatch_mega_combine/data_utils.cpp+16-11
@@ -131,14 +131,20 @@ CaseConfig LoadCaseConfig(const std::string& case_json_path)
131 cfg.aiv_num = ParseJsonUInt(text, "aiv_num");131 cfg.aiv_num = ParseJsonUInt(text, "aiv_num");
132 cfg.compare_atol = ParseJsonDouble(text, "compare_atol", 1e-3);132 cfg.compare_atol = ParseJsonDouble(text, "compare_atol", 1e-3);
133 cfg.compare_rtol = ParseJsonDouble(text, "compare_rtol", 1e-3);133 cfg.compare_rtol = ParseJsonDouble(text, "compare_rtol", 1e-3);
134- cfg.input_tokens_all_ranks =134+ const double defaultInputTokens = static_cast<double>(cfg.m) * cfg.world_size;
135- ParseJsonDouble(text, "input_tokens_all_ranks", static_cast<double>(cfg.m) * cfg.world_size);135+ cfg.input_tokens_all_ranks = ParseJsonDouble(text, "input_tokens_all_ranks", defaultInputTokens);
136- cfg.routed_tokens_all_ranks =136+ const double defaultRoutedTokens = cfg.input_tokens_all_ranks * cfg.topk;
137- ParseJsonDouble(text, "routed_tokens_all_ranks", static_cast<double>(cfg.m) * cfg.topk * cfg.world_size);137+ cfg.routed_tokens_all_ranks = ParseJsonDouble(text, "routed_tokens_all_ranks", defaultRoutedTokens);
138- cfg.remote_routed_tokens_all_ranks = ParseJsonDouble(text, "remote_routed_tokens_all_ranks", 0.0);138+ const double defaultRemoteRoutedTokens =
139- cfg.compute_flops_all_ranks =139+ cfg.world_size == 0U ? 0.0 :
140- ParseJsonDouble(text, "compute_flops_all_ranks", cfg.routed_tokens_all_ranks * 3.0 * cfg.k * cfg.n);140+ cfg.routed_tokens_all_ranks * static_cast<double>(cfg.world_size - 1U) / cfg.world_size;
141- cfg.comm_bytes_all_ranks = ParseJsonDouble(text, "comm_bytes_all_ranks", 0.0);141+ cfg.remote_routed_tokens_all_ranks =
142+ ParseJsonDouble(text, "remote_routed_tokens_all_ranks", defaultRemoteRoutedTokens);
143+ const double defaultComputeFlops = cfg.routed_tokens_all_ranks * 3.0 * cfg.k * cfg.n;
144+ cfg.compute_flops_all_ranks = ParseJsonDouble(text, "compute_flops_all_ranks", defaultComputeFlops);
145+ const double bytesPerRemoteRoute = static_cast<double>(cfg.k) * 3.0 + sizeof(float);
146+ cfg.comm_bytes_all_ranks =
147+ ParseJsonDouble(text, "comm_bytes_all_ranks", cfg.remote_routed_tokens_all_ranks * bytesPerRemoteRoute);
142 return cfg;148 return cfg;
143}149}
144 150 
@@ -146,9 +152,8 @@ RankFileSet BuildRankFileSet(const std::string& case_dir, int rank)
146{152{
147 const std::string prefix = case_dir + "/rank" + std::to_string(rank) + "_";153 const std::string prefix = case_dir + "/rank" + std::to_string(rank) + "_";
148 return RankFileSet{154 return RankFileSet{
149- prefix + "x.bin", prefix + "weight1.bin", prefix + "weight2.bin",155+ prefix + "x.bin", prefix + "weight1.bin", prefix + "weight2.bin", prefix + "expert_idx.bin",
150- prefix + "expert_idx.bin", prefix + "scale1.bin", prefix + "scale2.bin",156+ prefix + "scale1.bin", prefix + "scale2.bin", prefix + "probs.bin", prefix + "expected_out.bin",
151- prefix + "probs.bin", prefix + "x_active_mask.bin", prefix + "expected_out.bin",
152 };157 };
153}158}
154 159 
Mkernels/manual/a2a3/dispatch_mega_combine/data_utils.hpp+0-1
@@ -25,7 +25,6 @@ struct RankFileSet {
25 std::string scale1;25 std::string scale1;
26 std::string scale2;26 std::string scale2;
27 std::string probs;27 std::string probs;
28- std::string x_active_mask;
29 std::string expected_out;28 std::string expected_out;
30};29};
31 30 
Mkernels/manual/a2a3/dispatch_mega_combine/kernel_launch.cpp+15-11
@@ -49,11 +49,6 @@ extern "C" __global__ __aicore__ void dispatch_mega_combine_kernel(
49 GM_ADDR probs, GM_ADDR c, GM_ADDR expertTokenNums, GM_ADDR workspaceGM, GM_ADDR tilingGM, GM_ADDR profileGM,49 GM_ADDR probs, GM_ADDR c, GM_ADDR expertTokenNums, GM_ADDR workspaceGM, GM_ADDR tilingGM, GM_ADDR profileGM,
50 uint32_t startSyncDebug)50 uint32_t startSyncDebug)
51{51{
52-#ifdef _DEBUG
53- cce::printf(
54- "MegaMoe kernel enter block=%d sub=%d DAV_VEC=%d workspace=%d tiling=%d\n", int(get_block_idx()),
55- int(get_subblockid()), int(DAV_VEC), int(workspaceGM != nullptr), int(tilingGM != nullptr));
56-#endif
57 __gm__ uint64_t* profileEntry = nullptr;52 __gm__ uint64_t* profileEntry = nullptr;
58 if (profileGM != nullptr) {53 if (profileGM != nullptr) {
59 std::size_t profileOffset = static_cast<std::size_t>(get_block_idx()) * kMegaMoeProfileBytesPerBlock;54 std::size_t profileOffset = static_cast<std::size_t>(get_block_idx()) * kMegaMoeProfileBytesPerBlock;
@@ -66,12 +61,18 @@ extern "C" __global__ __aicore__ void dispatch_mega_combine_kernel(
66 set_ffts_base_addr(reinterpret_cast<uint64_t>(fftsAddr));61 set_ffts_base_addr(reinterpret_cast<uint64_t>(fftsAddr));
67 if (workspaceGM != nullptr && tilingGM != nullptr) {62 if (workspaceGM != nullptr && tilingGM != nullptr) {
68 const __gm__ MegaMoeTilingData* tilingData = reinterpret_cast<__gm__ MegaMoeTilingData*>(tilingGM);63 const __gm__ MegaMoeTilingData* tilingData = reinterpret_cast<__gm__ MegaMoeTilingData*>(tilingGM);
69-#ifdef _DEBUG64+#if defined(__DAV_VEC__)
70- cce::printf(65+ if (get_block_idx() == 0U && get_subblockid() == 0U) {
71- "MegaMoe kernel tiling block=%d sub=%d stage=%d M=%d topK=%d expertNum=%d\n", int(get_block_idx()),66+ ResetFixedSyncWorkspace(workspaceGM, tilingData);
72- int(get_subblockid()), int(tilingData->frontReorderTiling.stageNum), int(tilingData->megaMoeInfo.M),67+ if (tilingData->frontReorderTiling.stageNum >= 13U &&
73- int(tilingData->megaMoeInfo.topK), int(tilingData->frontReorderTiling.expertNum));68+ tilingData->unpermuteTiling.unpermuteImplMode == kMegaMoeUnpermuteImplRankStreaming) {
69+ PtoRemoteWindow remoteWindow;
70+ remoteWindow.Init(reinterpret_cast<GM_ADDR>(tilingData->runtimeInfo.remoteWindowContext));
71+ remoteWindow.PrepareDataReadyEpoch();
72+ }
73+ }
74#endif74#endif
75+ pto::SYNCALL<pto::SyncCoreType::Mix>();
75 if (startSyncDebug != 0U) {76 if (startSyncDebug != 0U) {
76 PtoRemoteWindow remoteWindow;77 PtoRemoteWindow remoteWindow;
77 remoteWindow.Init(reinterpret_cast<GM_ADDR>(tilingData->runtimeInfo.remoteWindowContext));78 remoteWindow.Init(reinterpret_cast<GM_ADDR>(tilingData->runtimeInfo.remoteWindowContext));
@@ -86,9 +87,12 @@ extern "C" __global__ __aicore__ void dispatch_mega_combine_kernel(
86 if (tilingData->megaMoeInfo.expertPerRank == 8U) {87 if (tilingData->megaMoeInfo.expertPerRank == 8U) {
87 RunMegaMoeSpecialized<8U>(88 RunMegaMoeSpecialized<8U>(
88 x, w1, w2, expertId, scale1, scale2, probs, c, expertTokenNums, workspaceGM, tilingData);89 x, w1, w2, expertId, scale1, scale2, probs, c, expertTokenNums, workspaceGM, tilingData);
89- } else {90+ } else if (tilingData->megaMoeInfo.expertPerRank == 16U) {
90 RunMegaMoeSpecialized<16U>(91 RunMegaMoeSpecialized<16U>(
91 x, w1, w2, expertId, scale1, scale2, probs, c, expertTokenNums, workspaceGM, tilingData);92 x, w1, w2, expertId, scale1, scale2, probs, c, expertTokenNums, workspaceGM, tilingData);
93+ } else if (tilingData->megaMoeInfo.expertPerRank == 32U) {
94+ RunMegaMoeSpecialized<32U>(
95+ x, w1, w2, expertId, scale1, scale2, probs, c, expertTokenNums, workspaceGM, tilingData);
92 }96 }
93 }97 }
94 98 
Mkernels/manual/a2a3/dispatch_mega_combine/main.cpp+71-307
@@ -1,5 +1,5 @@
1/**1/**
2-Copyright (c) 2025 Huawei Technologies Co., Ltd.2+Copyright (c) 2026 Huawei Technologies Co., Ltd.
3This program is free software, you can redistribute it and/or modify it under the terms and conditions of3This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4CANN Open Software License Agreement Version 2.0 (the "License").4CANN Open Software License Agreement Version 2.0 (the "License").
5Please refer to the License for details. You may not use this file except in compliance with the License.5Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -11,10 +11,10 @@ See LICENSE in the root of the software repository for the full text of the Lice
11#include <algorithm>11#include <algorithm>
12#include <cmath>12#include <cmath>
13#include <cstdlib>13#include <cstdlib>
14-#include <cstring>
15#include <iomanip>14#include <iomanip>
16#include <iostream>15#include <iostream>
17#include <limits>16#include <limits>
17+#include <memory>
18#include <numeric>18#include <numeric>
19#include <sstream>19#include <sstream>
20#include <stdexcept>20#include <stdexcept>
@@ -41,76 +41,33 @@ constexpr int kDefaultMeasureIters = 5;
41constexpr double kMicrosecondsPerSecond = 1000.0 * 1000.0;41constexpr double kMicrosecondsPerSecond = 1000.0 * 1000.0;
42constexpr double kBytesPerGiB = 1024.0 * 1024.0 * 1024.0;42constexpr double kBytesPerGiB = 1024.0 * 1024.0 * 1024.0;
43static double g_sys_cnt_multiple = 20.0; // Default A2/A3, in ns per SYS_CNT tick.43static double g_sys_cnt_multiple = 20.0; // Default A2/A3, in ns per SYS_CNT tick.
44-constexpr uint32_t kHostCombineImplDirectAuto = 3U;
45 44 
46-struct DeviceBuffer {45+struct DeviceMemoryReleaser {
47- void* ptr = nullptr;46+ void operator()(void* allocation) const noexcept { (void)aclrtFree(allocation); }
48- size_t bytes = 0;
49- 
50- DeviceBuffer() = default;
51- DeviceBuffer(const DeviceBuffer&) = delete;
52- DeviceBuffer& operator=(const DeviceBuffer&) = delete;
53- DeviceBuffer(DeviceBuffer&& other) noexcept : ptr(other.ptr), bytes(other.bytes)
54- {
55- other.ptr = nullptr;
56- other.bytes = 0;
57- }
58- DeviceBuffer& operator=(DeviceBuffer&& other) noexcept
59- {
60- if (this != &other) {
61- if (ptr != nullptr) {
62- aclrtFree(ptr);
63- }
64- ptr = other.ptr;
65- bytes = other.bytes;
66- other.ptr = nullptr;
67- other.bytes = 0;
68- }
69- return *this;
70- }
71- 
72- ~DeviceBuffer()
73- {
74- if (ptr != nullptr) {
75- aclrtFree(ptr);
76- }
77- }
78};47};
79 48 
80-struct HostBuffer {49+struct HostMemoryReleaser {
81- void* ptr = nullptr;50+ void operator()(void* allocation) const noexcept { (void)aclrtFreeHost(allocation); }
82- size_t bytes = 0;
83- 
84- HostBuffer() = default;
85- HostBuffer(const HostBuffer&) = delete;
86- HostBuffer& operator=(const HostBuffer&) = delete;
87- HostBuffer(HostBuffer&& other) noexcept : ptr(other.ptr), bytes(other.bytes)
88- {
89- other.ptr = nullptr;
90- other.bytes = 0;
91- }
92- HostBuffer& operator=(HostBuffer&& other) noexcept
93- {
94- if (this != &other) {
95- if (ptr != nullptr) {
96- aclrtFreeHost(ptr);
97- }
98- ptr = other.ptr;
99- bytes = other.bytes;
100- other.ptr = nullptr;
101- other.bytes = 0;
102- }
103- return *this;
104- }
105- 
106- ~HostBuffer()
107- {
108- if (ptr != nullptr) {
109- aclrtFreeHost(ptr);
110- }
111- }
112};51};
113 52 
53+template <typename Releaser>
54+class AclOwnedBuffer {
55+public:
56+ AclOwnedBuffer() = default;
57+ AclOwnedBuffer(void* allocation, size_t allocationBytes) : storage_(allocation), byteCount_(allocationBytes) {}
58+ 
59+ void* data() const noexcept { return storage_.get(); }
60+ 
61+ size_t size() const noexcept { return byteCount_; }
62+ 
63+private:
64+ std::unique_ptr<void, Releaser> storage_;
65+ size_t byteCount_ = 0;
66+};
67+ 
68+using DeviceBuffer = AclOwnedBuffer<DeviceMemoryReleaser>;
69+using HostBuffer = AclOwnedBuffer<HostMemoryReleaser>;
70+ 
114struct PerfStats {71struct PerfStats {
115 double avg = 0.0;72 double avg = 0.0;
116 double min = 0.0;73 double min = 0.0;
@@ -121,9 +78,7 @@ struct PerfStats {
121struct RunOptions {78struct RunOptions {
122 int warmup_iters = kDefaultWarmupIters;79 int warmup_iters = kDefaultWarmupIters;
123 int measure_iters = kDefaultMeasureIters;80 int measure_iters = kDefaultMeasureIters;
124- bool skip_accuracy = false;
125 bool start_sync_debug = false;81 bool start_sync_debug = false;
126- bool workload_audit = false;
127};82};
128 83 
129struct RankHostInputs {84struct RankHostInputs {
@@ -134,7 +89,6 @@ struct RankHostInputs {
134 std::vector<uint8_t> scale1;89 std::vector<uint8_t> scale1;
135 std::vector<uint8_t> scale2;90 std::vector<uint8_t> scale2;
136 std::vector<uint8_t> probs;91 std::vector<uint8_t> probs;
137- std::vector<uint8_t> x_active_mask;
138 std::vector<uint16_t> expected_out;92 std::vector<uint16_t> expected_out;
139};93};
140 94 
@@ -156,16 +110,16 @@ struct RankDeviceBuffers {
156 110 
157DeviceBuffer MakeDeviceBuffer(size_t bytes, const void* host_src = nullptr)111DeviceBuffer MakeDeviceBuffer(size_t bytes, const void* host_src = nullptr)
158{112{
159- DeviceBuffer buffer;
160- buffer.bytes = bytes;
161 if (bytes == 0) {113 if (bytes == 0) {
162- return buffer;114+ return {};
163 }115 }
164- if (aclrtMalloc(&buffer.ptr, bytes, ACL_MEM_MALLOC_HUGE_FIRST) != ACL_SUCCESS) {116+ void* allocation = nullptr;
117+ if (aclrtMalloc(&allocation, bytes, ACL_MEM_MALLOC_HUGE_FIRST) != ACL_SUCCESS) {
165 throw std::runtime_error("aclrtMalloc failed");118 throw std::runtime_error("aclrtMalloc failed");
166 }119 }
120+ DeviceBuffer buffer(allocation, bytes);
167 if (host_src != nullptr &&121 if (host_src != nullptr &&
168- aclrtMemcpy(buffer.ptr, bytes, host_src, bytes, ACL_MEMCPY_HOST_TO_DEVICE) != ACL_SUCCESS) {122+ aclrtMemcpy(buffer.data(), bytes, host_src, bytes, ACL_MEMCPY_HOST_TO_DEVICE) != ACL_SUCCESS) {
169 throw std::runtime_error("aclrtMemcpy host->device failed");123 throw std::runtime_error("aclrtMemcpy host->device failed");
170 }124 }
171 return buffer;125 return buffer;
@@ -173,15 +127,14 @@ DeviceBuffer MakeDeviceBuffer(size_t bytes, const void* host_src = nullptr)
173 127 
174HostBuffer MakeHostBuffer(size_t bytes)128HostBuffer MakeHostBuffer(size_t bytes)
175{129{
176- HostBuffer buffer;
177- buffer.bytes = bytes;
178 if (bytes == 0) {130 if (bytes == 0) {
179- return buffer;131+ return {};
180 }132 }
181- if (aclrtMallocHost(&buffer.ptr, bytes) != ACL_SUCCESS) {133+ void* allocation = nullptr;
134+ if (aclrtMallocHost(&allocation, bytes) != ACL_SUCCESS) {
182 throw std::runtime_error("aclrtMallocHost failed");135 throw std::runtime_error("aclrtMallocHost failed");
183 }136 }
184- return buffer;137+ return HostBuffer(allocation, bytes);
185}138}
186 139 
187std::vector<uint16_t> BytesToU16(const std::vector<uint8_t>& bytes)140std::vector<uint16_t> BytesToU16(const std::vector<uint8_t>& bytes)
@@ -211,8 +164,6 @@ int ParseEnvInt(const char* name, int default_value)
211 }164 }
212}165}
213 166 
214-bool TraceEnabled() { return ParseEnvInt("DISPATCH_MEGA_COMBINE_TRACE", 0) != 0; }
215- 
216uint64_t AlignUpU64(uint64_t value, uint64_t align)167uint64_t AlignUpU64(uint64_t value, uint64_t align)
217{168{
218 if (align == 0U) {169 if (align == 0U) {
@@ -234,14 +185,6 @@ uint64_t SwigluFullRowUbBytes(uint32_t n)
234 return ub_offset;185 return ub_offset;
235}186}
236 187 
237-void Trace(int rank_id, const std::string& message)
238-{
239- if (!TraceEnabled()) {
240- return;
241- }
242- std::cerr << "[trace] rank=" << rank_id << " " << message << std::endl;
243-}
244- 
245bool ZeroWindowMemory(const StandaloneRankRuntime& runtime)188bool ZeroWindowMemory(const StandaloneRankRuntime& runtime)
246{189{
247 const uint64_t window_bytes = runtime.hccl.WindowBytes();190 const uint64_t window_bytes = runtime.hccl.WindowBytes();
@@ -254,10 +197,10 @@ bool ZeroWindowMemory(const StandaloneRankRuntime& runtime)
254 197 
255void ZeroDeviceBuffer(const DeviceBuffer& buffer, const char* name)198void ZeroDeviceBuffer(const DeviceBuffer& buffer, const char* name)
256{199{
257- if (buffer.bytes == 0) {200+ if (buffer.size() == 0) {
258 return;201 return;
259 }202 }
260- if (aclrtMemset(buffer.ptr, buffer.bytes, 0, buffer.bytes) != ACL_SUCCESS) {203+ if (aclrtMemset(buffer.data(), buffer.size(), 0, buffer.size()) != ACL_SUCCESS) {
261 throw std::runtime_error(std::string("failed to zero ") + name);204 throw std::runtime_error(std::string("failed to zero ") + name);
262 }205 }
263}206}
@@ -332,18 +275,18 @@ std::vector<double> GatherMaxSamplesToRoot(const std::vector<double>& local_samp
332 275 
333double ReadKernelProfileUs(const DeviceBuffer& profile_dev, HostBuffer& profile_host, uint32_t block_dim)276double ReadKernelProfileUs(const DeviceBuffer& profile_dev, HostBuffer& profile_host, uint32_t block_dim)
334{277{
335- if (profile_dev.bytes == 0 || profile_host.bytes == 0 || block_dim == 0) {278+ if (profile_dev.size() == 0 || profile_host.size() == 0 || block_dim == 0) {
336 return 0.0;279 return 0.0;
337 }280 }
338 if (aclrtMemcpy(281 if (aclrtMemcpy(
339- profile_host.ptr, profile_host.bytes, profile_dev.ptr, profile_dev.bytes, ACL_MEMCPY_DEVICE_TO_HOST) !=282+ profile_host.data(), profile_host.size(), profile_dev.data(), profile_dev.size(),
340- ACL_SUCCESS) {283+ ACL_MEMCPY_DEVICE_TO_HOST) != ACL_SUCCESS) {
341 throw std::runtime_error("device->host profile copy failed");284 throw std::runtime_error("device->host profile copy failed");
342 }285 }
343 286 
344 uint64_t start_min = std::numeric_limits<uint64_t>::max();287 uint64_t start_min = std::numeric_limits<uint64_t>::max();
345 uint64_t end_max = 0;288 uint64_t end_max = 0;
346- const auto* profile = static_cast<const uint8_t*>(profile_host.ptr);289+ const auto* profile = static_cast<const uint8_t*>(profile_host.data());
347 for (uint32_t block = 0; block < block_dim; ++block) {290 for (uint32_t block = 0; block < block_dim; ++block) {
348 for (size_t profile_idx = 0; profile_idx < kMegaMoeProfileEntriesPerBlock; ++profile_idx) {291 for (size_t profile_idx = 0; profile_idx < kMegaMoeProfileEntriesPerBlock; ++profile_idx) {
349 const uint64_t* entry = reinterpret_cast<const uint64_t*>(292 const uint64_t* entry = reinterpret_cast<const uint64_t*>(
@@ -414,168 +357,14 @@ void PrintPerfSummary(
414 << " routed_tokens/s=" << ToTokensPerSecond(cfg.routed_tokens_all_ranks, kernel_stats.avg)357 << " routed_tokens/s=" << ToTokensPerSecond(cfg.routed_tokens_all_ranks, kernel_stats.avg)
415 << " eq_compute=" << ToTflops(cfg.compute_flops_all_ranks, kernel_stats.avg) << " TFLOPS"358 << " eq_compute=" << ToTflops(cfg.compute_flops_all_ranks, kernel_stats.avg) << " TFLOPS"
416 << " eq_comm=" << ToGbs(cfg.comm_bytes_all_ranks, kernel_stats.avg) << " GB/s\n";359 << " eq_comm=" << ToGbs(cfg.comm_bytes_all_ranks, kernel_stats.avg) << " GB/s\n";
417- std::cout360+ std::cout << " note: equivalent compute/comm are logical workload estimates, not hardware counters.\n";
418- << " note: equivalent compute/comm are derived from case.json logical workload, not hardware counters.\n";
419 std::cout << "===============================================================\n" << std::endl;361 std::cout << "===============================================================\n" << std::endl;
420}362}
421 363 
422-int32_t LoadI32(const std::vector<uint8_t>& bytes, size_t index)
423-{
424- const size_t byteOffset = index * sizeof(int32_t);
425- if (byteOffset + sizeof(int32_t) > bytes.size()) {
426- return 0;
427- }
428- uint32_t value = 0U;
429- for (size_t byteIdx = 0; byteIdx < sizeof(int32_t); ++byteIdx) {
430- value |= static_cast<uint32_t>(bytes[byteOffset + byteIdx]) << (byteIdx * 8U);
431- }
432- return static_cast<int32_t>(value);
433-}
434- 
435-struct WorkloadAuditLayout {
436- size_t global_expert_num = 0;
437- size_t fixed_fields = 6;
438- size_t per_rank_offset = 6;
439- size_t per_expert_offset = 0;
440- size_t fields = 0;
441-};
442- 
443-WorkloadAuditLayout BuildWorkloadAuditLayout(const CaseConfig& cfg)
444-{
445- WorkloadAuditLayout layout;
446- layout.global_expert_num = static_cast<size_t>(cfg.world_size) * cfg.expert_per_rank;
447- layout.per_rank_offset = layout.fixed_fields;
448- layout.per_expert_offset = layout.per_rank_offset + cfg.world_size;
449- layout.fields = layout.per_expert_offset + layout.global_expert_num;
450- return layout;
451-}
452- 
453-std::vector<uint64_t> BuildLocalWorkloadAudit(
454- const CaseConfig& cfg, const WorkloadAuditLayout& layout, int rank_id, const std::vector<uint8_t>& expert_idx,
455- const std::vector<uint8_t>& x_active_mask, size_t actual_output_elems)
456-{
457- std::vector<uint64_t> local(layout.fields, 0);
458- uint64_t active_tokens = 0;
459- uint64_t valid_routes = 0;
460- uint64_t remote_routes = 0;
461- uint64_t invalid_expert_routes = 0;
462- const size_t expert_idx_count = expert_idx.size() / sizeof(int32_t);
463- for (uint32_t token = 0; token < cfg.m; ++token) {
464- const bool active = token < x_active_mask.size() && x_active_mask[token] != 0;
465- if (!active) {
466- continue;
467- }
468- ++active_tokens;
469- for (uint32_t topk = 0; topk < cfg.topk; ++topk) {
470- const size_t slot = static_cast<size_t>(token) * cfg.topk + topk;
471- if (slot >= expert_idx_count) {
472- ++invalid_expert_routes;
473- continue;
474- }
475- const int32_t expert = LoadI32(expert_idx, slot);
476- if (expert < 0 || static_cast<size_t>(expert) >= layout.global_expert_num || cfg.expert_per_rank == 0U) {
477- ++invalid_expert_routes;
478- continue;
479- }
480- const uint32_t dst_rank = static_cast<uint32_t>(expert) / cfg.expert_per_rank;
481- ++valid_routes;
482- if (static_cast<int>(dst_rank) != rank_id) {
483- ++remote_routes;
484- }
485- ++local[layout.per_rank_offset + dst_rank];
486- ++local[layout.per_expert_offset + static_cast<size_t>(expert)];
487- }
488- }
489- local[0] = active_tokens;
490- local[1] = active_tokens * cfg.topk;
491- local[2] = valid_routes;
492- local[3] = remote_routes;
493- local[4] = invalid_expert_routes;
494- local[5] = actual_output_elems;
495- return local;
496-}
497- 
498-std::vector<uint64_t> GatherWorkloadAuditTotals(
499- const std::vector<uint64_t>& local, size_t fields, int rank_id, int world_size)
500-{
501- const int local_bytes = static_cast<int>(local.size() * sizeof(uint64_t));
502- std::vector<uint64_t> gathered(rank_id == 0 ? fields * static_cast<size_t>(world_size) : 0, 0);
503- CommMpiGather(
504- local.data(), local_bytes, COMM_MPI_CHAR, rank_id == 0 ? gathered.data() : nullptr, local_bytes, COMM_MPI_CHAR,
505- 0);
506- 
507- std::vector<uint64_t> total(rank_id == 0 ? fields : 0, 0);
508- if (rank_id != 0) {
509- return total;
510- }
511- for (int rank = 0; rank < world_size; ++rank) {
512- const uint64_t* rank_fields = gathered.data() + static_cast<size_t>(rank) * fields;
513- for (size_t idx = 0; idx < fields; ++idx) {
514- total[idx] += rank_fields[idx];
515- }
516- }
517- return total;
518-}
519- 
520-void PrintWorkloadAuditTotals(
521- const CaseConfig& cfg, const WorkloadAuditLayout& layout, const std::vector<uint64_t>& total, bool skip_accuracy)
522-{
523- const auto dest_begin = total.begin() + static_cast<std::ptrdiff_t>(layout.per_rank_offset);
524- const auto dest_end = dest_begin + cfg.world_size;
525- const auto expert_begin = total.begin() + static_cast<std::ptrdiff_t>(layout.per_expert_offset);
526- const auto expert_end = expert_begin + static_cast<std::ptrdiff_t>(layout.global_expert_num);
527- const uint64_t min_dest_rows = dest_begin == dest_end ? 0 : *std::min_element(dest_begin, dest_end);
528- const uint64_t max_dest_rows = dest_begin == dest_end ? 0 : *std::max_element(dest_begin, dest_end);
529- const uint64_t min_expert_rows = expert_begin == expert_end ? 0 : *std::min_element(expert_begin, expert_end);
530- const uint64_t max_expert_rows = expert_begin == expert_end ? 0 : *std::max_element(expert_begin, expert_end);
531- const size_t nonzero_experts =
532- static_cast<size_t>(std::count_if(expert_begin, expert_end, [](uint64_t rows) { return rows != 0; }));
533- const uint64_t expected_output_elems =
534- static_cast<uint64_t>(cfg.world_size) * static_cast<uint64_t>(cfg.m) * static_cast<uint64_t>(cfg.k);
535- const bool routes_match_case = static_cast<double>(total[0]) == cfg.input_tokens_all_ranks &&
536- static_cast<double>(total[2]) == cfg.routed_tokens_all_ranks &&
537- static_cast<double>(total[3]) == cfg.remote_routed_tokens_all_ranks;
538- const bool output_shape_match = total[5] == expected_output_elems;
539- 
540- std::cout << "[WORKLOAD_AUDIT] dispatch_mega_combine "
541- << (routes_match_case && output_shape_match && total[4] == 0 ? "PASS" : "CHECK")
542- << " accuracy=" << (skip_accuracy ? "SKIP" : "FULL") << '\n';
543- std::cout << " source active_tokens=" << total[0] << " topk_slots=" << total[1] << " valid_routes=" << total[2]
544- << " remote_routes=" << total[3] << " invalid_expert_routes=" << total[4] << '\n';
545- std::cout << " case_json input_tokens=" << static_cast<uint64_t>(cfg.input_tokens_all_ranks)
546- << " routed_tokens=" << static_cast<uint64_t>(cfg.routed_tokens_all_ranks)
547- << " remote_routed_tokens=" << static_cast<uint64_t>(cfg.remote_routed_tokens_all_ranks) << '\n';
548- std::cout << " dest_rows total=" << total[2] << " per_rank_min=" << min_dest_rows
549- << " per_rank_max=" << max_dest_rows << " max_output_size=" << cfg.max_output_size << '\n';
550- std::cout << " expert_rows nonzero=" << nonzero_experts << "/" << layout.global_expert_num
551- << " min=" << min_expert_rows << " max=" << max_expert_rows << '\n';
552- std::cout << " output_elements total=" << total[5] << " expected=" << expected_output_elems << std::endl;
553-}
554- 
555-void PrintWorkloadAuditIfEnabled(
556- const CaseConfig& cfg, int rank_id, int world_size, const std::vector<uint8_t>& expert_idx,
557- const std::vector<uint8_t>& x_active_mask, size_t actual_output_elems, bool skip_accuracy)
558-{
559- const WorkloadAuditLayout layout = BuildWorkloadAuditLayout(cfg);
560- const std::vector<uint64_t> local =
561- BuildLocalWorkloadAudit(cfg, layout, rank_id, expert_idx, x_active_mask, actual_output_elems);
562- const std::vector<uint64_t> total = GatherWorkloadAuditTotals(local, layout.fields, rank_id, world_size);
563- if (rank_id == 0) {
564- PrintWorkloadAuditTotals(cfg, layout, total, skip_accuracy);
565- }
566-}
567- 
568void ValidateFullPathConstraints(const CaseConfig& cfg)364void ValidateFullPathConstraints(const CaseConfig& cfg)
569{365{
570- if (cfg.expert_per_rank + 1U > MEGA_MOE_D2C_MAX_LOGICAL_GROUP_EVENTS) {366+ if (cfg.expert_per_rank != 8U && cfg.expert_per_rank != 16U && cfg.expert_per_rank != 32U) {
571- throw std::runtime_error(367+ throw std::runtime_error("expert_per_rank must be one of 8, 16 or 32");
572- "D2C hard flag budget exceeded: expert_per_rank=" + std::to_string(cfg.expert_per_rank) +
573- " max=" + std::to_string(MEGA_MOE_D2C_MAX_LOGICAL_GROUP_EVENTS));
574- }
575- if (cfg.expert_per_rank > MEGA_MOE_GMM2_TO_COMBINE_MAX_LOGICAL_GROUP_EVENTS) {
576- throw std::runtime_error(
577- "GMM2->combine hard flag budget exceeded: expert_per_rank=" + std::to_string(cfg.expert_per_rank) +
578- " max=" + std::to_string(MEGA_MOE_GMM2_TO_COMBINE_MAX_LOGICAL_GROUP_EVENTS));
579 }368 }
580 if (cfg.k % 128U != 0U) {369 if (cfg.k % 128U != 0U) {
581 throw std::runtime_error("GMM1 requires K % 128 == 0");370 throw std::runtime_error("GMM1 requires K % 128 == 0");
@@ -604,31 +393,25 @@ RunOptions LoadRunOptions()
604 RunOptions options;393 RunOptions options;
605 options.warmup_iters = ParseEnvInt("DISPATCH_MEGA_COMBINE_WARMUP_ITERS", kDefaultWarmupIters);394 options.warmup_iters = ParseEnvInt("DISPATCH_MEGA_COMBINE_WARMUP_ITERS", kDefaultWarmupIters);
606 options.measure_iters = ParseEnvInt("DISPATCH_MEGA_COMBINE_MEASURE_ITERS", kDefaultMeasureIters);395 options.measure_iters = ParseEnvInt("DISPATCH_MEGA_COMBINE_MEASURE_ITERS", kDefaultMeasureIters);
607- options.skip_accuracy = ParseEnvInt("DISPATCH_MEGA_COMBINE_SKIP_ACCURACY", 0) != 0;
608 options.start_sync_debug = ParseEnvInt("DISPATCH_MEGA_COMBINE_START_SYNC_DEBUG", 0) != 0;396 options.start_sync_debug = ParseEnvInt("DISPATCH_MEGA_COMBINE_START_SYNC_DEBUG", 0) != 0;
609- options.workload_audit = ParseEnvInt("DISPATCH_MEGA_COMBINE_WORKLOAD_AUDIT", 0) != 0;
610 if (options.warmup_iters < 0 || options.measure_iters < 0) {397 if (options.warmup_iters < 0 || options.measure_iters < 0) {
611 throw std::runtime_error("warmup/measure iters must be non-negative");398 throw std::runtime_error("warmup/measure iters must be non-negative");
612 }399 }
613 return options;400 return options;
614}401}
615 402 
616-MegaMoeBuildResult BuildAndValidateTiling(const CaseConfig& cfg, const StandaloneRankRuntime& runtime, int rank_id)403+MegaMoeBuildResult BuildAndValidateTiling(const CaseConfig& cfg, const StandaloneRankRuntime& runtime)
617{404{
618 MegaMoeBuildResult build = BuildMegaMoeTiling(cfg, runtime);405 MegaMoeBuildResult build = BuildMegaMoeTiling(cfg, runtime);
619 const auto& front = build.tiling.frontReorderTiling;406 const auto& front = build.tiling.frontReorderTiling;
620 if (!FrontCaseIsSupported(front.frontCase)) {407 if (!FrontCaseIsSupported(front.frontCase)) {
621 throw std::runtime_error("front unsupported case has no legacy fallback");408 throw std::runtime_error("front unsupported case has no legacy fallback");
622 }409 }
623- build.tiling.combineTiling.combineImplMode = kHostCombineImplDirectAuto;
624 ValidateFullPathConstraints(cfg);410 ValidateFullPathConstraints(cfg);
625- Trace(
626- rank_id,
627- "tiling built frontCase=" + std::to_string(front.frontCase) + " block_dim=" + std::to_string(build.block_dim));
628 return build;411 return build;
629}412}
630 413 
631-RankHostInputs LoadRankHostInputs(const RankFileSet& files, bool skip_accuracy)414+RankHostInputs LoadRankHostInputs(const RankFileSet& files)
632{415{
633 RankHostInputs inputs;416 RankHostInputs inputs;
634 inputs.x = ReadBinaryFile(files.x);417 inputs.x = ReadBinaryFile(files.x);
@@ -638,10 +421,7 @@ RankHostInputs LoadRankHostInputs(const RankFileSet& files, bool skip_accuracy)
638 inputs.scale1 = ReadBinaryFile(files.scale1);421 inputs.scale1 = ReadBinaryFile(files.scale1);
639 inputs.scale2 = ReadBinaryFile(files.scale2);422 inputs.scale2 = ReadBinaryFile(files.scale2);
640 inputs.probs = ReadBinaryFile(files.probs);423 inputs.probs = ReadBinaryFile(files.probs);
641- inputs.x_active_mask = ReadBinaryFile(files.x_active_mask);424+ inputs.expected_out = BytesToU16(ReadBinaryFile(files.expected_out));
642- if (!skip_accuracy) {
643- inputs.expected_out = BytesToU16(ReadBinaryFile(files.expected_out));
644- }
645 return inputs;425 return inputs;
646}426}
647 427 
@@ -677,56 +457,53 @@ MegaMoeLaunchArgs BuildLaunchArgs(
677 MegaMoeLaunchArgs args;457 MegaMoeLaunchArgs args;
678 args.ffts = reinterpret_cast<void*>(ffts_addr);458 args.ffts = reinterpret_cast<void*>(ffts_addr);
679 args.block_dim = build.block_dim;459 args.block_dim = build.block_dim;
680- args.tiling = buffers.tiling.ptr;460+ args.tiling = buffers.tiling.data();
681- args.workspace = buffers.workspace.ptr;461+ args.workspace = buffers.workspace.data();
682- args.x = buffers.x.ptr;462+ args.x = buffers.x.data();
683- args.weight1 = buffers.weight1.ptr;463+ args.weight1 = buffers.weight1.data();
684- args.weight2 = buffers.weight2.ptr;464+ args.weight2 = buffers.weight2.data();
685- args.expert_idx = buffers.expert_idx.ptr;465+ args.expert_idx = buffers.expert_idx.data();
686- args.scale1 = buffers.scale1.ptr;466+ args.scale1 = buffers.scale1.data();
687- args.scale2 = buffers.scale2.ptr;467+ args.scale2 = buffers.scale2.data();
688- args.probs = buffers.probs.ptr;468+ args.probs = buffers.probs.data();
689- args.out = buffers.out.ptr;469+ args.out = buffers.out.data();
690- args.expert_token_nums = buffers.expert_token_nums.ptr;470+ args.expert_token_nums = buffers.expert_token_nums.data();
691- args.profile_data = buffers.profile.ptr;471+ args.profile_data = buffers.profile.data();
692 args.start_sync_debug = start_sync_debug ? 1U : 0U;472 args.start_sync_debug = start_sync_debug ? 1U : 0U;
693 return args;473 return args;
694}474}
695 475 
696-void LaunchAndSync(int rank_id, const MegaMoeLaunchArgs& args, aclrtStream stream, const char* trace_tag)476+void LaunchAndSync(const MegaMoeLaunchArgs& args, aclrtStream stream)
697{477{
698- Trace(rank_id, std::string(trace_tag) + " launch begin");
699 launchMegaMoe(args, stream);478 launchMegaMoe(args, stream);
700- Trace(rank_id, std::string(trace_tag) + " launch submitted");
701 if (aclrtSynchronizeStream(stream) != ACL_SUCCESS) {479 if (aclrtSynchronizeStream(stream) != ACL_SUCCESS) {
702 throw std::runtime_error("stream sync failed");480 throw std::runtime_error("stream sync failed");
703 }481 }
704- Trace(rank_id, std::string(trace_tag) + " launch synced");
705}482}
706 483 
707void RunWarmupIterations(484void RunWarmupIterations(
708 const StandaloneRankRuntime& runtime, const RankDeviceBuffers& buffers, const MegaMoeLaunchArgs& args,485 const StandaloneRankRuntime& runtime, const RankDeviceBuffers& buffers, const MegaMoeLaunchArgs& args,
709- int warmup_iters, int rank_id)486+ int warmup_iters)
710{487{
711 CommMpiBarrier();488 CommMpiBarrier();
712 for (int iter = 0; iter < warmup_iters; ++iter) {489 for (int iter = 0; iter < warmup_iters; ++iter) {
713 PrepareIterationState(runtime, buffers.out, buffers.expert_token_nums, buffers.workspace, buffers.profile);490 PrepareIterationState(runtime, buffers.out, buffers.expert_token_nums, buffers.workspace, buffers.profile);
714 CommMpiBarrier();491 CommMpiBarrier();
715- LaunchAndSync(rank_id, args, runtime.compute_stream, "warmup");492+ LaunchAndSync(args, runtime.compute_stream);
716 CommMpiBarrier();493 CommMpiBarrier();
717 }494 }
718}495}
719 496 
720std::vector<double> RunMeasureIterations(497std::vector<double> RunMeasureIterations(
721 const StandaloneRankRuntime& runtime, RankDeviceBuffers& buffers, const MegaMoeLaunchArgs& args,498 const StandaloneRankRuntime& runtime, RankDeviceBuffers& buffers, const MegaMoeLaunchArgs& args,
722- const MegaMoeBuildResult& build, int measure_iters, int rank_id)499+ const MegaMoeBuildResult& build, int measure_iters)
723{500{
724 std::vector<double> kernel_times_us;501 std::vector<double> kernel_times_us;
725 kernel_times_us.reserve(static_cast<size_t>(measure_iters));502 kernel_times_us.reserve(static_cast<size_t>(measure_iters));
726 for (int iter = 0; iter < measure_iters; ++iter) {503 for (int iter = 0; iter < measure_iters; ++iter) {
727 PrepareIterationState(runtime, buffers.out, buffers.expert_token_nums, buffers.workspace, buffers.profile);504 PrepareIterationState(runtime, buffers.out, buffers.expert_token_nums, buffers.workspace, buffers.profile);
728 CommMpiBarrier();505 CommMpiBarrier();
729- LaunchAndSync(rank_id, args, runtime.compute_stream, "measure");506+ LaunchAndSync(args, runtime.compute_stream);
730 CommMpiBarrier();507 CommMpiBarrier();
731 kernel_times_us.push_back(ReadKernelProfileUs(buffers.profile, buffers.profile_host, build.block_dim));508 kernel_times_us.push_back(ReadKernelProfileUs(buffers.profile, buffers.profile_host, build.block_dim));
732 }509 }
@@ -737,8 +514,8 @@ std::vector<uint16_t> CopyActualOutputToHost(const CaseConfig& cfg, const Device
737{514{
738 std::vector<uint16_t> actual_out(static_cast<size_t>(cfg.m) * cfg.k);515 std::vector<uint16_t> actual_out(static_cast<size_t>(cfg.m) * cfg.k);
739 if (aclrtMemcpy(516 if (aclrtMemcpy(
740- actual_out.data(), actual_out.size() * sizeof(uint16_t), out_dev.ptr, actual_out.size() * sizeof(uint16_t),517+ actual_out.data(), actual_out.size() * sizeof(uint16_t), out_dev.data(),
741- ACL_MEMCPY_DEVICE_TO_HOST) != ACL_SUCCESS) {518+ actual_out.size() * sizeof(uint16_t), ACL_MEMCPY_DEVICE_TO_HOST) != ACL_SUCCESS) {
742 throw std::runtime_error("device->host output copy failed");519 throw std::runtime_error("device->host output copy failed");
743 }520 }
744 return actual_out;521 return actual_out;
@@ -746,14 +523,8 @@ std::vector<uint16_t> CopyActualOutputToHost(const CaseConfig& cfg, const Device
746 523 
747bool ReportRankAccuracy(524bool ReportRankAccuracy(
748 int rank_id, int world_size, const CaseConfig& cfg, const RankHostInputs& inputs,525 int rank_id, int world_size, const CaseConfig& cfg, const RankHostInputs& inputs,
749- const std::vector<uint16_t>& actual_out, bool skip_accuracy)526+ const std::vector<uint16_t>& actual_out)
750{527{
751- if (skip_accuracy) {
752- PrintOrderedByRank(
753- rank_id, world_size,
754- "rank=" + std::to_string(rank_id) + " accuracy=SKIP\nPASS rank=" + std::to_string(rank_id));
755- return true;
756- }
757 const AccuracyReport report = CompareFp16File(inputs.expected_out, actual_out, cfg.compare_atol, cfg.compare_rtol);528 const AccuracyReport report = CompareFp16File(inputs.expected_out, actual_out, cfg.compare_atol, cfg.compare_rtol);
758 PrintOrderedByRank(529 PrintOrderedByRank(
759 rank_id, world_size,530 rank_id, world_size,
@@ -768,22 +539,20 @@ bool RunOneRank(int rank_id, int world_size, const std::string& case_dir, const
768 if (!InitStandaloneRankRuntime(runtime, rank_id, world_size, root_info)) {539 if (!InitStandaloneRankRuntime(runtime, rank_id, world_size, root_info)) {
769 return false;540 return false;
770 }541 }
771- Trace(rank_id, "runtime initialized");
772 542 
773 bool ok = false;543 bool ok = false;
774 try {544 try {
775 const RunOptions options = LoadRunOptions();545 const RunOptions options = LoadRunOptions();
776 const CaseConfig cfg = LoadCaseConfig(case_dir + "/case.json");546 const CaseConfig cfg = LoadCaseConfig(case_dir + "/case.json");
777 const RankFileSet files = BuildRankFileSet(case_dir, rank_id);547 const RankFileSet files = BuildRankFileSet(case_dir, rank_id);
778- const MegaMoeBuildResult build = BuildAndValidateTiling(cfg, runtime, rank_id);548+ const MegaMoeBuildResult build = BuildAndValidateTiling(cfg, runtime);
779- const RankHostInputs inputs = LoadRankHostInputs(files, options.skip_accuracy);549+ const RankHostInputs inputs = LoadRankHostInputs(files);
780 RankDeviceBuffers buffers = AllocateRankDeviceBuffers(cfg, build, inputs);550 RankDeviceBuffers buffers = AllocateRankDeviceBuffers(cfg, build, inputs);
781- Trace(rank_id, "buffers allocated");
782 551 
783 const MegaMoeLaunchArgs args = BuildLaunchArgs(build, buffers, options.start_sync_debug);552 const MegaMoeLaunchArgs args = BuildLaunchArgs(build, buffers, options.start_sync_debug);
784- RunWarmupIterations(runtime, buffers, args, options.warmup_iters, rank_id);553+ RunWarmupIterations(runtime, buffers, args, options.warmup_iters);
785 const std::vector<double> kernel_times_us =554 const std::vector<double> kernel_times_us =
786- RunMeasureIterations(runtime, buffers, args, build, options.measure_iters, rank_id);555+ RunMeasureIterations(runtime, buffers, args, build, options.measure_iters);
787 556 
788 const std::vector<double> kernel_max_samples = GatherMaxSamplesToRoot(kernel_times_us, rank_id, world_size);557 const std::vector<double> kernel_max_samples = GatherMaxSamplesToRoot(kernel_times_us, rank_id, world_size);
789 if (rank_id == 0) {558 if (rank_id == 0) {
@@ -792,19 +561,14 @@ bool RunOneRank(int rank_id, int world_size, const std::string& case_dir, const
792 561 
793 PrepareIterationState(runtime, buffers.out, buffers.expert_token_nums, buffers.workspace, buffers.profile);562 PrepareIterationState(runtime, buffers.out, buffers.expert_token_nums, buffers.workspace, buffers.profile);
794 CommMpiBarrier();563 CommMpiBarrier();
795- LaunchAndSync(rank_id, args, runtime.compute_stream, "final");564+ LaunchAndSync(args, runtime.compute_stream);
796 CommMpiBarrier();565 CommMpiBarrier();
797 566 
798 const std::vector<uint16_t> actual_out = CopyActualOutputToHost(cfg, buffers.out);567 const std::vector<uint16_t> actual_out = CopyActualOutputToHost(cfg, buffers.out);
799- if (options.workload_audit) {
800- PrintWorkloadAuditIfEnabled(
801- cfg, rank_id, world_size, inputs.expert_idx, inputs.x_active_mask, actual_out.size(),
802- options.skip_accuracy);
803- }
804 WriteBinaryFile(568 WriteBinaryFile(
805 case_dir + "/output_rank" + std::to_string(rank_id) + ".bin", actual_out.data(),569 case_dir + "/output_rank" + std::to_string(rank_id) + ".bin", actual_out.data(),
806 actual_out.size() * sizeof(uint16_t));570 actual_out.size() * sizeof(uint16_t));
807- ok = ReportRankAccuracy(rank_id, world_size, cfg, inputs, actual_out, options.skip_accuracy);571+ ok = ReportRankAccuracy(rank_id, world_size, cfg, inputs, actual_out);
808 } catch (const std::exception& ex) {572 } catch (const std::exception& ex) {
809 std::cerr << "rank=" << rank_id << " error: " << ex.what() << std::endl;573 std::cerr << "rank=" << rank_id << " error: " << ex.what() << std::endl;
810 ok = false;574 ok = false;
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/combine.h+154-272
@@ -21,38 +21,24 @@ See LICENSE in the root of the software repository for the full text of the Lice
21#include "utils/common_helpers.hpp"21#include "utils/common_helpers.hpp"
22#include "utils/const_args.hpp"22#include "utils/const_args.hpp"
23#include "utils/hccl_window.hpp"23#include "utils/hccl_window.hpp"
24-#include "utils/pto_sync_substrate.hpp"24+#include "utils/mega_expert_sync.hpp"
25#include "utils/pto_vector.hpp"25#include "utils/pto_vector.hpp"
26 26 
27constexpr uint32_t kCombineVecTileElems = 8192U;27constexpr uint32_t kCombineVecTileElems = 8192U;
28constexpr uint32_t kCombineBufferNum = 2U;28constexpr uint32_t kCombineBufferNum = 2U;
29-constexpr uint32_t kCombineInvalidTask = 0xFFFFFFFFU;
30-constexpr uint32_t kCombineSmallTokenThreshold = 4096U;
31-constexpr uint32_t kCombineSmallTokenSubtileRows = 16U;
32-constexpr uint32_t kCombineSmallTokenSubtileCols = 256U;
33-constexpr uint32_t kCombineSmallMaxElems = 8192U;
34-constexpr uint32_t kCombineSmallScaleElems = kCombineSmallMaxElems / kCombineSmallTokenSubtileCols;
35-constexpr uint32_t kCombineImplDirectLarge = 1U;
36-constexpr uint32_t kCombineImplDirectSmall = 2U;
37-constexpr uint32_t kCombineImplDirectAuto = 3U;
38-constexpr uint32_t kCombineDirectLargeUbStages = 2U;
39-constexpr uint32_t kCombineDirectSmallUbStages = 2U;
40-constexpr uint32_t kCombineLargeLanesPerRank = 2U;
41 29 
42template <typename OutputElement>30template <typename OutputElement>
43class Combine {31class Combine {
44public:32public:
45 AICORE inline void Init(GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData);33 AICORE inline void Init(GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData);
46- AICORE inline void Process();34+ AICORE inline void ProcessFixed(uint32_t groupLocalId, uint32_t groupSize);
35+ AICORE inline void ProcessFixedFinalBoundary(uint32_t role, uint32_t flatAivId, bool combineActive);
47 36 
48private:37private:
49 static_assert(38 static_assert(
50 std::is_same_v<OutputElement, half> || std::is_same_v<OutputElement, bfloat16_t>,39 std::is_same_v<OutputElement, half> || std::is_same_v<OutputElement, bfloat16_t>,
51 "combine output must be half or bfloat16");40 "combine output must be half or bfloat16");
52 41 
53- using VectorShape = pto::Shape<1, 1, 1, 1, pto::DYNAMIC>;
54- using VectorStride = pto::Stride<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, 1>;
55- using ScaleGlobal = pto::GlobalTensor<float, VectorShape, VectorStride, pto::Layout::ND>;
56 using BlockShape = pto::Shape<1, 1, 1, pto::DYNAMIC, pto::DYNAMIC>;42 using BlockShape = pto::Shape<1, 1, 1, pto::DYNAMIC, pto::DYNAMIC>;
57 using BlockStride = pto::Stride<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, 1>;43 using BlockStride = pto::Stride<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, 1>;
58 using CBlockGlobal = pto::GlobalTensor<half, BlockShape, BlockStride, pto::Layout::ND>;44 using CBlockGlobal = pto::GlobalTensor<half, BlockShape, BlockStride, pto::Layout::ND>;
@@ -60,32 +46,9 @@ private:
60 using TileC = pto::Tile<pto::TileType::Vec, half, 1, kCombineVecTileElems, pto::BLayout::RowMajor, -1, -1>;46 using TileC = pto::Tile<pto::TileType::Vec, half, 1, kCombineVecTileElems, pto::BLayout::RowMajor, -1, -1>;
61 using TileFp32 = pto::Tile<pto::TileType::Vec, float, 1, kCombineVecTileElems, pto::BLayout::RowMajor, -1, -1>;47 using TileFp32 = pto::Tile<pto::TileType::Vec, float, 1, kCombineVecTileElems, pto::BLayout::RowMajor, -1, -1>;
62 using TileD = pto::Tile<pto::TileType::Vec, OutputElement, 1, kCombineVecTileElems, pto::BLayout::RowMajor, -1, -1>;48 using TileD = pto::Tile<pto::TileType::Vec, OutputElement, 1, kCombineVecTileElems, pto::BLayout::RowMajor, -1, -1>;
63- using SmallTileC = pto::Tile<49+ AICORE inline bool RankStreamingEnabled() const
64- pto::TileType::Vec, half, kCombineSmallTokenSubtileRows, kCombineSmallTokenSubtileCols, pto::BLayout::RowMajor,
65- -1, -1>;
66- using SmallTileFp32 = pto::Tile<
67- pto::TileType::Vec, float, kCombineSmallTokenSubtileRows, kCombineSmallTokenSubtileCols, pto::BLayout::RowMajor,
68- -1, -1>;
69- using SmallTileD = pto::Tile<
70- pto::TileType::Vec, OutputElement, kCombineSmallTokenSubtileRows, kCombineSmallTokenSubtileCols,
71- pto::BLayout::RowMajor, -1, -1>;
72- 
73- AICORE inline uint64_t TokenVolume() const { return static_cast<uint64_t>(problemM_) * topK_; }
74- AICORE inline bool IsSmallTokenPath() const { return TokenVolume() <= kCombineSmallTokenThreshold; }
75- AICORE inline uint32_t CombineImplMode() const { return tilingData_->combineTiling.combineImplMode; }
76- AICORE inline bool DirectLargeEnabled() const
77 {50 {
78- const uint32_t mode = CombineImplMode();51+ return tilingData_->unpermuteTiling.unpermuteImplMode == kMegaMoeUnpermuteImplRankStreaming;
79- return (mode == kCombineImplDirectLarge || mode == kCombineImplDirectAuto) && !IsSmallTokenPath();
80- }
81- AICORE inline bool DirectSmallEnabled() const
82- {
83- const uint32_t mode = CombineImplMode();
84- return (mode == kCombineImplDirectSmall || mode == kCombineImplDirectAuto) && IsSmallTokenPath();
85- }
86- AICORE inline uint16_t Gmm2ToCombineFlagId(uint32_t groupIdx) const
87- {
88- return MEGA_MOE_GMM2_TO_COMBINE_HARD_FLAG_BASE + groupIdx / CROSS_CORE_FLAG_MAX_SET_COUNT;
89 }52 }
90 AICORE inline uint32_t CurrentM(uint32_t groupIdx) const53 AICORE inline uint32_t CurrentM(uint32_t groupIdx) const
91 {54 {
@@ -125,7 +88,23 @@ private:
125 if (lanes == 0U) {88 if (lanes == 0U) {
126 lanes = 1U;89 lanes = 1U;
127 }90 }
128- return lanes > kCombineLargeLanesPerRank ? kCombineLargeLanesPerRank : lanes;91+ const uint32_t configuredLanes = tilingData_->fixedGroupTiling.combineLargeLanesPerRank;
92+ return lanes > configuredLanes ? configuredLanes : lanes;
93+ }
94+ AICORE inline uint32_t DirectLargeTaskCount() const { return rankSize_ * LargeLanesPerRank(); }
95+ AICORE inline uint32_t DirectLargeWorkerCount() const
96+ {
97+ const uint32_t taskCount = DirectLargeTaskCount();
98+ return coreNum_ < taskCount ? coreNum_ : taskCount;
99+ }
100+ AICORE inline uint32_t ReadyCoordinatorCore() const
101+ {
102+ const uint32_t workerCount = DirectLargeWorkerCount();
103+ if (workerCount == 0U) {
104+ return 0U;
105+ }
106+ const uint64_t firstLocalTask = static_cast<uint64_t>(rank_) * LargeLanesPerRank();
107+ return static_cast<uint32_t>(firstLocalTask % workerCount);
129 }108 }
130 AICORE inline uint32_t LargeLaneRowBegin(uint32_t rows, uint32_t laneIdx, uint32_t lanesPerRank) const109 AICORE inline uint32_t LargeLaneRowBegin(uint32_t rows, uint32_t laneIdx, uint32_t lanesPerRank) const
131 {110 {
@@ -145,26 +124,18 @@ private:
145 AICORE inline void InitUbLayout();124 AICORE inline void InitUbLayout();
146 AICORE inline void SetInitialFlags() const;125 AICORE inline void SetInitialFlags() const;
147 AICORE inline void FinalizeLocalPipe() const;126 AICORE inline void FinalizeLocalPipe() const;
127+ AICORE inline void FinalizeExpertStores() const;
148 AICORE inline uint32_t TokenPerExpertResetElems() const;128 AICORE inline uint32_t TokenPerExpertResetElems() const;
149- AICORE inline bool ResetTokenPerExpert(uint32_t elems) const;129+ AICORE inline bool ResetTokenPerExpertByOwner(uint32_t elems, bool resetOwner) const;
150- AICORE inline void ProcessFinalBoundary();130+ AICORE inline void PublishAssignedExpertProgress(uint32_t readyExpertCount) const;
151- AICORE inline void WaitGmm2Ready(uint32_t groupIdx, bool aivSyncAfterWait) const;131+ AICORE inline void FinalizeRankStreamingLane();
132+ AICORE inline uint32_t Gmm2ProducerCount(uint32_t groupIdx) const;
133+ AICORE inline void WaitGmm2Ready(uint32_t groupIdx) const;
152 AICORE inline void ProcessDirectLargeSegmentRows(134 AICORE inline void ProcessDirectLargeSegmentRows(
153 uint32_t srcRank, uint32_t srcRowOffset, uint32_t rows, uint32_t dstRowOffset);135 uint32_t srcRank, uint32_t srcRowOffset, uint32_t rows, uint32_t dstRowOffset);
154 AICORE inline void ProcessDirectLargeTokenPath();136 AICORE inline void ProcessDirectLargeTokenPath();
155- AICORE inline void LoadSmallSubtile(
156- uint32_t bufferId, uint32_t srcRowOffset, uint32_t rowNum, uint32_t colBegin, uint32_t colNum) const;
157- AICORE inline void DequantDirectSmallSubtile(uint32_t bufferId, uint32_t rowNum, uint32_t colNum);
158- AICORE inline void StoreSmallSubtileIntersection(
159- uint32_t bufferId, __gm__ OutputElement* dstBase, uint32_t dstRowOffset, uint32_t ubRowOffset, uint32_t rowNum,
160- uint32_t colBegin, uint32_t colNum);
161- AICORE inline void StoreSmallSubtileToRanks(
162- uint32_t groupIdx, const GmmCommonTileInfo& tileInfo, uint32_t tileRowBegin, uint32_t rows, uint32_t bufferId);
163- AICORE inline void ProcessDirectSmallTile(
164- uint32_t groupIdx, uint32_t groupBase, const GmmCommonTileInfo& tileInfo, uint32_t subtileBegin,
165- uint32_t subtileCount);
166- AICORE inline void ProcessDirectSmallTokenPath();
167 137 
138+ GM_ADDR workspaceGM_ = nullptr;
168 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;139 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;
169 140 
170 PtoRemoteWindow remoteWindow_;141 PtoRemoteWindow remoteWindow_;
@@ -175,9 +146,7 @@ private:
175 __gm__ int32_t* preSumBeforeRankPtr_ = nullptr;146 __gm__ int32_t* preSumBeforeRankPtr_ = nullptr;
176 __gm__ int32_t* tokenPerExpertPtr_ = nullptr;147 __gm__ int32_t* tokenPerExpertPtr_ = nullptr;
177 148 
178- uint32_t problemM_ = 0;
179 uint32_t problemK_ = 0;149 uint32_t problemK_ = 0;
180- uint32_t topK_ = 0;
181 uint32_t maxOutputSize_ = 0;150 uint32_t maxOutputSize_ = 0;
182 uint32_t expertPerRank_ = 0;151 uint32_t expertPerRank_ = 0;
183 uint32_t expertNumAligned_ = 0;152 uint32_t expertNumAligned_ = 0;
@@ -186,22 +155,20 @@ private:
186 uint32_t coreIdx_ = 0;155 uint32_t coreIdx_ = 0;
187 uint32_t coreNum_ = 1;156 uint32_t coreNum_ = 1;
188 uint32_t pingpongId_ = 0;157 uint32_t pingpongId_ = 0;
158+ int32_t dataReadyEpoch_ = 0;
189 uint64_t ubCOffset_[kCombineBufferNum] = {0, 0};159 uint64_t ubCOffset_[kCombineBufferNum] = {0, 0};
190 uint64_t ubFp32Offset_[kCombineBufferNum] = {0, 0};160 uint64_t ubFp32Offset_[kCombineBufferNum] = {0, 0};
191 uint64_t ubDOffset_[kCombineBufferNum] = {0, 0};161 uint64_t ubDOffset_[kCombineBufferNum] = {0, 0};
192- uint64_t ubScaleOffset_[kCombineBufferNum] = {0, 0};
193- mutable uint32_t smallScaleSourceOffset_[kCombineBufferNum] = {kCombineInvalidTask, kCombineInvalidTask};
194};162};
195 163 
196template <typename OutputElement>164template <typename OutputElement>
197AICORE inline void Combine<OutputElement>::Init(GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData)165AICORE inline void Combine<OutputElement>::Init(GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData)
198{166{
167+ workspaceGM_ = workspaceGM;
199 tilingData_ = tilingData;168 tilingData_ = tilingData;
200 pingpongId_ = 0;169 pingpongId_ = 0;
201 170 
202- problemM_ = tilingData_->megaMoeInfo.M;
203 problemK_ = tilingData_->megaMoeInfo.K;171 problemK_ = tilingData_->megaMoeInfo.K;
204- topK_ = tilingData_->megaMoeInfo.topK;
205 maxOutputSize_ = tilingData_->megaMoeInfo.maxOutputSize;172 maxOutputSize_ = tilingData_->megaMoeInfo.maxOutputSize;
206 expertPerRank_ = tilingData_->megaMoeInfo.expertPerRank;173 expertPerRank_ = tilingData_->megaMoeInfo.expertPerRank;
207 expertNumAligned_ = tilingData_->frontReorderTiling.expertNumAligned;174 expertNumAligned_ = tilingData_->frontReorderTiling.expertNumAligned;
@@ -224,6 +191,7 @@ AICORE inline void Combine<OutputElement>::Init(GM_ADDR workspaceGM, const __gm_
224 reinterpret_cast<__gm__ int32_t*>(workspaceGM + tilingData_->frontReorderTiling.preSumBeforeRankOffset);191 reinterpret_cast<__gm__ int32_t*>(workspaceGM + tilingData_->frontReorderTiling.preSumBeforeRankOffset);
225 tokenPerExpertPtr_ =192 tokenPerExpertPtr_ =
226 reinterpret_cast<__gm__ int32_t*>(remoteWindow_.LocalBase() + peerMemoryLayout_.offsetPeerTokenPerExpert);193 reinterpret_cast<__gm__ int32_t*>(remoteWindow_.LocalBase() + peerMemoryLayout_.offsetPeerTokenPerExpert);
194+ dataReadyEpoch_ = RankStreamingEnabled() ? remoteWindow_.DataReadyEpoch() : 0;
227 195 
228 InitUbLayout();196 InitUbLayout();
229}197}
@@ -232,30 +200,13 @@ template <typename OutputElement>
232AICORE inline void Combine<OutputElement>::InitUbLayout()200AICORE inline void Combine<OutputElement>::InitUbLayout()
233{201{
234 uint64_t ubOffset = 0;202 uint64_t ubOffset = 0;
235- if (DirectLargeEnabled()) {
236- for (uint32_t i = 0; i < kCombineBufferNum; ++i) {
237- ubCOffset_[i] = ubOffset;
238- ubOffset += alignUp(static_cast<uint64_t>(problemK_) * sizeof(half), UB_ALIGN);
239- ubDOffset_[i] = ubOffset;
240- ubOffset += alignUp(static_cast<uint64_t>(problemK_) * sizeof(OutputElement), UB_ALIGN);
241- ubFp32Offset_[i] = ubOffset;
242- ubOffset += alignUp(static_cast<uint64_t>(problemK_) * sizeof(float), UB_ALIGN);
243- ubScaleOffset_[i] = 0U;
244- smallScaleSourceOffset_[i] = kCombineInvalidTask;
245- }
246- return;
247- }
248- 
249 for (uint32_t i = 0; i < kCombineBufferNum; ++i) {203 for (uint32_t i = 0; i < kCombineBufferNum; ++i) {
250 ubCOffset_[i] = ubOffset;204 ubCOffset_[i] = ubOffset;
251- ubOffset += alignUp(static_cast<uint64_t>(kCombineSmallMaxElems) * sizeof(half), UB_ALIGN);205+ ubOffset += alignUp(static_cast<uint64_t>(problemK_) * sizeof(half), UB_ALIGN);
252 ubDOffset_[i] = ubOffset;206 ubDOffset_[i] = ubOffset;
253- ubOffset += alignUp(static_cast<uint64_t>(kCombineSmallMaxElems) * sizeof(OutputElement), UB_ALIGN);207+ ubOffset += alignUp(static_cast<uint64_t>(problemK_) * sizeof(OutputElement), UB_ALIGN);
254 ubFp32Offset_[i] = ubOffset;208 ubFp32Offset_[i] = ubOffset;
255- ubOffset += alignUp(static_cast<uint64_t>(kCombineSmallMaxElems) * sizeof(float), UB_ALIGN);209+ ubOffset += alignUp(static_cast<uint64_t>(problemK_) * sizeof(float), UB_ALIGN);
256- ubScaleOffset_[i] = ubOffset;
257- ubOffset += alignUp(static_cast<uint64_t>(kCombineSmallScaleElems) * sizeof(float), UB_ALIGN);
258- smallScaleSourceOffset_[i] = kCombineInvalidTask;
259 }210 }
260}211}
261 212 
@@ -277,6 +228,19 @@ AICORE inline void Combine<OutputElement>::FinalizeLocalPipe() const
277 }228 }
278}229}
279 230 
231+template <typename OutputElement>
232+AICORE inline void Combine<OutputElement>::FinalizeExpertStores() const
233+{
234+ for (uint32_t i = 0; i < kCombineBufferNum; ++i) {
235+ wait_flag(PIPE_MTE3, PIPE_V, StoreFreeEvent(i));
236+ }
237+ for (uint32_t i = 0; i < kCombineBufferNum; ++i) {
238+ set_flag(PIPE_MTE3, PIPE_V, StoreFreeEvent(i));
239+ }
240+ pipe_barrier(PIPE_ALL);
241+ dsb(DSB_DDR);
242+}
243+ 
280template <typename OutputElement>244template <typename OutputElement>
281AICORE inline uint32_t Combine<OutputElement>::TokenPerExpertResetElems() const245AICORE inline uint32_t Combine<OutputElement>::TokenPerExpertResetElems() const
282{246{
@@ -284,9 +248,9 @@ AICORE inline uint32_t Combine<OutputElement>::TokenPerExpertResetElems() const
284}248}
285 249 
286template <typename OutputElement>250template <typename OutputElement>
287-AICORE inline bool Combine<OutputElement>::ResetTokenPerExpert(uint32_t elems) const251+AICORE inline bool Combine<OutputElement>::ResetTokenPerExpertByOwner(uint32_t elems, bool resetOwner) const
288{252{
289- if (coreIdx_ != coreNum_ - 1U) {253+ if (!resetOwner) {
290 return false;254 return false;
291 }255 }
292 PtoFillUb<int32_t>(0U, 0, elems);256 PtoFillUb<int32_t>(0U, 0, elems);
@@ -298,20 +262,91 @@ AICORE inline bool Combine<OutputElement>::ResetTokenPerExpert(uint32_t elems) c
298}262}
299 263 
300template <typename OutputElement>264template <typename OutputElement>
301-AICORE inline void Combine<OutputElement>::ProcessFinalBoundary()265+AICORE inline void Combine<OutputElement>::ProcessFixedFinalBoundary(
266+ uint32_t role, uint32_t flatAivId, bool combineActive)
302{267{
303- FinalizeLocalPipe();268+ if (role == kMegaMoeFixedRoleCombine && combineActive) {
269+ FinalizeLocalPipe();
270+ }
304 pto::SYNCALL<pto::SyncCoreType::AIVOnly>();271 pto::SYNCALL<pto::SyncCoreType::AIVOnly>();
305- (void)ResetTokenPerExpert(TokenPerExpertResetElems());272+ ResetTokenPerExpertByOwner(TokenPerExpertResetElems(), flatAivId + 1U == kMegaMoeFixedPhysicalAivNum);
306 remoteWindow_.CrossRankSync();273 remoteWindow_.CrossRankSync();
307}274}
308 275 
309template <typename OutputElement>276template <typename OutputElement>
310-AICORE inline void Combine<OutputElement>::WaitGmm2Ready(uint32_t groupIdx, bool aivSyncAfterWait) const277+AICORE inline void Combine<OutputElement>::PublishAssignedExpertProgress(uint32_t readyExpertCount) const
311{278{
312- CrossCoreWaitFlag<0x2>(Gmm2ToCombineFlagId(groupIdx));279+ const uint32_t lanesPerRank = LargeLanesPerRank();
313- if (aivSyncAfterWait) {280+ const uint32_t taskCount = DirectLargeTaskCount();
314- pto::SYNCALL<pto::SyncCoreType::AIVOnly>();281+ for (uint32_t taskIdx = coreIdx_; taskIdx < taskCount; taskIdx += coreNum_) {
282+ const uint32_t ownerRank = taskIdx / lanesPerRank;
283+ remoteWindow_.PublishExpertProgress(static_cast<int32_t>(ownerRank), readyExpertCount, dataReadyEpoch_);
284+ }
285+}
286+ 
287+template <typename OutputElement>
288+AICORE inline void Combine<OutputElement>::FinalizeRankStreamingLane()
289+{
290+ const uint32_t taskCount = DirectLargeTaskCount();
291+ const uint32_t combineWorkerCount = DirectLargeWorkerCount();
292+ if (coreIdx_ >= combineWorkerCount) {
293+ return;
294+ }
295+ 
296+ remoteWindow_.AcquireDataReady();
297+ const uint32_t lanesPerRank = LargeLanesPerRank();
298+ for (uint32_t taskIdx = coreIdx_; taskIdx < taskCount; taskIdx += coreNum_) {
299+ const uint32_t ownerRank = taskIdx / lanesPerRank;
300+ remoteWindow_.PublishExpertProgress(static_cast<int32_t>(ownerRank), expertPerRank_, dataReadyEpoch_);
301+ remoteWindow_.PublishDataReady(static_cast<int32_t>(ownerRank), dataReadyEpoch_);
302+ }
303+ remoteWindow_.PublishLocalCombineDone(coreIdx_, dataReadyEpoch_);
304+ 
305+ if (coreIdx_ != 0U) {
306+ return;
307+ }
308+ remoteWindow_.WaitLocalCombineDoneMte(combineWorkerCount, dataReadyEpoch_);
309+ ResetTokenPerExpertByOwner(TokenPerExpertResetElems(), true);
310+ const uint32_t workerCount = tilingData_->fixedGroupTiling.physicalAivNum;
311+ const uint32_t initialWorkerCount = tilingData_->fixedGroupTiling.gmm1GroupSize * 2U;
312+ const uint32_t helperCount = workerCount > initialWorkerCount ? workerCount - initialWorkerCount : 0U;
313+ remoteWindow_.PublishUnpermuteStartRangeMte(initialWorkerCount, helperCount, dataReadyEpoch_);
314+}
315+ 
316+template <typename OutputElement>
317+AICORE inline uint32_t Combine<OutputElement>::Gmm2ProducerCount(uint32_t groupIdx) const
318+{
319+ const __gm__ MegaMoeFixedGroupTiling& fixed = tilingData_->fixedGroupTiling;
320+ if (groupIdx < fixed.gmm2JoinCheckStartExpert) {
321+ return fixed.gmm2GroupSize;
322+ }
323+ 
324+ const int32_t decision = WaitEpochAcquire(
325+ FixedSyncSlot(workspaceGM_, tilingData_, FixedSyncLayout(tilingData_).gmm2JoinSlot),
326+ static_cast<int32_t>(groupIdx + 1U));
327+ const uint32_t encodedExpert = static_cast<uint32_t>(decision & kMegaMoeFixedGmm2JoinDecisionMask);
328+ const bool joined =
329+ (decision & kMegaMoeFixedGmm2JoinDecisionBit) != 0 && encodedExpert != 0U && groupIdx + 1U >= encodedExpert;
330+ return joined ? fixed.physicalAicNum : fixed.gmm2GroupSize;
331+}
332+ 
333+template <typename OutputElement>
334+AICORE inline void Combine<OutputElement>::WaitGmm2Ready(uint32_t groupIdx) const
335+{
336+ const MegaMoeSyncLayout sync = FixedSyncLayout(tilingData_);
337+ const uint32_t readyLocalId = coreIdx_ % tilingData_->fixedGroupTiling.gmm2GroupSize;
338+ const uint32_t readySlot = sync.combineReadyBase + readyLocalId;
339+ const int32_t expectedEpoch = static_cast<int32_t>(groupIdx * 2U + 2U);
340+ if (coreIdx_ == ReadyCoordinatorCore()) {
341+ const uint32_t producerCount = Gmm2ProducerCount(groupIdx);
342+ const uint32_t producerBaseOffset = producerCount == tilingData_->fixedGroupTiling.physicalAicNum ?
343+ 0U :
344+ tilingData_->fixedGroupTiling.gmm1GroupSize;
345+ CoordinateGroupConsumersMte(
346+ workspaceGM_, tilingData_, sync.gmm2ArrivalBase + producerBaseOffset, sync.combineReadyBase, producerCount,
347+ tilingData_->fixedGroupTiling.gmm2GroupSize, groupIdx);
348+ } else {
349+ WaitEpochAcquire(FixedSyncSlot(workspaceGM_, tilingData_, readySlot), expectedEpoch);
315 }350 }
316}351}
317 352 
@@ -375,14 +410,23 @@ template <typename OutputElement>
375AICORE inline void Combine<OutputElement>::ProcessDirectLargeTokenPath()410AICORE inline void Combine<OutputElement>::ProcessDirectLargeTokenPath()
376{411{
377 uint32_t groupBase = 0;412 uint32_t groupBase = 0;
413+ const bool delayedStart = expertPerRank_ != 0U;
414+ const uint32_t initialReadyExpert = expertPerRank_ > tilingData_->fixedGroupTiling.combineStartAfterGmm2Expert ?
415+ tilingData_->fixedGroupTiling.combineStartAfterGmm2Expert :
416+ expertPerRank_ - 1U;
417+ if (delayedStart && coreIdx_ < DirectLargeTaskCount()) {
418+ WaitGmm2Ready(initialReadyExpert);
419+ }
378 for (uint32_t groupIdx = 0; groupIdx < expertPerRank_; ++groupIdx) { // 逐个group遍历420 for (uint32_t groupIdx = 0; groupIdx < expertPerRank_; ++groupIdx) { // 逐个group遍历
379 const uint32_t currentM = CurrentM(groupIdx); // expert 总共有多少输出 row421 const uint32_t currentM = CurrentM(groupIdx); // expert 总共有多少输出 row
380- WaitGmm2Ready(groupIdx, true);
381 const uint32_t lanesPerRank = LargeLanesPerRank();422 const uint32_t lanesPerRank = LargeLanesPerRank();
423+ const uint32_t taskCount = rankSize_ * lanesPerRank;
424+ if (!delayedStart || groupIdx > initialReadyExpert) {
425+ WaitGmm2Ready(groupIdx);
426+ }
382 if (lanesPerRank == 0U) {427 if (lanesPerRank == 0U) {
383 continue;428 continue;
384 }429 }
385- const uint32_t taskCount = rankSize_ * lanesPerRank; // 卡按照lane再切分，当前默认每卡 2 lane
386 for (uint32_t taskIdx = coreIdx_; taskIdx < taskCount; taskIdx += coreNum_) {430 for (uint32_t taskIdx = coreIdx_; taskIdx < taskCount; taskIdx += coreNum_) {
387 const uint32_t safeLanes = lanesPerRank == 0U ? 1U : lanesPerRank;431 const uint32_t safeLanes = lanesPerRank == 0U ? 1U : lanesPerRank;
388 const uint32_t srcRank = taskIdx / safeLanes;432 const uint32_t srcRank = taskIdx / safeLanes;
@@ -395,195 +439,33 @@ AICORE inline void Combine<OutputElement>::ProcessDirectLargeTokenPath()
395 const uint32_t dstRowOffset = DstRowOffset(srcRank, groupIdx);439 const uint32_t dstRowOffset = DstRowOffset(srcRank, groupIdx);
396 ProcessDirectLargeSegmentRows(srcRank, srcRowOffset + rowBegin, rowNum, dstRowOffset + rowBegin);440 ProcessDirectLargeSegmentRows(srcRank, srcRowOffset + rowBegin, rowNum, dstRowOffset + rowBegin);
397 }441 }
442+ const uint32_t readyExpertCount = groupIdx + 1U;
443+ const uint32_t phase1ReadyExpertCount = tilingData_->fixedGroupTiling.unpermutePhase1ReadyExpertCount;
444+ if (RankStreamingEnabled() && lanesPerRank == 1U && coreIdx_ < DirectLargeWorkerCount() &&
445+ readyExpertCount == phase1ReadyExpertCount && readyExpertCount < expertPerRank_) {
446+ FinalizeExpertStores();
447+ PublishAssignedExpertProgress(readyExpertCount);
448+ }
398 groupBase += currentM;449 groupBase += currentM;
399 }450 }
400}451}
401 452 
402template <typename OutputElement>453template <typename OutputElement>
403-AICORE inline void Combine<OutputElement>::LoadSmallSubtile(454+AICORE inline void Combine<OutputElement>::ProcessFixed(uint32_t groupLocalId, uint32_t groupSize)
404- uint32_t bufferId, uint32_t srcRowOffset, uint32_t rowNum, uint32_t colBegin, uint32_t colNum) const
405-{
406- wait_flag(PIPE_V, PIPE_MTE2, LoadFreeEvent(bufferId));
407- SmallTileC cTile(rowNum, colNum);
408- pto::TASSIGN(cTile, ubCOffset_[bufferId]);
409- BlockShape cShape(rowNum, colNum);
410- BlockStride cStride(
411- static_cast<int64_t>(rowNum) * problemK_, static_cast<int64_t>(rowNum) * problemK_,
412- static_cast<int64_t>(rowNum) * problemK_, problemK_);
413- CBlockGlobal cGlobal(gmm2OutputPtr_ + static_cast<uint64_t>(srcRowOffset) * problemK_ + colBegin, cShape, cStride);
414- pto::TLOAD(cTile, cGlobal);
415- 
416- if (smallScaleSourceOffset_[bufferId] != srcRowOffset) {
417- SmallTileFp32 scaleTile(1, rowNum);
418- pto::TASSIGN(scaleTile, ubScaleOffset_[bufferId]);
419- VectorShape scaleShape(rowNum);
420- VectorStride scaleStride(rowNum, rowNum, rowNum, rowNum);
421- ScaleGlobal scaleGlobal(perTokenScale2Ptr_ + srcRowOffset, scaleShape, scaleStride);
422- pto::TLOAD(scaleTile, scaleGlobal);
423- smallScaleSourceOffset_[bufferId] = srcRowOffset;
424- }
425- set_flag(PIPE_MTE2, PIPE_V, LoadReadyEvent(bufferId));
426- set_flag(PIPE_MTE2, PIPE_S, LoadReadyEvent(bufferId));
427-}
428- 
429-template <typename OutputElement>
430-AICORE inline void Combine<OutputElement>::DequantDirectSmallSubtile(
431- uint32_t bufferId, uint32_t rowNum, uint32_t colNum)
432-{
433- wait_flag(PIPE_MTE2, PIPE_V, LoadReadyEvent(bufferId));
434- SmallTileFp32 fp32Tile(rowNum, colNum);
435- SmallTileC cTile(rowNum, colNum);
436- pto::TASSIGN(fp32Tile, ubFp32Offset_[bufferId]);
437- pto::TASSIGN(cTile, ubCOffset_[bufferId]);
438- pto::TCVT(fp32Tile, cTile, pto::RoundMode::CAST_NONE);
439- pipe_barrier(PIPE_V);
440- set_flag(PIPE_V, PIPE_MTE2, LoadFreeEvent(bufferId));
441- 
442- wait_flag(PIPE_MTE2, PIPE_S, LoadReadyEvent(bufferId));
443- SmallTileFp32 scaleTile(1, rowNum);
444- pto::TASSIGN(scaleTile, ubScaleOffset_[bufferId]);
445- for (uint32_t row = 0; row < rowNum; ++row) {
446- const float scale = scaleTile.GetValue(row);
447- SmallTileFp32 rowTile(1, colNum);
448- pto::TASSIGN(
449- rowTile,
450- ubFp32Offset_[bufferId] + static_cast<uint64_t>(row) * kCombineSmallTokenSubtileCols * sizeof(float));
451- pto::TMULS(rowTile, rowTile, scale);
452- }
453- pipe_barrier(PIPE_V);
454- wait_flag(PIPE_MTE3, PIPE_V, StoreFreeEvent(bufferId));
455- SmallTileD dTile(rowNum, colNum);
456- pto::TASSIGN(dTile, ubDOffset_[bufferId]);
457- pto::TCVT(dTile, fp32Tile, pto::RoundMode::CAST_RINT);
458- set_flag(PIPE_V, PIPE_MTE3, StoreReadyEvent(bufferId));
459-}
460- 
461-template <typename OutputElement>
462-AICORE inline void Combine<OutputElement>::StoreSmallSubtileIntersection(
463- uint32_t bufferId, __gm__ OutputElement* dstBase, uint32_t dstRowOffset, uint32_t ubRowOffset, uint32_t rowNum,
464- uint32_t colBegin, uint32_t colNum)
465-{
466- SmallTileD dTile(rowNum, colNum);
467- pto::TASSIGN(
468- dTile, ubDOffset_[bufferId] +
469- static_cast<uint64_t>(ubRowOffset) * kCombineSmallTokenSubtileCols * sizeof(OutputElement));
470- BlockShape dShape(rowNum, colNum);
471- BlockStride dStride(
472- static_cast<int64_t>(rowNum) * problemK_, static_cast<int64_t>(rowNum) * problemK_,
473- static_cast<int64_t>(rowNum) * problemK_, problemK_);
474- DBlockGlobal dGlobal(dstBase + static_cast<uint64_t>(dstRowOffset) * problemK_ + colBegin, dShape, dStride);
475- pto::TSTORE(dGlobal, dTile);
476-}
477- 
478-template <typename OutputElement>
479-AICORE inline void Combine<OutputElement>::StoreSmallSubtileToRanks(
480- uint32_t groupIdx, const GmmCommonTileInfo& tileInfo, uint32_t tileRowBegin, uint32_t rows, uint32_t bufferId)
481-{
482- const uint32_t stTile = tileRowBegin;
483- const uint32_t edTile = tileRowBegin + rows;
484- uint32_t preSumRankInExpert = 0U;
485- uint32_t tileOffset = 0U;
486- wait_flag(PIPE_V, PIPE_MTE3, StoreReadyEvent(bufferId));
487- for (uint32_t srcRank = 0; srcRank < rankSize_; ++srcRank) {
488- const uint32_t lenRankInExpert = RowsRaw(srcRank, groupIdx);
489- const uint32_t dstExpertOffset = DstRowOffset(srcRank, groupIdx);
490- const uint32_t stRankInExpert = preSumRankInExpert;
491- const uint32_t edRankInExpert = stRankInExpert + lenRankInExpert;
492- preSumRankInExpert += lenRankInExpert;
493- if (stRankInExpert >= edTile) {
494- break;
495- }
496- if (edRankInExpert <= stTile) {
497- continue;
498- }
499- const uint32_t stData = stRankInExpert > stTile ? stRankInExpert : stTile;
500- const uint32_t edData = edRankInExpert < edTile ? edRankInExpert : edTile;
501- if (edData <= stData) {
502- continue;
503- }
504- const uint32_t lenData = edData - stData;
505- const uint32_t dstOffsetInExpert = stTile > stRankInExpert ? stTile - stRankInExpert : 0U;
506- __gm__ OutputElement* dstBase = reinterpret_cast<__gm__ OutputElement*>(
507- remoteWindow_.RemoteBase(peerMemoryLayout_.offsetD, static_cast<int32_t>(srcRank)));
508- if (dstBase != nullptr) {
509- StoreSmallSubtileIntersection(
510- bufferId, dstBase, dstExpertOffset + dstOffsetInExpert, tileOffset, lenData, tileInfo.blockColStart,
511- tileInfo.actualN);
512- }
513- tileOffset += lenData;
514- }
515- set_flag(PIPE_MTE3, PIPE_V, StoreFreeEvent(bufferId));
516-}
517- 
518-template <typename OutputElement>
519-AICORE inline void Combine<OutputElement>::ProcessDirectSmallTile(
520- uint32_t groupIdx, uint32_t groupBase, const GmmCommonTileInfo& tileInfo, uint32_t subtileBegin,
521- uint32_t subtileCount)
522-{
523- for (uint32_t subtile = 0; subtile < subtileCount; ++subtile) {
524- const uint32_t rowInTile = (subtileBegin + subtile) * kCombineSmallTokenSubtileRows;
525- if (rowInTile >= tileInfo.actualM) {
526- continue;
527- }
528- const uint32_t rows = (tileInfo.actualM - rowInTile > kCombineSmallTokenSubtileRows) ?
529- kCombineSmallTokenSubtileRows :
530- (tileInfo.actualM - rowInTile);
531- const uint32_t tileRowBegin = tileInfo.blockRowStart + rowInTile;
532- const uint32_t srcRow = groupBase + tileRowBegin;
533- const uint32_t bufferId = pingpongId_;
534- pingpongId_ = (pingpongId_ + 1U) % kCombineBufferNum;
535- LoadSmallSubtile(bufferId, srcRow, rows, tileInfo.blockColStart, tileInfo.actualN);
536- DequantDirectSmallSubtile(bufferId, rows, tileInfo.actualN);
537- StoreSmallSubtileToRanks(groupIdx, tileInfo, tileRowBegin, rows, bufferId);
538- }
539-}
540- 
541-template <typename OutputElement>
542-AICORE inline void Combine<OutputElement>::ProcessDirectSmallTokenPath()
543-{
544- uint32_t groupBase = 0;
545- uint32_t startCoreIdx = 0;
546- const uint32_t aicCoreIdx = get_block_idx();
547- const uint32_t aicCoreNum = get_block_num();
548- const uint32_t aivSubCoreIdx = get_subblockid();
549- const uint32_t l1TileM = tilingData_->gmm2Tiling.l1TileM;
550- const uint32_t l1TileN = tilingData_->gmm2Tiling.l1TileN;
551- for (uint32_t groupIdx = 0; groupIdx < expertPerRank_; ++groupIdx) {
552- const uint32_t currentM = MoeClipCurrentM(CurrentM(groupIdx), groupBase, maxOutputSize_);
553- WaitGmm2Ready(groupIdx, false);
554- const uint32_t coreLoops = GmmCommonCoreLoops(currentM, problemK_, l1TileM, l1TileN); // 按照L1的size做切tile
555- const uint32_t startLoopIdx =
556- aicCoreNum == 0U ? 0U : GmmCommonStartLoopIdx(aicCoreIdx, aicCoreNum, startCoreIdx);
557- for (uint32_t loopIdx = startLoopIdx; aicCoreNum != 0U && loopIdx < coreLoops; loopIdx += aicCoreNum) {
558- const GmmCommonTileInfo tileInfo = GmmCommonBuildTileInfo(currentM, problemK_, l1TileM, l1TileN, loopIdx);
559- const uint32_t subtileCount = static_cast<uint32_t>(
560- ceilDiv(tileInfo.actualM, kCombineSmallTokenSubtileRows)); // 按照m维度做切分成subtile
561- const uint32_t firstHalfSubtiles = subtileCount / 2U;
562- const uint32_t firstSubtile = aivSubCoreIdx == 0U ? 0U : firstHalfSubtiles;
563- uint32_t assignedSubtiles = subtileCount / 2U;
564- if (aivSubCoreIdx == 1U && assignedSubtiles * 2U < subtileCount) {
565- ++assignedSubtiles;
566- }
567- ProcessDirectSmallTile(groupIdx, groupBase, tileInfo, firstSubtile, assignedSubtiles);
568- }
569- startCoreIdx = aicCoreNum == 0U ? 0U : (startCoreIdx + coreLoops) % aicCoreNum;
570- groupBase += currentM;
571- }
572-}
573- 
574-template <typename OutputElement>
575-AICORE inline void Combine<OutputElement>::Process()
576{455{
577 if ASCEND_IS_AIC {456 if ASCEND_IS_AIC {
578 return;457 return;
579 }458 }
459+ coreIdx_ = groupLocalId;
460+ coreNum_ = groupSize;
580 SetInitialFlags();461 SetInitialFlags();
581- if (DirectSmallEnabled()) { // problemM_ * topK_ 小于4096的时候走small case462+ if (coreIdx_ < DirectLargeTaskCount()) {
582- ProcessDirectSmallTokenPath();
583- } else if (DirectLargeEnabled()) {
584 ProcessDirectLargeTokenPath();463 ProcessDirectLargeTokenPath();
585 }464 }
586- ProcessFinalBoundary();465+ if (RankStreamingEnabled()) {
466+ FinalizeLocalPipe();
467+ FinalizeRankStreamingLane();
468+ }
587}469}
588 470 
589#endif // DISPATCH_MEGA_COMBINE_COMBINE_H471#endif // DISPATCH_MEGA_COMBINE_COMBINE_H
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/dispatch.h+29-33
@@ -19,7 +19,7 @@ See LICENSE in the root of the software repository for the full text of the Lice
19#include "utils/common_helpers.hpp"19#include "utils/common_helpers.hpp"
20#include "utils/const_args.hpp"20#include "utils/const_args.hpp"
21#include "utils/hccl_window.hpp"21#include "utils/hccl_window.hpp"
22-#include "utils/pto_sync_substrate.hpp"22+#include "utils/mega_expert_sync.hpp"
23#include "utils/pto_vector.hpp"23#include "utils/pto_vector.hpp"
24 24 
25constexpr uint32_t kDispatchBufferNum = 2U;25constexpr uint32_t kDispatchBufferNum = 2U;
@@ -42,6 +42,7 @@ public:
42 AICORE inline void Init(GM_ADDR expertTokenNumsGM, GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData)42 AICORE inline void Init(GM_ADDR expertTokenNumsGM, GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData)
43 {43 {
44 (void)expertTokenNumsGM;44 (void)expertTokenNumsGM;
45+ workspaceGM_ = workspaceGM;
45 tilingData_ = tilingData;46 tilingData_ = tilingData;
46 47 
47 const auto& info = tilingData_->megaMoeInfo;48 const auto& info = tilingData_->megaMoeInfo;
@@ -74,11 +75,12 @@ public:
74 offsetAPtr_ = reinterpret_cast<__gm__ int8_t*>(remoteWindow_.LocalBase() + peerMemoryLayout_.offsetA);75 offsetAPtr_ = reinterpret_cast<__gm__ int8_t*>(remoteWindow_.LocalBase() + peerMemoryLayout_.offsetA);
75 }76 }
76 77 
77- AICORE inline void Process() const78+ AICORE inline void ProcessFixed(uint32_t groupLocalId, uint32_t groupSize)
78 {79 {
79- if ASCEND_IS_AIV {80+ coreIdx_ = groupLocalId;
80- ProcessRankSplitCopy();81+ coreNum_ = groupSize;
81- }82+ dispatchGatherScratchPtr_ = reinterpret_cast<__gm__ uint8_t*>(workspaceGM_ + DispatchGatherScratchCoreOffset());
83+ ProcessRankSplitCopy();
82 }84 }
83 85 
84private:86private:
@@ -97,11 +99,6 @@ private:
97 problemK_ / sizeof(uint32_t) <= kDispatchGatherPackedWordTileCols;99 problemK_ / sizeof(uint32_t) <= kDispatchGatherPackedWordTileCols;
98 }100 }
99 101 
100- AICORE inline void SetGmm1ReadyByLogicalEvent(uint32_t logicalGroupEventIdx) const
101- {
102- CrossCoreSetFlag<0x2, PIPE_MTE3>(MegaMoeD2CHardFlagId(logicalGroupEventIdx));
103- }
104- 
105 AICORE inline uint64_t DispatchGatherScratchCoreOffset() const102 AICORE inline uint64_t DispatchGatherScratchCoreOffset() const
106 {103 {
107 return tilingData_->dispatchTiling.dispatchGatherScratchOffset +104 return tilingData_->dispatchTiling.dispatchGatherScratchOffset +
@@ -132,10 +129,6 @@ private:
132 static_cast<uint64_t>(bufferId) * tilingData_->dispatchTiling.dispatchGatherTileBytes);129 static_cast<uint64_t>(bufferId) * tilingData_->dispatchTiling.dispatchGatherTileBytes);
133 }130 }
134 131 
135- AICORE inline void SetGmm1InitialReady() const { SetGmm1ReadyByLogicalEvent(0U); }
136- 
137- AICORE inline void SetGmm1GroupReady(uint32_t groupIdx) const { SetGmm1ReadyByLogicalEvent(groupIdx + 1U); }
138- 
139 AICORE inline void PrepareDispatchGatherCopyEvents() const132 AICORE inline void PrepareDispatchGatherCopyEvents() const
140 {133 {
141 set_flag(PIPE_MTE3, PIPE_MTE2, DispatchGatherBufferEvent(0U));134 set_flag(PIPE_MTE3, PIPE_MTE2, DispatchGatherBufferEvent(0U));
@@ -170,22 +163,22 @@ private:
170 }163 }
171 164 
172 AICORE inline void FetchRankGroupRows(165 AICORE inline void FetchRankGroupRows(
173- uint32_t srcRank, uint32_t groupIdx, uint32_t prevGroupSum, uint32_t& prevSum, int32_t& pingpongIdx) const166+ uint32_t srcRank, uint32_t groupIdx, uint32_t prevGroupSum, int32_t& pingpongIdx) const
174 {167 {
175 const uint32_t rawRows = RawRowsForLocalGroup(srcRank, groupIdx);168 const uint32_t rawRows = RawRowsForLocalGroup(srcRank, groupIdx);
176 const uint32_t dstRowBase = prevGroupSum + CopyCumsumBeforeSource(srcRank, groupIdx);169 const uint32_t dstRowBase = prevGroupSum + CopyCumsumBeforeSource(srcRank, groupIdx);
177- const uint32_t srcRowBase = prevSum;
178 uint32_t rows = 0U;170 uint32_t rows = 0U;
179 if (dstRowBase < maxOutputSize_) {171 if (dstRowBase < maxOutputSize_) {
180 rows = rawRows;172 rows = rawRows;
181 if (dstRowBase + rows > maxOutputSize_) {173 if (dstRowBase + rows > maxOutputSize_) {
182 rows = maxOutputSize_ - dstRowBase;174 rows = maxOutputSize_ - dstRowBase;
183 }175 }
184- prevSum += rows;
185 }176 }
186 if (rows == 0U) {177 if (rows == 0U) {
187 return;178 return;
188 }179 }
180+ const uint32_t srcRowBase =
181+ static_cast<uint32_t>(preSumBeforeRankPtr_[static_cast<uint64_t>(srcRank) * expertPerRank_ + groupIdx]);
189 __gm__ int8_t* remotePackedRows = reinterpret_cast<__gm__ int8_t*>(182 __gm__ int8_t* remotePackedRows = reinterpret_cast<__gm__ int8_t*>(
190 remoteWindow_.RemoteBase(peerMemoryLayout_.offsetA, static_cast<int32_t>(srcRank)));183 remoteWindow_.RemoteBase(peerMemoryLayout_.offsetA, static_cast<int32_t>(srcRank)));
191 __gm__ int8_t* remoteSrc = remotePackedRows + static_cast<uint64_t>(srcRowBase) * PackedRowStride();184 __gm__ int8_t* remoteSrc = remotePackedRows + static_cast<uint64_t>(srcRowBase) * PackedRowStride();
@@ -309,32 +302,35 @@ private:
309 302 
310 AICORE inline void ProcessRankSplitCopy() const303 AICORE inline void ProcessRankSplitCopy() const
311 {304 {
312- SetGmm1InitialReady();305+ if ASCEND_IS_AIV {
313- if (coreIdx_ < rankSize_) {306+ const bool activeCopyCore = rankSize_ != 0U && coreIdx_ < rankSize_;
314- uint32_t prevSum =
315- static_cast<uint32_t>(preSumBeforeRankPtr_[static_cast<uint64_t>(coreIdx_) * expertPerRank_]);
316 uint32_t prevGroupSum = 0U;307 uint32_t prevGroupSum = 0U;
317 int32_t pingpongIdx = 0;308 int32_t pingpongIdx = 0;
318 for (uint32_t groupIdx = 0U; groupIdx < expertPerRank_; ++groupIdx) {309 for (uint32_t groupIdx = 0U; groupIdx < expertPerRank_; ++groupIdx) {
319- PrepareDispatchGatherCopyEvents();
320 const uint32_t currentM = static_cast<uint32_t>(310 const uint32_t currentM = static_cast<uint32_t>(
321 cumsumMMPtr_[static_cast<uint64_t>(rankSize_ - 1U) * expertPerRank_ + groupIdx]);311 cumsumMMPtr_[static_cast<uint64_t>(rankSize_ - 1U) * expertPerRank_ + groupIdx]);
322- for (uint32_t srcRank = coreIdx_; srcRank < rankSize_; srcRank += coreNum_) {312+ if (activeCopyCore) {
323- FetchRankGroupRows(srcRank, groupIdx, prevGroupSum, prevSum, pingpongIdx);313+ PrepareDispatchGatherCopyEvents();
314+ for (uint32_t srcRank = coreIdx_; srcRank < rankSize_; srcRank += coreNum_) {
315+ FetchRankGroupRows(srcRank, groupIdx, prevGroupSum, pingpongIdx);
316+ }
317+ prevGroupSum += currentM;
318+ WaitDispatchGatherCopyEvents();
324 }319 }
325- prevGroupSum += currentM;320+ 
326- WaitDispatchGatherCopyEvents();321+ const __gm__ MegaMoeFixedGroupTiling& fixed = tilingData_->fixedGroupTiling;
327- pto::SYNCALL<pto::SyncCoreType::AIVOnly>();322+ const MegaMoeSyncLayout sync = FixedSyncLayout(tilingData_);
328- SetGmm1GroupReady(groupIdx);323+ const uint32_t gmm1ConsumerCount =
329- }324+ groupIdx < fixed.fullAicGmm1ExpertCount ? fixed.physicalAicNum : fixed.gmm1GroupSize;
330- } else {325+ const uint32_t coordinatorLocalId = rankSize_ == 0U ? 0U : rank_ % rankSize_;
331- for (uint32_t groupIdx = 0U; groupIdx < expertPerRank_; ++groupIdx) {326+ NotifyGroupConsumersMte(
332- pto::SYNCALL<pto::SyncCoreType::AIVOnly>();327+ workspaceGM_, tilingData_, sync.dispatchArrivalBase, sync.dispatchReadyBase, rankSize_,
333- SetGmm1GroupReady(groupIdx);328+ gmm1ConsumerCount, coreIdx_, coordinatorLocalId, groupIdx);
334 }329 }
335 }330 }
336 }331 }
337 332 
333+ GM_ADDR workspaceGM_ = nullptr;
338 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;334 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;
339 335 
340 __gm__ int8_t* gmAPtr_ = nullptr;336 __gm__ int8_t* gmAPtr_ = nullptr;
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/dispatch_mega_combine.h+133-22
@@ -24,6 +24,7 @@ See LICENSE in the root of the software repository for the full text of the Lice
24#include "kernel_launch.hpp"24#include "kernel_launch.hpp"
25#include "swiglu.h"25#include "swiglu.h"
26#include "unpermute.h"26#include "unpermute.h"
27+#include "utils/mega_expert_sync.hpp"
27 28 
28template <typename InputElement, uint32_t ExpertPerRank>29template <typename InputElement, uint32_t ExpertPerRank>
29AICORE inline void FrontRunVmsSort(FrontReorderVmsSort<InputElement>& path)30AICORE inline void FrontRunVmsSort(FrontReorderVmsSort<InputElement>& path)
@@ -77,6 +78,9 @@ public:
77 __aicore__ inline void Process();78 __aicore__ inline void Process();
78 79 
79private:80private:
81+ __aicore__ inline void ProcessFixedGroups(uint16_t stageNum);
82+ __aicore__ inline void ProcessFixedGmm1(uint32_t physicalBlockId);
83+ 
80 GM_ADDR xGM_ = nullptr;84 GM_ADDR xGM_ = nullptr;
81 GM_ADDR weight1GM_ = nullptr;85 GM_ADDR weight1GM_ = nullptr;
82 GM_ADDR weight2GM_ = nullptr;86 GM_ADDR weight2GM_ = nullptr;
@@ -110,35 +114,142 @@ __aicore__ inline void MegaMoe<CType_, ExpertPerRank>::Init(
110}114}
111 115 
112template <typename CType_, uint32_t ExpertPerRank>116template <typename CType_, uint32_t ExpertPerRank>
113-__aicore__ inline void MegaMoe<CType_, ExpertPerRank>::Process()117+__aicore__ inline void MegaMoe<CType_, ExpertPerRank>::ProcessFixedGmm1(uint32_t physicalBlockId)
114{118{
115- using OutputElement = half;
116- 
117- FrontReorderProcess<CType_, ExpertPerRank>(xGM_, expertIdGM_, expertTokenNumsGM_, workspaceGM_, tilingData_);
118- 
119- DispatchGather<CType_> dispatchGather;
120- dispatchGather.Init(expertTokenNumsGM_, workspaceGM_, tilingData_);
121- dispatchGather.Process();
122- 
123 Gmm1<CType_> gmm1;119 Gmm1<CType_> gmm1;
124 gmm1.Init(weight1GM_, scale1GM_, expertTokenNumsGM_, workspaceGM_, tilingData_);120 gmm1.Init(weight1GM_, scale1GM_, expertTokenNumsGM_, workspaceGM_, tilingData_);
125- gmm1.Process();121+ gmm1.ProcessFixed(physicalBlockId, tilingData_->fixedGroupTiling.physicalAicNum);
122+}
126 123 
127- Swiglu<CType_> swiglu;124+template <typename CType_, uint32_t ExpertPerRank>
128- swiglu.Init(expertTokenNumsGM_, workspaceGM_, tilingData_);125+__aicore__ inline void MegaMoe<CType_, ExpertPerRank>::ProcessFixedGroups(uint16_t stageNum)
129- swiglu.Process();126+{
127+ const MegaMoeFixedCoreRoleInfo role = FixedCoreRole(tilingData_);
128+ const __gm__ MegaMoeFixedGroupTiling& fixed = tilingData_->fixedGroupTiling;
129+ const bool rankStreaming = tilingData_->unpermuteTiling.unpermuteImplMode == kMegaMoeUnpermuteImplRankStreaming;
130+ const uint32_t rankStreamingWorkerCount = fixed.physicalAivNum;
130 131 
131- Gmm2<CType_> gmm2;132+ if (role.role == kMegaMoeFixedRoleDispatch && role.groupLocalId < tilingData_->runtimeInfo.rankSize &&
132- gmm2.Init(weight2GM_, scale2GM_, expertTokenNumsGM_, workspaceGM_, tilingData_);133+ stageNum >= 9U) {
133- gmm2.Process();134+ DispatchGather<CType_> dispatchGather;
135+ dispatchGather.Init(expertTokenNumsGM_, workspaceGM_, tilingData_);
136+ dispatchGather.ProcessFixed(role.groupLocalId, role.groupSize);
137+ } else if ((role.role == kMegaMoeFixedRoleGmm1 || role.role == kMegaMoeFixedRoleGmm2) && stageNum >= 10U) {
138+ ProcessFixedGmm1(role.physicalBlockId);
139+ const bool dynamicGmm2Join = stageNum >= 12U;
140+ if (dynamicGmm2Join && role.role == kMegaMoeFixedRoleGmm1 && role.groupLocalId == 0U) {
141+ pipe_barrier(PIPE_ALL);
142+ dsb(DSB_DDR);
143+ PublishScalarEpoch(
144+ FixedSyncSlot(workspaceGM_, tilingData_, FixedSyncLayout(tilingData_).gmm1DoneSlot),
145+ kMegaMoeFixedGmm1DoneMarker);
146+ }
147+ if (stageNum >= 12U && (role.role == kMegaMoeFixedRoleGmm2 || dynamicGmm2Join)) {
148+ Gmm2<CType_> gmm2;
149+ gmm2.Init(weight2GM_, scale2GM_, expertTokenNumsGM_, workspaceGM_, tilingData_);
150+ if (role.role == kMegaMoeFixedRoleGmm2) {
151+ gmm2.ProcessFixed(role.groupLocalId, role.groupSize);
152+ } else {
153+ gmm2.ProcessFixedHelper(role.groupLocalId);
154+ }
155+ }
156+ } else if (
157+ role.role == kMegaMoeFixedRoleSwiglu && role.groupLocalId < fixed.swigluActiveGroupSize && stageNum >= 11U) {
158+ Swiglu<CType_> swiglu;
159+ swiglu.Init(expertTokenNumsGM_, workspaceGM_, tilingData_);
160+ swiglu.ProcessFixed(role.groupLocalId, fixed.swigluActiveGroupSize);
161+ } else if (role.role == kMegaMoeFixedRoleCombine && role.groupLocalId < role.groupSize && stageNum >= 13U) {
162+ Combine<half> combine;
163+ combine.Init(workspaceGM_, tilingData_);
164+ combine.ProcessFixed(role.groupLocalId, role.groupSize);
165+ }
134 166 
135- Combine<OutputElement> combine;167+ if ASCEND_IS_AIV {
136- combine.Init(workspaceGM_, tilingData_);168+ if (!rankStreaming && stageNum >= 13U) {
137- combine.Process();169+ Combine<half> finalBoundary;
170+ finalBoundary.Init(workspaceGM_, tilingData_);
171+ finalBoundary.ProcessFixedFinalBoundary(
172+ role.role, role.flatAivId, role.role != kMegaMoeFixedRoleCombine || role.groupLocalId < role.groupSize);
173+ }
174+ if (rankStreaming && stageNum == 13U && role.physicalBlockId == 0U && role.subblockId == 0U) {
175+ PtoRemoteWindow remoteWindow;
176+ remoteWindow.Init(reinterpret_cast<GM_ADDR>(tilingData_->runtimeInfo.remoteWindowContext));
177+ const int32_t epoch = remoteWindow.DataReadyEpoch();
178+ for (uint32_t producerRank = 0U; producerRank < tilingData_->runtimeInfo.rankSize; ++producerRank) {
179+ WaitEpochAcquire(remoteWindow.LocalDataReadySlot(producerRank), epoch);
180+ }
181+ }
138 182 
139- Unpermute<OutputElement> unpermute;183+ const uint32_t initialUnpermuteGroupSize = fixed.gmm1GroupSize * 2U;
140- unpermute.Init(workspaceGM_, probsGM_, outGM_, tilingData_);184+ const bool initialUnpermuteGroup = role.physicalBlockId < fixed.gmm1GroupSize;
141- unpermute.Process();185+ const uint32_t rankStreamingWorkerIdx = initialUnpermuteGroup ?
186+ role.physicalBlockId + role.subblockId * fixed.gmm1GroupSize :
187+ initialUnpermuteGroupSize + role.groupLocalId;
188+ const bool unpermuteWorker = !rankStreaming || rankStreamingWorkerIdx < rankStreamingWorkerCount;
189+ if (rankStreaming && stageNum >= 14U && initialUnpermuteGroup) {
190+ PtoRemoteWindow remoteWindow;
191+ remoteWindow.Init(reinterpret_cast<GM_ADDR>(tilingData_->runtimeInfo.remoteWindowContext));
192+ const int32_t epoch = remoteWindow.DataReadyEpoch();
193+ if (role.role == kMegaMoeFixedRoleDispatch) {
194+ remoteWindow.PublishDispatchRelease(role.groupLocalId, epoch);
195+ } else if (role.role == kMegaMoeFixedRoleSwiglu) {
196+ remoteWindow.PublishSwigluRelease(role.groupLocalId, epoch);
197+ }
198+ 
199+ if (rankStreamingWorkerIdx == 0U) {
200+ remoteWindow.WaitDispatchReleaseMte(fixed.dispatchGroupSize, epoch);
201+ remoteWindow.WaitSwigluReleaseMte(fixed.swigluGroupSize, epoch);
202+ 
203+ uint32_t readyExpertCounts[COMBINE_EXPERT_PROGRESS_MAX_RANKS] = {0U};
204+ const uint32_t rankCount = tilingData_->runtimeInfo.rankSize;
205+ const uint32_t expertPerRank = tilingData_->megaMoeInfo.expertPerRank;
206+ const uint32_t configuredCut = fixed.unpermutePhase1ReadyExpertCount;
207+ const uint32_t readyCut = configuredCut < expertPerRank ? configuredCut : expertPerRank;
208+ uint32_t minimumReady = 0U;
209+ while (minimumReady < readyCut) {
210+ minimumReady = remoteWindow.ReadExpertProgressMte(epoch, expertPerRank, readyExpertCounts);
211+ if (minimumReady < readyCut) {
212+ RemoteWindowSyncPollBackoff();
213+ }
214+ }
215+ remoteWindow.AcquireDataReady();
216+ uint32_t readyRankMask = 0U;
217+ for (uint32_t producerRank = 0U; producerRank < rankCount; ++producerRank) {
218+ if (readyExpertCounts[producerRank] >= readyCut) {
219+ readyRankMask |= 1U << producerRank;
220+ }
221+ }
222+ remoteWindow.PublishUnpermutePhase1Progress(readyExpertCounts, rankCount, readyRankMask, epoch);
223+ const uint32_t initialWorkerCount = rankStreamingWorkerCount < initialUnpermuteGroupSize ?
224+ rankStreamingWorkerCount :
225+ initialUnpermuteGroupSize;
226+ remoteWindow.PublishUnpermuteStartRangeMte(0U, initialWorkerCount, epoch);
227+ }
228+ }
229+ if (stageNum >= 14U && unpermuteWorker) {
230+ uint32_t workerIdx = role.flatAivId;
231+ uint32_t workerCount = tilingData_->fixedGroupTiling.physicalAivNum;
232+ if (rankStreaming) {
233+ workerIdx = rankStreamingWorkerIdx;
234+ workerCount = rankStreamingWorkerCount;
235+ PtoRemoteWindow remoteWindow;
236+ remoteWindow.Init(reinterpret_cast<GM_ADDR>(tilingData_->runtimeInfo.remoteWindowContext));
237+ WaitEpochAcquire(remoteWindow.LocalUnpermuteStartSlot(workerIdx), remoteWindow.DataReadyEpoch());
238+ }
239+ Unpermute<half> unpermute;
240+ unpermute.Init(workspaceGM_, expertIdGM_, probsGM_, outGM_, tilingData_, workerIdx, workerCount);
241+ unpermute.Process();
242+ }
243+ }
244+}
245+ 
246+template <typename CType_, uint32_t ExpertPerRank>
247+__aicore__ inline void MegaMoe<CType_, ExpertPerRank>::Process()
248+{
249+ const uint16_t stageNum = tilingData_->frontReorderTiling.stageNum;
250+ 
251+ FrontReorderProcess<CType_, ExpertPerRank>(xGM_, expertIdGM_, expertTokenNumsGM_, workspaceGM_, tilingData_);
252+ ProcessFixedGroups(stageNum);
142}253}
143 254 
144#endif // DISPATCH_MEGA_COMBINE_H255#endif // DISPATCH_MEGA_COMBINE_H
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/dispatch_mega_combine_tiling.h+88-26
@@ -114,29 +114,7 @@ struct MegaMoeGmm1Tiling {
114struct MegaMoeSwigluTiling {114struct MegaMoeSwigluTiling {
115 uint64_t gmPermutedTokenOffset = 0;115 uint64_t gmPermutedTokenOffset = 0;
116 uint64_t perTokenScale2Offset = 0;116 uint64_t perTokenScale2Offset = 0;
117- uint64_t swigluSegmentMetaOffset = 0;117+ uint32_t reserved[12] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
118- uint64_t swigluSegmentMetaBytes = 0;
119- uint32_t reservedSwigluParams[5] = {0, 0, 0, 0, 0};
120- uint32_t reserved0[2] = {0, 0};
121-};
122- 
123-struct MegaMoeSwigluSegmentRuntimeMeta {
124- uint32_t segmentIdx = 0;
125- uint32_t segmentStartExpert = 0;
126- uint32_t segmentEndExpert = 0;
127- uint32_t segmentRowBase = 0;
128- uint32_t segmentRows = 0;
129- uint32_t cumsumRows = 0;
130- uint32_t expertTokenRows = 0;
131- uint32_t rowSplitBase = 0;
132- uint32_t rowSplitRem = 0;
133- uint32_t valid = 0;
134- uint32_t generation = 0;
135- uint32_t producerCoreIdx = 0;
136- uint32_t metadataMode = 0;
137- uint32_t segmentNum = 0;
138- uint32_t epilogueGranularity = 0;
139- uint32_t marker = 0;
140};118};
141 119 
142struct MegaMoeGmm2Tiling {120struct MegaMoeGmm2Tiling {
@@ -157,21 +135,104 @@ struct MegaMoeCombineTiling {
157 uint64_t reservedCombineScratchBytes = 0;135 uint64_t reservedCombineScratchBytes = 0;
158 uint64_t reservedCombineScratchBytesPerAiv = 0;136 uint64_t reservedCombineScratchBytesPerAiv = 0;
159 uint32_t reservedCombineTileCols = 0;137 uint32_t reservedCombineTileCols = 0;
160- uint32_t combineImplMode = 0;138+ uint32_t reserved0 = 0;
161};139};
162 140 
163struct MegaMoeUnpermuteTiling {141struct MegaMoeUnpermuteTiling {
164 uint32_t unpermuteTileCols = 1024;142 uint32_t unpermuteTileCols = 1024;
165 uint32_t unpermuteTokenBatch = 256;143 uint32_t unpermuteTokenBatch = 256;
166- uint32_t reservedUnpermuteLayoutVersion = 0;144+ uint32_t unpermuteImplMode = 0;
167 uint32_t reserved0[11] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};145 uint32_t reserved0[11] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
168};146};
169 147 
148+constexpr uint32_t kMegaMoeUnpermuteImplBarrier = 0U;
149+constexpr uint32_t kMegaMoeUnpermuteImplRankStreaming = 1U;
150+constexpr uint32_t kMegaMoeFixedLayoutVersion = 11U;
151+constexpr uint32_t kMegaMoeFixedPhysicalAicNum = 24U;
152+constexpr uint32_t kMegaMoeFixedPhysicalAivNum = 48U;
153+constexpr uint32_t kMegaMoeExpertProgressMaxRanks = 16U;
154+constexpr uint32_t kMegaMoeRankStreamingMaxTokensPerWorker = 256U;
155+constexpr uint32_t kMegaMoeFixedMaxExperts = 32U;
156+ 
157+constexpr uint32_t kMegaMoeFixedSyncSlotBytes = 64U;
158+constexpr uint32_t kMegaMoeFixedSyncHeadCanarySlot = 0U;
159+constexpr uint32_t kMegaMoeFixedSyncBytes = 8704U;
160+constexpr uint32_t kMegaMoeFixedHeadCanary = 0x13579BDFU;
161+constexpr uint32_t kMegaMoeFixedTailCanary = 0x2468ACE0U;
162+constexpr int32_t kMegaMoeFixedGmm1DoneMarker = 1;
163+constexpr int32_t kMegaMoeFixedGmm2JoinDecisionBit = 0x40000000;
164+constexpr int32_t kMegaMoeFixedGmm2JoinDecisionMask = kMegaMoeFixedGmm2JoinDecisionBit - 1;
165+ 
166+enum MegaMoeFixedRole : uint32_t {
167+ kMegaMoeFixedRoleGmm1 = 1U,
168+ kMegaMoeFixedRoleGmm2 = 2U,
169+ kMegaMoeFixedRoleSwiglu = 3U,
170+ kMegaMoeFixedRoleCombine = 4U,
171+ kMegaMoeFixedRoleDispatch = 5U,
172+};
173+ 
174+struct MegaMoeFixedGroupTiling {
175+ uint64_t syncOffset = 0;
176+ uint64_t syncBytes = 0;
177+ uint32_t layoutVersion = kMegaMoeFixedLayoutVersion;
178+ uint32_t physicalAicNum = 0;
179+ uint32_t physicalAivNum = 0;
180+ uint32_t dispatchGroupSize = 16U;
181+ uint32_t gmm1GroupSize = 16U;
182+ uint32_t gmm2GroupSize = 8U;
183+ uint32_t swigluGroupSize = 16U;
184+ uint32_t swigluActiveGroupSize = 16U;
185+ uint32_t combineGroupSize = 8U;
186+ uint32_t shapeConfigM = 0;
187+ uint32_t fullAicGmm1ExpertCount = 2U;
188+ uint32_t unpermutePhase1ReadyExpertCount = 13U;
189+ uint32_t gmm2JoinCheckStartExpert = 13U;
190+ uint32_t combineStartAfterGmm2Expert = 2U;
191+ uint32_t combineLargeLanesPerRank = 1U;
192+ uint32_t reserved0 = 0;
193+ uint32_t syncSlotBytes = kMegaMoeFixedSyncSlotBytes;
194+ uint32_t maxExperts = kMegaMoeFixedMaxExperts;
195+};
196+ 
197+struct MegaMoeSyncLayout {
198+ uint32_t dispatchArrivalBase = 0;
199+ uint32_t dispatchReadyBase = 0;
200+ uint32_t gmm1ArrivalBase = 0;
201+ uint32_t swigluReadyBase = 0;
202+ uint32_t swigluArrivalBase = 0;
203+ uint32_t gmm2ReadyBase = 0;
204+ uint32_t gmm2ArrivalBase = 0;
205+ uint32_t combineReadyBase = 0;
206+ uint32_t gmm1DoneSlot = 0;
207+ uint32_t gmm2JoinSlot = 0;
208+ uint32_t tailCanarySlot = 0;
209+ uint32_t slotCount = 0;
210+};
211+ 
212+template <typename FixedT>
213+inline MegaMoeSyncLayout MakeMegaMoeSyncLayout(const FixedT& fixed)
214+{
215+ MegaMoeSyncLayout layout;
216+ layout.dispatchArrivalBase = kMegaMoeFixedSyncHeadCanarySlot + 1U;
217+ layout.dispatchReadyBase = layout.dispatchArrivalBase + fixed.dispatchGroupSize;
218+ layout.gmm1ArrivalBase = layout.dispatchReadyBase + fixed.physicalAicNum;
219+ layout.swigluReadyBase = layout.gmm1ArrivalBase + fixed.gmm1GroupSize;
220+ layout.swigluArrivalBase = layout.swigluReadyBase + fixed.swigluGroupSize;
221+ layout.gmm2ReadyBase = layout.swigluArrivalBase + fixed.swigluGroupSize;
222+ layout.gmm2ArrivalBase = layout.gmm2ReadyBase + fixed.gmm2GroupSize;
223+ layout.combineReadyBase = layout.gmm2ArrivalBase + fixed.physicalAicNum;
224+ layout.gmm1DoneSlot = layout.combineReadyBase + fixed.gmm2GroupSize;
225+ layout.gmm2JoinSlot = layout.gmm1DoneSlot + 1U;
226+ layout.tailCanarySlot = layout.gmm2JoinSlot + 1U;
227+ layout.slotCount = layout.tailCanarySlot + 1U;
228+ return layout;
229+}
230+ 
170static_assert(sizeof(MegaMoeSwigluTiling) == 64);231static_assert(sizeof(MegaMoeSwigluTiling) == 64);
171-static_assert(sizeof(MegaMoeSwigluSegmentRuntimeMeta) == 64);
172static_assert(sizeof(MegaMoeGmm2Tiling) == 56);232static_assert(sizeof(MegaMoeGmm2Tiling) == 56);
173static_assert(sizeof(MegaMoeCombineTiling) == 48);233static_assert(sizeof(MegaMoeCombineTiling) == 48);
174static_assert(sizeof(MegaMoeUnpermuteTiling) == 56);234static_assert(sizeof(MegaMoeUnpermuteTiling) == 56);
235+static_assert(sizeof(MegaMoeFixedGroupTiling) == 88);
175 236 
176struct MegaMoeTilingData {237struct MegaMoeTilingData {
177 MegaMoeInfo megaMoeInfo;238 MegaMoeInfo megaMoeInfo;
@@ -183,4 +244,5 @@ struct MegaMoeTilingData {
183 MegaMoeGmm2Tiling gmm2Tiling;244 MegaMoeGmm2Tiling gmm2Tiling;
184 MegaMoeCombineTiling combineTiling;245 MegaMoeCombineTiling combineTiling;
185 MegaMoeUnpermuteTiling unpermuteTiling;246 MegaMoeUnpermuteTiling unpermuteTiling;
247+ MegaMoeFixedGroupTiling fixedGroupTiling;
186};248};
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/front_reorder.h+4-2
@@ -1011,7 +1011,7 @@ AICORE inline void FrontEndBuildCountExchangeAndPreSum(const FrontReorderCommonS
1011template <uint32_t Pitch>1011template <uint32_t Pitch>
1012AICORE inline void FrontEndBuildCumsumForPitch(const FrontReorderCommonState& op)1012AICORE inline void FrontEndBuildCumsumForPitch(const FrontReorderCommonState& op)
1013{1013{
1014- static_assert(Pitch == 8U || Pitch == 16U, "front cumsum supports expertPerRank 8 or 16");1014+ static_assert(Pitch == 8U || Pitch == 16U || Pitch == 32U, "front cumsum supports expertPerRank 8, 16 or 32");
1015 using CumsumTile = pto::Tile<pto::TileType::Vec, int32_t, 16, Pitch, pto::BLayout::RowMajor, -1, -1>;1015 using CumsumTile = pto::Tile<pto::TileType::Vec, int32_t, 16, Pitch, pto::BLayout::RowMajor, -1, -1>;
1016 using CumsumGlobal = pto::GlobalTensor<int32_t, FrontEndShapeDyn, FrontEndStrideDyn, pto::Layout::ND>;1016 using CumsumGlobal = pto::GlobalTensor<int32_t, FrontEndShapeDyn, FrontEndStrideDyn, pto::Layout::ND>;
1017 const uint32_t localBegin = op.rank_ * Pitch;1017 const uint32_t localBegin = op.rank_ * Pitch;
@@ -1050,7 +1050,9 @@ AICORE inline void FrontEndBuildCumsumForPitch(const FrontReorderCommonState& op
1050template <uint32_t ExpertPerRank>1050template <uint32_t ExpertPerRank>
1051AICORE inline void FrontEndBuildCumsumAndExpertTokenNums(const FrontReorderCommonState& op)1051AICORE inline void FrontEndBuildCumsumAndExpertTokenNums(const FrontReorderCommonState& op)
1052{1052{
1053- static_assert(ExpertPerRank == 8U || ExpertPerRank == 16U, "front cumsum supports expertPerRank 8 or 16");1053+ static_assert(
1054+ ExpertPerRank == 8U || ExpertPerRank == 16U || ExpertPerRank == 32U,
1055+ "front cumsum supports expertPerRank 8, 16 or 32");
1054 if (op.coreIdx_ == 0U) {1056 if (op.coreIdx_ == 0U) {
1055 FrontEndBuildCumsumForPitch<ExpertPerRank>(op);1057 FrontEndBuildCumsumForPitch<ExpertPerRank>(op);
1056 }1058 }
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/gmm1.h+78-22
@@ -17,9 +17,8 @@ See LICENSE in the root of the software repository for the full text of the Lice
17#include "gmm_common.h"17#include "gmm_common.h"
18#include "utils/common_helpers.hpp"18#include "utils/common_helpers.hpp"
19#include "utils/const_args.hpp"19#include "utils/const_args.hpp"
20-#include "utils/pto_sync_substrate.hpp"20+#include "utils/mega_expert_sync.hpp"
21 21 
22-constexpr uint32_t kGmm1InvalidTask = kGmmCommonInvalidTask;
23using Gmm1Pipeline = GmmCommonPipeline;22using Gmm1Pipeline = GmmCommonPipeline;
24 23 
25template <typename InputElement>24template <typename InputElement>
@@ -28,11 +27,23 @@ public:
28 AICORE inline void Init(27 AICORE inline void Init(
29 GM_ADDR weight1GM, GM_ADDR scale1GM, GM_ADDR expertTokenNumsGM, GM_ADDR workspaceGM,28 GM_ADDR weight1GM, GM_ADDR scale1GM, GM_ADDR expertTokenNumsGM, GM_ADDR workspaceGM,
30 const __gm__ MegaMoeTilingData* tilingData);29 const __gm__ MegaMoeTilingData* tilingData);
31- AICORE inline void Process();30+ AICORE inline void ProcessFixed(uint32_t groupLocalId, uint32_t groupSize);
32 31 
33private:32private:
34- AICORE inline uint32_t CoreLoops(uint32_t currentM) const33+ AICORE inline bool ParticipantBalancedNPartition(uint32_t currentM) const
35 {34 {
35+ const uint32_t tileM = tilingData_->gmm1Tiling.l1TileM;
36+ const uint32_t tileN = tilingData_->gmm1Tiling.l1TileN;
37+ const uint32_t defaultLoops = GmmCommonCoreLoops(currentM, problemN_, tileM, tileN);
38+ constexpr uint32_t kNAlign = 32U;
39+ return GmmCommonTileM(currentM, tileM) == 1U && defaultLoops < coreNum_ && problemN_ % kNAlign == 0U &&
40+ problemN_ / kNAlign >= coreNum_;
41+ }
42+ AICORE inline uint32_t CoreLoops(uint32_t currentM, bool participantBalanced) const
43+ {
44+ if (participantBalanced) {
45+ return coreNum_;
46+ }
36 return GmmCommonCoreLoops(47 return GmmCommonCoreLoops(
37 currentM, problemN_, tilingData_->gmm1Tiling.l1TileM, tilingData_->gmm1Tiling.l1TileN);48 currentM, problemN_, tilingData_->gmm1Tiling.l1TileM, tilingData_->gmm1Tiling.l1TileN);
38 }49 }
@@ -42,18 +53,48 @@ private:
42 }53 }
43 AICORE inline void WaitDispatchGroupReady(uint32_t groupIdx) const54 AICORE inline void WaitDispatchGroupReady(uint32_t groupIdx) const
44 {55 {
45- CrossCoreWaitFlag<0x2>(MegaMoeD2CHardFlagId(groupIdx + 1U));56+ WaitEpochAcquire(
57+ FixedSyncSlot(workspaceGM_, tilingData_, FixedSyncLayout(tilingData_).dispatchReadyBase + coreIdx_),
58+ static_cast<int32_t>(groupIdx * 2U + 2U));
59+ }
60+ AICORE inline GmmCommonTileInfo ParticipantBalancedTileInfo(uint32_t currentM, uint32_t loopIdx) const
61+ {
62+ constexpr uint32_t kNAlign = 32U;
63+ const uint32_t nUnits = problemN_ / kNAlign;
64+ const uint32_t baseUnits = nUnits / coreNum_;
65+ const uint32_t wideTileCount = nUnits % coreNum_;
66+ const uint32_t widePrefix = loopIdx < wideTileCount ? loopIdx : wideTileCount;
67+ 
68+ GmmCommonTileInfo tileInfo;
69+ tileInfo.tileM = 1U;
70+ tileInfo.tileN = coreNum_;
71+ tileInfo.blockM = 0U;
72+ tileInfo.blockN = loopIdx;
73+ tileInfo.actualM = currentM;
74+ tileInfo.actualN = (baseUnits + (loopIdx < wideTileCount ? 1U : 0U)) * kNAlign;
75+ tileInfo.blockRowStart = 0U;
76+ tileInfo.blockColStart = (loopIdx * baseUnits + widePrefix) * kNAlign;
77+ return tileInfo;
46 }78 }
47 AICORE inline void RunGmmTile(79 AICORE inline void RunGmmTile(
48- Gmm1Pipeline& gmmPipeline, uint32_t groupIdx, uint32_t groupBase, uint32_t currentM, uint32_t loopIdx) const80+ Gmm1Pipeline& gmmPipeline, uint32_t groupIdx, uint32_t groupBase, uint32_t currentM, bool participantBalanced,
81+ uint32_t loopIdx) const
49 {82 {
83+ if (participantBalanced) {
84+ const GmmCommonTileInfo tileInfo = ParticipantBalancedTileInfo(currentM, loopIdx);
85+ GmmCommonRunTileInfo(
86+ gmmPipeline, gmAPtr_, weight1Ptr_, gmCPtr_, scale1Ptr_, groupIdx, groupBase, tileInfo, problemN_,
87+ problemK_, problemK_, problemK_, problemN_, problemN_);
88+ return;
89+ }
50 GmmCommonRunTile(90 GmmCommonRunTile(
51 gmmPipeline, gmAPtr_, weight1Ptr_, gmCPtr_, scale1Ptr_, groupIdx, groupBase, currentM, loopIdx, problemN_,91 gmmPipeline, gmAPtr_, weight1Ptr_, gmCPtr_, scale1Ptr_, groupIdx, groupBase, currentM, loopIdx, problemN_,
52 problemK_, problemK_, problemK_, problemN_, problemN_, tilingData_->gmm1Tiling.l1TileM,92 problemK_, problemK_, problemK_, problemN_, problemN_, tilingData_->gmm1Tiling.l1TileM,
53 tilingData_->gmm1Tiling.l1TileN);93 tilingData_->gmm1Tiling.l1TileN);
54 }94 }
55- AICORE inline void SetC2VReady(uint32_t segmentIdx) const;95+ AICORE inline void ProcessImpl();
56 96 
97+ GM_ADDR workspaceGM_ = nullptr;
57 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;98 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;
58 99 
59 __gm__ int8_t* gmAPtr_ = nullptr;100 __gm__ int8_t* gmAPtr_ = nullptr;
@@ -77,6 +118,7 @@ AICORE inline void Gmm1<InputElement>::Init(
77 const __gm__ MegaMoeTilingData* tilingData)118 const __gm__ MegaMoeTilingData* tilingData)
78{119{
79 (void)expertTokenNumsGM;120 (void)expertTokenNumsGM;
121+ workspaceGM_ = workspaceGM;
80 tilingData_ = tilingData;122 tilingData_ = tilingData;
81 123 
82 problemK_ = tilingData_->megaMoeInfo.K;124 problemK_ = tilingData_->megaMoeInfo.K;
@@ -93,42 +135,56 @@ AICORE inline void Gmm1<InputElement>::Init(
93 scale1Ptr_ = reinterpret_cast<__gm__ uint64_t*>(scale1GM);135 scale1Ptr_ = reinterpret_cast<__gm__ uint64_t*>(scale1GM);
94 cumsumMMPtr_ = reinterpret_cast<__gm__ int32_t*>(workspaceGM + tilingData_->frontReorderTiling.cumsumMMOffset);136 cumsumMMPtr_ = reinterpret_cast<__gm__ int32_t*>(workspaceGM + tilingData_->frontReorderTiling.cumsumMMOffset);
95}137}
138+ 
96template <typename InputElement>139template <typename InputElement>
97-AICORE inline void Gmm1<InputElement>::SetC2VReady(uint32_t segmentIdx) const140+AICORE inline void Gmm1<InputElement>::ProcessFixed(uint32_t groupLocalId, uint32_t groupSize)
98{141{
99- CrossCoreSetFlag<0x2, PIPE_FIX>(MegaMoeC2VHardFlagId(segmentIdx));142+ coreIdx_ = groupLocalId;
143+ coreNum_ = groupSize;
144+ ProcessImpl();
100}145}
101 146 
102template <typename InputElement>147template <typename InputElement>
103-AICORE inline void Gmm1<InputElement>::Process()148+AICORE inline void Gmm1<InputElement>::ProcessImpl()
104{149{
105 if ASCEND_IS_AIV {150 if ASCEND_IS_AIV {
106 return;151 return;
107 }152 }
108- CrossCoreWaitFlag<0x2>(MegaMoeD2CHardFlagId(0U));
109 Gmm1Pipeline gmmPipeline;153 Gmm1Pipeline gmmPipeline;
110 uint32_t groupBase = 0;154 uint32_t groupBase = 0;
111 uint32_t startCoreIdx = 0;155 uint32_t startCoreIdx = 0;
112- uint32_t segmentIdx = 0;
113- const uint32_t firstSegmentEnd = MoeSwigluEpilogueGranularity(expertPerRank_);
114 for (uint32_t groupIdx = 0; groupIdx < expertPerRank_; ++groupIdx) {156 for (uint32_t groupIdx = 0; groupIdx < expertPerRank_; ++groupIdx) {
157+ const __gm__ MegaMoeFixedGroupTiling& fixed = tilingData_->fixedGroupTiling;
158+ coreNum_ = groupIdx < fixed.fullAicGmm1ExpertCount ? fixed.physicalAicNum : fixed.gmm1GroupSize;
159+ if (coreIdx_ >= coreNum_) {
160+ break;
161+ }
162+ startCoreIdx %= coreNum_;
115 WaitDispatchGroupReady(groupIdx);163 WaitDispatchGroupReady(groupIdx);
116 const uint32_t currentMRaw = MoeCurrentMRaw(cumsumMMPtr_, rankSize_, expertPerRank_, groupIdx);164 const uint32_t currentMRaw = MoeCurrentMRaw(cumsumMMPtr_, rankSize_, expertPerRank_, groupIdx);
117 const uint32_t currentM = MoeClipCurrentM(currentMRaw, groupBase, maxOutputSize_);165 const uint32_t currentM = MoeClipCurrentM(currentMRaw, groupBase, maxOutputSize_);
118- const uint32_t coreLoops = CoreLoops(currentM); // 当前 expert 一共有多少个 GMM tile166+ const bool participantBalanced = ParticipantBalancedNPartition(currentM);
167+ const uint32_t coreLoops = CoreLoops(currentM, participantBalanced); // 当前 expert 一共有多少个 GMM tile
119 const uint32_t startLoopIdx = StartLoopIdx(startCoreIdx); // 当前 AIC 在这个 expert 里的第一个 tile id168 const uint32_t startLoopIdx = StartLoopIdx(startCoreIdx); // 当前 AIC 在这个 expert 里的第一个 tile id
120 for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += coreNum_) {169 for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += coreNum_) {
121- RunGmmTile(gmmPipeline, groupIdx, groupBase, currentM, loopIdx);170+ RunGmmTile(gmmPipeline, groupIdx, groupBase, currentM, participantBalanced, loopIdx);
122 }171 }
123- const uint32_t groupEnd = groupIdx + 1U;172+ gmmPipeline.SynchronizeBlock();
124- if (groupEnd == firstSegmentEnd || groupEnd == expertPerRank_) {173+ if (groupIdx < fixed.fullAicGmm1ExpertCount) {
125- gmmPipeline.SynchronizeBlock();174+ pto::SYNCALL<pto::SyncCoreType::AICOnly>();
126- SetC2VReady(segmentIdx);175+ if (coreIdx_ == 0U) {
127- ++segmentIdx;176+ pipe_barrier(PIPE_ALL);
177+ dsb(DSB_DDR);
178+ PublishScalarEpochRange(
179+ workspaceGM_, tilingData_, FixedSyncLayout(tilingData_).swigluReadyBase,
180+ fixed.swigluActiveGroupSize, static_cast<int32_t>(groupIdx * 2U + 2U));
181+ }
182+ } else {
183+ PublishGroupArrival(
184+ workspaceGM_, tilingData_, FixedSyncLayout(tilingData_).gmm1ArrivalBase, coreIdx_, groupIdx);
128 }185 }
129 groupBase += currentM;186 groupBase += currentM;
130- // 从上一个expert分配结束后的下一个 AIC 继续开始下一个exprt,负载均衡187+ startCoreIdx = GmmCommonNextStartCoreIdx(startCoreIdx, coreNum_, coreLoops);
131- startCoreIdx = (startCoreIdx + coreLoops) % coreNum_;
132 }188 }
133}189}
134 190 
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/gmm2.h+130-33
@@ -17,9 +17,8 @@ See LICENSE in the root of the software repository for the full text of the Lice
17#include "gmm_common.h"17#include "gmm_common.h"
18#include "utils/common_helpers.hpp"18#include "utils/common_helpers.hpp"
19#include "utils/const_args.hpp"19#include "utils/const_args.hpp"
20-#include "utils/pto_sync_substrate.hpp"20+#include "utils/mega_expert_sync.hpp"
21 21 
22-constexpr uint32_t kGmm2InvalidTask = kGmmCommonInvalidTask;
23using Gmm2Pipeline = GmmCommonPipeline;22using Gmm2Pipeline = GmmCommonPipeline;
24 23 
25template <typename InputElement>24template <typename InputElement>
@@ -28,7 +27,8 @@ public:
28 AICORE inline void Init(27 AICORE inline void Init(
29 GM_ADDR weight2GM, GM_ADDR scale2GM, GM_ADDR expertTokenNumsGM, GM_ADDR workspaceGM,28 GM_ADDR weight2GM, GM_ADDR scale2GM, GM_ADDR expertTokenNumsGM, GM_ADDR workspaceGM,
30 const __gm__ MegaMoeTilingData* tilingData);29 const __gm__ MegaMoeTilingData* tilingData);
31- AICORE inline void Process();30+ AICORE inline void ProcessFixed(uint32_t groupLocalId, uint32_t groupSize);
31+ AICORE inline void ProcessFixedHelper(uint32_t groupLocalId);
32 32 
33private:33private:
34 AICORE inline uint32_t CoreLoops(uint32_t currentM) const34 AICORE inline uint32_t CoreLoops(uint32_t currentM) const
@@ -47,7 +47,13 @@ private:
47 loopIdx, outputN_, inputK_, inputK_, inputK_, outputN_, outputN_, tilingData_->gmm2Tiling.l1TileM,47 loopIdx, outputN_, inputK_, inputK_, inputK_, outputN_, outputN_, tilingData_->gmm2Tiling.l1TileM,
48 tilingData_->gmm2Tiling.l1TileN);48 tilingData_->gmm2Tiling.l1TileN);
49 }49 }
50+ AICORE inline uint32_t GroupBaseBefore(uint32_t groupIdx) const;
51+ AICORE inline uint32_t StartCoreBefore(uint32_t groupIdx, uint32_t coreNum) const;
52+ AICORE inline int32_t PrimaryJoinDecision(uint32_t groupIdx) const;
53+ AICORE inline uint32_t HelperJoinExpert() const;
54+ AICORE inline void ProcessImpl(bool helperGroup);
50 55 
56+ GM_ADDR workspaceGM_ = nullptr;
51 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;57 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;
52 58 
53 __gm__ int8_t* gmPermutedTokenPtr_ = nullptr;59 __gm__ int8_t* gmPermutedTokenPtr_ = nullptr;
@@ -55,7 +61,6 @@ private:
55 __gm__ int8_t* weight2Ptr_ = nullptr;61 __gm__ int8_t* weight2Ptr_ = nullptr;
56 __gm__ uint64_t* scale2Ptr_ = nullptr;62 __gm__ uint64_t* scale2Ptr_ = nullptr;
57 __gm__ int32_t* cumsumMMPtr_ = nullptr;63 __gm__ int32_t* cumsumMMPtr_ = nullptr;
58- __gm__ int32_t* expertTokenNumsPtr_ = nullptr;
59 64 
60 uint32_t inputK_ = 0;65 uint32_t inputK_ = 0;
61 uint32_t outputN_ = 0;66 uint32_t outputN_ = 0;
@@ -64,6 +69,7 @@ private:
64 uint32_t rankSize_ = 0;69 uint32_t rankSize_ = 0;
65 uint32_t coreIdx_ = 0;70 uint32_t coreIdx_ = 0;
66 uint32_t coreNum_ = 1;71 uint32_t coreNum_ = 1;
72+ uint32_t primaryLocalId_ = 0;
67};73};
68 74 
69template <typename InputElement>75template <typename InputElement>
@@ -72,6 +78,8 @@ AICORE inline void Gmm2<InputElement>::Init(
72 const __gm__ MegaMoeTilingData* tilingData)78 const __gm__ MegaMoeTilingData* tilingData)
73{79{
74 (void)sizeof(InputElement);80 (void)sizeof(InputElement);
81+ (void)expertTokenNumsGM;
82+ workspaceGM_ = workspaceGM;
75 tilingData_ = tilingData;83 tilingData_ = tilingData;
76 84 
77 const uint32_t problemN = tilingData_->megaMoeInfo.N;85 const uint32_t problemN = tilingData_->megaMoeInfo.N;
@@ -90,46 +98,135 @@ AICORE inline void Gmm2<InputElement>::Init(
90 weight2Ptr_ = reinterpret_cast<__gm__ int8_t*>(weight2GM);98 weight2Ptr_ = reinterpret_cast<__gm__ int8_t*>(weight2GM);
91 scale2Ptr_ = reinterpret_cast<__gm__ uint64_t*>(scale2GM);99 scale2Ptr_ = reinterpret_cast<__gm__ uint64_t*>(scale2GM);
92 cumsumMMPtr_ = reinterpret_cast<__gm__ int32_t*>(workspaceGM + tilingData_->frontReorderTiling.cumsumMMOffset);100 cumsumMMPtr_ = reinterpret_cast<__gm__ int32_t*>(workspaceGM + tilingData_->frontReorderTiling.cumsumMMOffset);
93- expertTokenNumsPtr_ = reinterpret_cast<__gm__ int32_t*>(expertTokenNumsGM);
94}101}
95 102 
96template <typename InputElement>103template <typename InputElement>
97-AICORE inline void Gmm2<InputElement>::Process()104+AICORE inline void Gmm2<InputElement>::ProcessFixed(uint32_t groupLocalId, uint32_t groupSize)
105+{
106+ coreIdx_ = groupLocalId;
107+ coreNum_ = groupSize;
108+ primaryLocalId_ = groupLocalId;
109+ ProcessImpl(false);
110+}
111+ 
112+template <typename InputElement>
113+AICORE inline void Gmm2<InputElement>::ProcessFixedHelper(uint32_t groupLocalId)
114+{
115+ coreIdx_ = groupLocalId;
116+ coreNum_ = tilingData_->fixedGroupTiling.physicalAicNum;
117+ primaryLocalId_ = groupLocalId;
118+ ProcessImpl(true);
119+}
120+ 
121+template <typename InputElement>
122+AICORE inline uint32_t Gmm2<InputElement>::GroupBaseBefore(uint32_t groupIdx) const
123+{
124+ uint32_t groupBase = 0U;
125+ for (uint32_t expert = 0U; expert < groupIdx; ++expert) {
126+ const uint32_t currentMRaw = MoeCurrentMRaw(cumsumMMPtr_, rankSize_, expertPerRank_, expert);
127+ groupBase += MoeClipCurrentM(currentMRaw, groupBase, maxOutputSize_);
128+ }
129+ return groupBase;
130+}
131+ 
132+template <typename InputElement>
133+AICORE inline uint32_t Gmm2<InputElement>::StartCoreBefore(uint32_t groupIdx, uint32_t coreNum) const
134+{
135+ uint32_t groupBase = 0U;
136+ uint32_t startCoreIdx = 0U;
137+ for (uint32_t expert = 0U; expert < groupIdx; ++expert) {
138+ const uint32_t currentMRaw = MoeCurrentMRaw(cumsumMMPtr_, rankSize_, expertPerRank_, expert);
139+ const uint32_t currentM = MoeClipCurrentM(currentMRaw, groupBase, maxOutputSize_);
140+ groupBase += currentM;
141+ startCoreIdx = GmmCommonNextStartCoreIdx(startCoreIdx, coreNum, CoreLoops(currentM));
142+ }
143+ return startCoreIdx;
144+}
145+ 
146+template <typename InputElement>
147+AICORE inline int32_t Gmm2<InputElement>::PrimaryJoinDecision(uint32_t groupIdx) const
148+{
149+ volatile __gm__ int32_t* joinSlot =
150+ FixedSyncSlot(workspaceGM_, tilingData_, FixedSyncLayout(tilingData_).gmm2JoinSlot);
151+ const int32_t expectedDecision = static_cast<int32_t>(groupIdx + 1U);
152+ if (primaryLocalId_ != 0U) {
153+ return WaitEpochAcquire(joinSlot, expectedDecision);
154+ }
155+ 
156+ const int32_t done =
157+ ReadScalarEpoch(FixedSyncSlot(workspaceGM_, tilingData_, FixedSyncLayout(tilingData_).gmm1DoneSlot));
158+ const int32_t decision =
159+ done >= kMegaMoeFixedGmm1DoneMarker ? (kMegaMoeFixedGmm2JoinDecisionBit | expectedDecision) : expectedDecision;
160+ PublishScalarEpoch(joinSlot, decision);
161+ return decision;
162+}
163+ 
164+template <typename InputElement>
165+AICORE inline uint32_t Gmm2<InputElement>::HelperJoinExpert() const
166+{
167+ const int32_t decision = WaitEpochAcquire(
168+ FixedSyncSlot(workspaceGM_, tilingData_, FixedSyncLayout(tilingData_).gmm2JoinSlot),
169+ kMegaMoeFixedGmm2JoinDecisionBit);
170+ const int32_t encodedExpert = decision & kMegaMoeFixedGmm2JoinDecisionMask;
171+ return encodedExpert > 0 ? static_cast<uint32_t>(encodedExpert - 1) : expertPerRank_;
172+}
173+ 
174+template <typename InputElement>
175+AICORE inline void Gmm2<InputElement>::ProcessImpl(bool helperGroup)
98{176{
99 if ASCEND_IS_AIV {177 if ASCEND_IS_AIV {
100 return;178 return;
101 }179 }
102- 
103 Gmm2Pipeline gmmPipeline;180 Gmm2Pipeline gmmPipeline;
104 uint32_t groupBase = 0;181 uint32_t groupBase = 0;
105 uint32_t startCoreIdx = 0;182 uint32_t startCoreIdx = 0;
106- const uint32_t segmentNum = MoeSwigluSegmentNum(expertPerRank_);183+ uint32_t firstGroupIdx = 0U;
107- for (uint32_t segmentIdx = 0; segmentIdx < segmentNum; ++segmentIdx) {184+ bool joinedGroup = helperGroup;
108- CrossCoreWaitFlag<0x2>(MegaMoeV2CHardFlagId(segmentIdx));185+ if (helperGroup) {
109- 186+ firstGroupIdx = HelperJoinExpert();
110- uint32_t segmentStartExpert = 0;187+ if (firstGroupIdx >= expertPerRank_) {
111- uint32_t segmentEndExpert = 0;188+ return;
112- uint32_t segmentRowBase = 0;
113- uint32_t segmentRows = 0;
114- uint32_t cumsumRows = 0;
115- uint32_t expertTokenRows = 0;
116- MoeBuildSegmentMetadata(
117- segmentIdx, expertPerRank_, maxOutputSize_, cumsumMMPtr_, expertTokenNumsPtr_, rankSize_,
118- segmentStartExpert, segmentEndExpert, segmentRowBase, segmentRows, cumsumRows, expertTokenRows);
119- 
120- for (uint32_t groupIdx = segmentStartExpert; groupIdx < segmentEndExpert; ++groupIdx) {
121- const uint32_t currentMRaw = MoeCurrentMRaw(cumsumMMPtr_, rankSize_, expertPerRank_, groupIdx);
122- const uint32_t currentM = MoeClipCurrentM(currentMRaw, groupBase, maxOutputSize_);
123- const uint32_t coreLoops = CoreLoops(currentM);
124- const uint32_t startLoopIdx = StartLoopIdx(startCoreIdx);
125- for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += coreNum_) {
126- RunGmmTile(gmmPipeline, groupIdx, groupBase, currentM, loopIdx);
127- }
128- gmmPipeline.SynchronizeBlock();
129- gmmPipeline.Finalize(static_cast<int32_t>(groupIdx), MEGA_MOE_GMM2_TO_COMBINE_HARD_FLAG_BASE);
130- groupBase += currentM;
131- startCoreIdx = (startCoreIdx + coreLoops) % coreNum_;
132 }189 }
190+ groupBase = GroupBaseBefore(firstGroupIdx);
191+ startCoreIdx = StartCoreBefore(firstGroupIdx, tilingData_->fixedGroupTiling.gmm2GroupSize);
192+ }
193+ for (uint32_t groupIdx = firstGroupIdx; groupIdx < expertPerRank_; ++groupIdx) {
194+ const MegaMoeSyncLayout sync = FixedSyncLayout(tilingData_);
195+ const uint32_t readyLocalId = coreIdx_ % tilingData_->fixedGroupTiling.gmm2GroupSize;
196+ WaitEpochAcquire(
197+ FixedSyncSlot(workspaceGM_, tilingData_, sync.gmm2ReadyBase + readyLocalId),
198+ static_cast<int32_t>(groupIdx * 2U + 2U));
199+ if (!helperGroup && !joinedGroup && groupIdx >= tilingData_->fixedGroupTiling.gmm2JoinCheckStartExpert) {
200+ const int32_t decision = PrimaryJoinDecision(groupIdx);
201+ const uint32_t encodedJoinExpert = static_cast<uint32_t>(decision & kMegaMoeFixedGmm2JoinDecisionMask);
202+ const bool joinThisExpert = (decision & kMegaMoeFixedGmm2JoinDecisionBit) != 0 && encodedJoinExpert != 0U &&
203+ groupIdx + 1U >= encodedJoinExpert;
204+ if (joinThisExpert) {
205+ joinedGroup = true;
206+ coreIdx_ = tilingData_->fixedGroupTiling.gmm1GroupSize + primaryLocalId_;
207+ coreNum_ = tilingData_->fixedGroupTiling.physicalAicNum;
208+ }
209+ }
210+ const uint32_t currentMRaw = MoeCurrentMRaw(cumsumMMPtr_, rankSize_, expertPerRank_, groupIdx);
211+ const uint32_t currentM = MoeClipCurrentM(currentMRaw, groupBase, maxOutputSize_);
212+ const uint32_t coreLoops = CoreLoops(currentM);
213+ const uint32_t startLoopIdx = StartLoopIdx(startCoreIdx);
214+ for (uint32_t loopIdx = startLoopIdx; loopIdx < coreLoops; loopIdx += coreNum_) {
215+ RunGmmTile(gmmPipeline, groupIdx, groupBase, currentM, loopIdx);
216+ }
217+ gmmPipeline.SynchronizeBlock();
218+ if (tilingData_->frontReorderTiling.stageNum >= 13U) {
219+ const uint32_t arrivalLocalId =
220+ helperGroup ? primaryLocalId_ : tilingData_->fixedGroupTiling.gmm1GroupSize + primaryLocalId_;
221+ PublishGroupArrival(workspaceGM_, tilingData_, sync.gmm2ArrivalBase, arrivalLocalId, groupIdx);
222+ }
223+ groupBase += currentM;
224+ startCoreIdx = GmmCommonNextStartCoreIdx(startCoreIdx, coreNum_, coreLoops);
225+ }
226+ if (!helperGroup && !joinedGroup && primaryLocalId_ == 0U) {
227+ const int32_t sentinel = kMegaMoeFixedGmm2JoinDecisionBit + static_cast<int32_t>(expertPerRank_ + 1U);
228+ PublishScalarEpoch(
229+ FixedSyncSlot(workspaceGM_, tilingData_, FixedSyncLayout(tilingData_).gmm2JoinSlot), sentinel);
133 }230 }
134}231}
135 232 
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/gmm_common.h+37-41
@@ -17,7 +17,6 @@ See LICENSE in the root of the software repository for the full text of the Lice
17#include "utils/const_args.hpp"17#include "utils/const_args.hpp"
18#include "utils/pto_gmm_preload_async_fixpipe_quant.hpp"18#include "utils/pto_gmm_preload_async_fixpipe_quant.hpp"
19 19 
20-constexpr uint32_t kGmmCommonInvalidTask = 0xFFFFFFFFU;
21constexpr uint32_t kGmmCommonSwizzleOffset = 9U;20constexpr uint32_t kGmmCommonSwizzleOffset = 9U;
22 21 
23using GmmCommonPipeline =22using GmmCommonPipeline =
@@ -54,16 +53,36 @@ AICORE inline uint32_t GmmCommonStartLoopIdx(uint32_t coreIdx, uint32_t coreNum,
54 return ((coreIdx < startCoreIdx) ? (coreIdx + coreNum) : coreIdx) - startCoreIdx;53 return ((coreIdx < startCoreIdx) ? (coreIdx + coreNum) : coreIdx) - startCoreIdx;
55}54}
56 55 
56+AICORE inline uint32_t GmmCommonNextStartCoreIdx(uint32_t startCoreIdx, uint32_t coreNum, uint32_t coreLoops)
57+{
58+ if (coreNum > 1U && coreLoops != 0U) {
59+ const uint32_t remainder = coreLoops % coreNum;
60+ const uint32_t advance = remainder == 0U ? 1U : remainder;
61+ return (startCoreIdx + advance) % coreNum;
62+ }
63+ return coreNum <= 1U ? 0U : startCoreIdx;
64+}
65+ 
57AICORE inline void GmmCommonGetBlockCoordMN(66AICORE inline void GmmCommonGetBlockCoordMN(
58 uint32_t loopIdx, uint32_t tileM, uint32_t tileN, uint32_t& blockM, uint32_t& blockN)67 uint32_t loopIdx, uint32_t tileM, uint32_t tileN, uint32_t& blockM, uint32_t& blockN)
59{68{
69+ blockM = 0U;
70+ blockN = 0U;
71+ if (tileM == 0U || tileN == 0U) {
72+ return;
73+ }
74+ const uint32_t swizzleSpan = kGmmCommonSwizzleOffset * tileM;
60 const uint32_t tileBlockLoop = static_cast<uint32_t>(ceilDiv(tileN, kGmmCommonSwizzleOffset));75 const uint32_t tileBlockLoop = static_cast<uint32_t>(ceilDiv(tileN, kGmmCommonSwizzleOffset));
61- const uint32_t tileBlockIdx = loopIdx / (kGmmCommonSwizzleOffset * tileM);76+ const uint32_t tileBlockIdx = loopIdx / swizzleSpan;
62- const uint32_t inTileBlockIdx = loopIdx % (kGmmCommonSwizzleOffset * tileM);77+ const uint32_t inTileBlockIdx = loopIdx % swizzleSpan;
63 uint32_t nCol = kGmmCommonSwizzleOffset;78 uint32_t nCol = kGmmCommonSwizzleOffset;
64 if (tileBlockIdx + 1U == tileBlockLoop) {79 if (tileBlockIdx + 1U == tileBlockLoop) {
65 nCol = tileN - kGmmCommonSwizzleOffset * tileBlockIdx;80 nCol = tileN - kGmmCommonSwizzleOffset * tileBlockIdx;
66 }81 }
82+ if (nCol == 0U) {
83+ blockN = tileBlockIdx * kGmmCommonSwizzleOffset;
84+ return;
85+ }
67 blockM = inTileBlockIdx / nCol;86 blockM = inTileBlockIdx / nCol;
68 blockN = tileBlockIdx * kGmmCommonSwizzleOffset + inTileBlockIdx % nCol;87 blockN = tileBlockIdx * kGmmCommonSwizzleOffset + inTileBlockIdx % nCol;
69 if ((tileBlockIdx & 1U) != 0U) {88 if ((tileBlockIdx & 1U) != 0U) {
@@ -107,11 +126,6 @@ AICORE inline uint32_t MoeCurrentMRaw(
107 return static_cast<uint32_t>(cumsumMMPtr[static_cast<uint64_t>(rankSize - 1U) * expertPerRank + groupIdx]);126 return static_cast<uint32_t>(cumsumMMPtr[static_cast<uint64_t>(rankSize - 1U) * expertPerRank + groupIdx]);
108}127}
109 128 
110-AICORE inline uint32_t MoeExpertTokenNums(__gm__ int32_t* expertTokenNumsPtr, uint32_t groupIdx)
111-{
112- return static_cast<uint32_t>(expertTokenNumsPtr[groupIdx]);
113-}
114- 
115AICORE inline uint32_t MoeClipCurrentM(uint32_t currentMRaw, uint32_t groupBase, uint32_t maxOutputSize)129AICORE inline uint32_t MoeClipCurrentM(uint32_t currentMRaw, uint32_t groupBase, uint32_t maxOutputSize)
116{130{
117 if (groupBase >= maxOutputSize) {131 if (groupBase >= maxOutputSize) {
@@ -121,42 +135,12 @@ AICORE inline uint32_t MoeClipCurrentM(uint32_t currentMRaw, uint32_t groupBase,
121 return currentMRaw > remaining ? remaining : currentMRaw;135 return currentMRaw > remaining ? remaining : currentMRaw;
122}136}
123 137 
124-AICORE inline void MoeBuildSegmentMetadata(138+AICORE inline void GmmCommonRunTileInfo(
125- uint32_t segmentIdx, uint32_t expertPerRank, uint32_t maxOutputSize, __gm__ int32_t* cumsumMMPtr,
126- __gm__ int32_t* expertTokenNumsPtr, uint32_t rankSize, uint32_t& segmentStartExpert, uint32_t& segmentEndExpert,
127- uint32_t& segmentRowBase, uint32_t& segmentRows, uint32_t& cumsumRows, uint32_t& expertTokenRows)
128-{
129- segmentStartExpert = MoeSwigluSegmentStartExpert(expertPerRank, segmentIdx);
130- segmentEndExpert = MoeSwigluSegmentEndExpert(expertPerRank, segmentIdx);
131- segmentRowBase = 0U;
132- segmentRows = 0U;
133- cumsumRows = 0U;
134- expertTokenRows = 0U;
135- 
136- uint32_t groupBase = 0U;
137- for (uint32_t groupIdx = 0U; groupIdx < segmentEndExpert; ++groupIdx) {
138- const uint32_t currentMRaw = MoeCurrentMRaw(cumsumMMPtr, rankSize, expertPerRank, groupIdx);
139- const uint32_t currentM = MoeClipCurrentM(currentMRaw, groupBase, maxOutputSize);
140- if (groupIdx == segmentStartExpert) {
141- segmentRowBase = groupBase;
142- }
143- if (groupIdx >= segmentStartExpert) {
144- segmentRows += currentM;
145- cumsumRows += currentMRaw;
146- expertTokenRows += MoeExpertTokenNums(expertTokenNumsPtr, groupIdx);
147- }
148- groupBase += currentM;
149- }
150-}
151- 
152-AICORE inline void GmmCommonRunTile(
153 GmmCommonPipeline& gmmPipeline, __gm__ int8_t* gmAPtr, __gm__ int8_t* gmWeightPtr, __gm__ half* gmCPtr,139 GmmCommonPipeline& gmmPipeline, __gm__ int8_t* gmAPtr, __gm__ int8_t* gmWeightPtr, __gm__ half* gmCPtr,
154- __gm__ uint64_t* gmScalePtr, uint32_t groupIdx, uint32_t groupBase, uint32_t currentM, uint32_t loopIdx,140+ __gm__ uint64_t* gmScalePtr, uint32_t groupIdx, uint32_t groupBase, const GmmCommonTileInfo& tileInfo,
155 uint32_t problemN, uint32_t actualK, uint32_t aLeadingDim, uint32_t bFullRows, uint32_t cLeadingDim,141 uint32_t problemN, uint32_t actualK, uint32_t aLeadingDim, uint32_t bFullRows, uint32_t cLeadingDim,
156- uint32_t scaleGroupStride, uint32_t l1TileM, uint32_t l1TileN)142+ uint32_t scaleGroupStride)
157{143{
158- const GmmCommonTileInfo tileInfo = GmmCommonBuildTileInfo(currentM, problemN, l1TileM, l1TileN, loopIdx);
159- 
160 const uint64_t gmOffsetA = static_cast<uint64_t>(groupBase + tileInfo.blockRowStart) * aLeadingDim;144 const uint64_t gmOffsetA = static_cast<uint64_t>(groupBase + tileInfo.blockRowStart) * aLeadingDim;
161 const uint64_t gmOffsetB = static_cast<uint64_t>(groupIdx) * GmmCommonPackedWeightExpertStride(bFullRows, problemN);145 const uint64_t gmOffsetB = static_cast<uint64_t>(groupIdx) * GmmCommonPackedWeightExpertStride(bFullRows, problemN);
162 const uint64_t gmOffsetC = static_cast<uint64_t>(groupBase + tileInfo.blockRowStart) * cLeadingDim +146 const uint64_t gmOffsetC = static_cast<uint64_t>(groupBase + tileInfo.blockRowStart) * cLeadingDim +
@@ -178,4 +162,16 @@ AICORE inline void GmmCommonRunTile(
178 gmmPipeline.RunTile(blockParams);162 gmmPipeline.RunTile(blockParams);
179}163}
180 164 
165+AICORE inline void GmmCommonRunTile(
166+ GmmCommonPipeline& gmmPipeline, __gm__ int8_t* gmAPtr, __gm__ int8_t* gmWeightPtr, __gm__ half* gmCPtr,
167+ __gm__ uint64_t* gmScalePtr, uint32_t groupIdx, uint32_t groupBase, uint32_t currentM, uint32_t loopIdx,
168+ uint32_t problemN, uint32_t actualK, uint32_t aLeadingDim, uint32_t bFullRows, uint32_t cLeadingDim,
169+ uint32_t scaleGroupStride, uint32_t l1TileM, uint32_t l1TileN)
170+{
171+ const GmmCommonTileInfo tileInfo = GmmCommonBuildTileInfo(currentM, problemN, l1TileM, l1TileN, loopIdx);
172+ GmmCommonRunTileInfo(
173+ gmmPipeline, gmAPtr, gmWeightPtr, gmCPtr, gmScalePtr, groupIdx, groupBase, tileInfo, problemN, actualK,
174+ aLeadingDim, bFullRows, cLeadingDim, scaleGroupStride);
175+}
176+ 
181#endif // DISPATCH_MEGA_COMBINE_GMM_COMMON_H177#endif // DISPATCH_MEGA_COMBINE_GMM_COMMON_H
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/swiglu.h+34-102
@@ -19,12 +19,9 @@ See LICENSE in the root of the software repository for the full text of the Lice
19#include "gmm_common.h"19#include "gmm_common.h"
20#include "utils/common_helpers.hpp"20#include "utils/common_helpers.hpp"
21#include "utils/const_args.hpp"21#include "utils/const_args.hpp"
22+#include "utils/mega_expert_sync.hpp"
22#include "utils/pto_vector.hpp"23#include "utils/pto_vector.hpp"
23-#include "utils/pto_sync_substrate.hpp"
24 24 
25-constexpr uint32_t kSwigluWaitSourceC2VOnly = 1U;
26-constexpr uint32_t kSwigluPipelineModeInputOutputSplit = 1U;
27-constexpr uint32_t kSwigluMetadataModeSharedSegmentMeta = 1U;
28constexpr uint32_t kSwigluVecTileElems = 1024U;25constexpr uint32_t kSwigluVecTileElems = 1024U;
29constexpr uint32_t kSwigluFullRowIoBlockChunks = 4U;26constexpr uint32_t kSwigluFullRowIoBlockChunks = 4U;
30constexpr uint32_t kSwigluUbStageNum = 2U;27constexpr uint32_t kSwigluUbStageNum = 2U;
@@ -37,19 +34,9 @@ template <typename InputElement>
37class Swiglu {34class Swiglu {
38public:35public:
39 AICORE inline void Init(GM_ADDR expertTokenNumsGM, GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData);36 AICORE inline void Init(GM_ADDR expertTokenNumsGM, GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData);
40- AICORE inline void Process();37+ AICORE inline void ProcessFixed(uint32_t groupLocalId, uint32_t groupSize);
41 38 
42private:39private:
43- AICORE inline __gm__ MegaMoeSwigluSegmentRuntimeMeta* SegmentMetaPtr() const
44- {
45- return reinterpret_cast<__gm__ MegaMoeSwigluSegmentRuntimeMeta*>(
46- workspaceGM_ + tilingData_->swigluTiling.swigluSegmentMetaOffset);
47- }
48- AICORE inline void WriteSharedSegmentMetadata(uint32_t segmentIdx) const;
49- AICORE inline void ReadSharedSegmentMetadata(
50- uint32_t segmentIdx, uint32_t& segmentStartExpert, uint32_t& segmentEndExpert, uint32_t& segmentRowBase,
51- uint32_t& segmentRows, uint32_t& cumsumRows, uint32_t& expertTokenRows, uint32_t& rowSplitBase,
52- uint32_t& rowSplitRem) const;
53 AICORE inline uint64_t AlignUbBytes(uint64_t value) const { return (value + 31U) / 32U * 32U; }40 AICORE inline uint64_t AlignUbBytes(uint64_t value) const { return (value + 31U) / 32U * 32U; }
54 AICORE inline uint64_t SwigluMaxScratchBytes() const41 AICORE inline uint64_t SwigluMaxScratchBytes() const
55 {42 {
@@ -118,6 +105,7 @@ private:
118 AICORE inline void StoreFullRowOutput(uint32_t rowIdx, uint32_t bufferId) const;105 AICORE inline void StoreFullRowOutput(uint32_t rowIdx, uint32_t bufferId) const;
119 AICORE inline float ComputeAndStorePreparedFullRow(uint32_t rowIdx, uint32_t bufferId, float perTokenScale) const;106 AICORE inline float ComputeAndStorePreparedFullRow(uint32_t rowIdx, uint32_t bufferId, float perTokenScale) const;
120 AICORE inline void IssueStoreScale2Chunk(uint32_t rowStart, uint32_t rowCount, uint32_t scaleBufferId) const;107 AICORE inline void IssueStoreScale2Chunk(uint32_t rowStart, uint32_t rowCount, uint32_t scaleBufferId) const;
108+ AICORE inline void ProcessImpl();
121 109 
122 GM_ADDR workspaceGM_ = nullptr;110 GM_ADDR workspaceGM_ = nullptr;
123 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;111 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;
@@ -126,7 +114,6 @@ private:
126 __gm__ int8_t* gmPermutedTokenPtr_ = nullptr;114 __gm__ int8_t* gmPermutedTokenPtr_ = nullptr;
127 __gm__ float* perTokenScale2Ptr_ = nullptr;115 __gm__ float* perTokenScale2Ptr_ = nullptr;
128 __gm__ int32_t* cumsumMMPtr_ = nullptr;116 __gm__ int32_t* cumsumMMPtr_ = nullptr;
129- __gm__ int32_t* expertTokenNumsPtr_ = nullptr;
130 uint32_t problemN_ = 0;117 uint32_t problemN_ = 0;
131 uint32_t outputN_ = 0;118 uint32_t outputN_ = 0;
132 uint32_t maxOutputSize_ = 0;119 uint32_t maxOutputSize_ = 0;
@@ -134,7 +121,6 @@ private:
134 uint32_t rankSize_ = 0;121 uint32_t rankSize_ = 0;
135 uint32_t coreIdx_ = 0;122 uint32_t coreIdx_ = 0;
136 uint32_t coreNum_ = 1;123 uint32_t coreNum_ = 1;
137- uint32_t stageNum_ = 0;
138};124};
139 125 
140template <typename InputElement>126template <typename InputElement>
@@ -142,6 +128,7 @@ AICORE inline void Swiglu<InputElement>::Init(
142 GM_ADDR expertTokenNumsGM, GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData)128 GM_ADDR expertTokenNumsGM, GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData)
143{129{
144 (void)sizeof(InputElement);130 (void)sizeof(InputElement);
131+ (void)expertTokenNumsGM;
145 workspaceGM_ = workspaceGM;132 workspaceGM_ = workspaceGM;
146 tilingData_ = tilingData;133 tilingData_ = tilingData;
147 problemN_ = tilingData_->megaMoeInfo.N;134 problemN_ = tilingData_->megaMoeInfo.N;
@@ -149,7 +136,6 @@ AICORE inline void Swiglu<InputElement>::Init(
149 maxOutputSize_ = tilingData_->megaMoeInfo.maxOutputSize;136 maxOutputSize_ = tilingData_->megaMoeInfo.maxOutputSize;
150 expertPerRank_ = tilingData_->megaMoeInfo.expertPerRank;137 expertPerRank_ = tilingData_->megaMoeInfo.expertPerRank;
151 rankSize_ = tilingData_->runtimeInfo.rankSize;138 rankSize_ = tilingData_->runtimeInfo.rankSize;
152- stageNum_ = tilingData_->frontReorderTiling.stageNum;
153 139 
154 coreIdx_ = get_block_idx();140 coreIdx_ = get_block_idx();
155 coreNum_ = get_block_num();141 coreNum_ = get_block_num();
@@ -164,66 +150,6 @@ AICORE inline void Swiglu<InputElement>::Init(
164 reinterpret_cast<__gm__ int8_t*>(workspaceGM_ + tilingData_->swigluTiling.gmPermutedTokenOffset);150 reinterpret_cast<__gm__ int8_t*>(workspaceGM_ + tilingData_->swigluTiling.gmPermutedTokenOffset);
165 perTokenScale2Ptr_ = reinterpret_cast<__gm__ float*>(workspaceGM_ + tilingData_->swigluTiling.perTokenScale2Offset);151 perTokenScale2Ptr_ = reinterpret_cast<__gm__ float*>(workspaceGM_ + tilingData_->swigluTiling.perTokenScale2Offset);
166 cumsumMMPtr_ = reinterpret_cast<__gm__ int32_t*>(workspaceGM_ + tilingData_->frontReorderTiling.cumsumMMOffset);152 cumsumMMPtr_ = reinterpret_cast<__gm__ int32_t*>(workspaceGM_ + tilingData_->frontReorderTiling.cumsumMMOffset);
167- expertTokenNumsPtr_ = reinterpret_cast<__gm__ int32_t*>(expertTokenNumsGM);
168-}
169-template <typename InputElement>
170-AICORE inline void Swiglu<InputElement>::WriteSharedSegmentMetadata(uint32_t segmentIdx) const
171-{
172- if (coreIdx_ != 0U) {
173- return;
174- }
175- 
176- uint32_t segmentStartExpert = 0;
177- uint32_t segmentEndExpert = 0;
178- uint32_t segmentRowBase = 0;
179- uint32_t segmentRows = 0;
180- uint32_t cumsumRows = 0;
181- uint32_t expertTokenRows = 0;
182- MoeBuildSegmentMetadata(
183- segmentIdx, expertPerRank_, maxOutputSize_, cumsumMMPtr_, expertTokenNumsPtr_, rankSize_, segmentStartExpert,
184- segmentEndExpert, segmentRowBase, segmentRows, cumsumRows, expertTokenRows);
185- const uint32_t rowSplitBase = segmentRows / coreNum_;
186- const uint32_t rowSplitRem = segmentRows - rowSplitBase * coreNum_;
187- 
188- volatile __gm__ MegaMoeSwigluSegmentRuntimeMeta* entry = SegmentMetaPtr() + segmentIdx;
189- entry->valid = 0U;
190- entry->segmentIdx = segmentIdx;
191- entry->segmentStartExpert = segmentStartExpert;
192- entry->segmentEndExpert = segmentEndExpert;
193- entry->segmentRowBase = segmentRowBase;
194- entry->segmentRows = segmentRows;
195- entry->cumsumRows = cumsumRows;
196- entry->expertTokenRows = expertTokenRows;
197- entry->rowSplitBase = rowSplitBase;
198- entry->rowSplitRem = rowSplitRem;
199- entry->generation = stageNum_;
200- entry->producerCoreIdx = coreIdx_;
201- entry->metadataMode = kSwigluMetadataModeSharedSegmentMeta;
202- entry->segmentNum = MoeSwigluSegmentNum(expertPerRank_);
203- entry->epilogueGranularity = MoeSwigluEpilogueGranularity(expertPerRank_);
204- entry->marker = 1U;
205- pipe_barrier(PIPE_ALL);
206- entry->valid = 1U;
207- pipe_barrier(PIPE_ALL);
208- V5DcciGmRange(
209- reinterpret_cast<__gm__ void*>(SegmentMetaPtr() + segmentIdx), sizeof(MegaMoeSwigluSegmentRuntimeMeta));
210-}
211- 
212-template <typename InputElement>
213-AICORE inline void Swiglu<InputElement>::ReadSharedSegmentMetadata(
214- uint32_t segmentIdx, uint32_t& segmentStartExpert, uint32_t& segmentEndExpert, uint32_t& segmentRowBase,
215- uint32_t& segmentRows, uint32_t& cumsumRows, uint32_t& expertTokenRows, uint32_t& rowSplitBase,
216- uint32_t& rowSplitRem) const
217-{
218- volatile __gm__ MegaMoeSwigluSegmentRuntimeMeta* entry = SegmentMetaPtr() + segmentIdx;
219- segmentStartExpert = entry->segmentStartExpert;
220- segmentEndExpert = entry->segmentEndExpert;
221- segmentRowBase = entry->segmentRowBase;
222- segmentRows = entry->segmentRows;
223- cumsumRows = entry->cumsumRows;
224- expertTokenRows = entry->expertTokenRows;
225- rowSplitBase = entry->rowSplitBase;
226- rowSplitRem = entry->rowSplitRem;
227}153}
228template <typename InputElement>154template <typename InputElement>
229AICORE inline uint64_t Swiglu<InputElement>::SwigluStageBytes() const155AICORE inline uint64_t Swiglu<InputElement>::SwigluStageBytes() const
@@ -618,38 +544,44 @@ AICORE inline float Swiglu<InputElement>::ComputeAndStorePreparedFullRow(
618}544}
619 545 
620template <typename InputElement>546template <typename InputElement>
621-AICORE inline void Swiglu<InputElement>::Process()547+AICORE inline void Swiglu<InputElement>::ProcessFixed(uint32_t groupLocalId, uint32_t groupSize)
548+{
549+ coreIdx_ = groupLocalId;
550+ coreNum_ = groupSize;
551+ ProcessImpl();
552+}
553+ 
554+template <typename InputElement>
555+AICORE inline void Swiglu<InputElement>::ProcessImpl()
622{556{
623 if ASCEND_IS_AIC {557 if ASCEND_IS_AIC {
624 return;558 return;
625 }559 }
626 560 
627- const uint32_t segmentNum = MoeSwigluSegmentNum(expertPerRank_);561+ uint32_t groupBase = 0U;
628- for (uint32_t segmentIdx = 0; segmentIdx < segmentNum; ++segmentIdx) {562+ for (uint32_t groupIdx = 0U; groupIdx < expertPerRank_; ++groupIdx) {
629- CrossCoreWaitFlag<0x2>(MegaMoeC2VHardFlagId(segmentIdx));563+ const int32_t readyEpoch = static_cast<int32_t>(groupIdx * 2U + 2U);
630- WriteSharedSegmentMetadata(segmentIdx); // core 0负责分配任务给多个aiv564+ const MegaMoeSyncLayout sync = FixedSyncLayout(tilingData_);
631- pto::SYNCALL<pto::SyncCoreType::AIVOnly>();565+ if (groupIdx >= tilingData_->fixedGroupTiling.fullAicGmm1ExpertCount && coreIdx_ == 0U) {
566+ CoordinateGroupConsumersMte(
567+ workspaceGM_, tilingData_, sync.gmm1ArrivalBase, sync.swigluReadyBase,
568+ tilingData_->fixedGroupTiling.gmm1GroupSize, coreNum_, groupIdx);
569+ } else {
570+ WaitEpochAcquire(FixedSyncSlot(workspaceGM_, tilingData_, sync.swigluReadyBase + coreIdx_), readyEpoch);
571+ }
632 572 
633- uint32_t segmentStartExpert = 0;573+ const uint32_t currentMRaw = MoeCurrentMRaw(cumsumMMPtr_, rankSize_, expertPerRank_, groupIdx);
634- uint32_t segmentEndExpert = 0;574+ const uint32_t currentM = MoeClipCurrentM(currentMRaw, groupBase, maxOutputSize_);
635- uint32_t segmentRowBase = 0;575+ const uint32_t rowSplitBase = coreNum_ == 0U ? 0U : currentM / coreNum_;
636- uint32_t segmentRows = 0;576+ const uint32_t rowSplitRem = currentM - rowSplitBase * coreNum_;
637- uint32_t cumsumRows = 0;577+ const uint32_t localRows = rowSplitBase + (coreIdx_ < rowSplitRem ? 1U : 0U);
638- uint32_t expertTokenRows = 0;
639- uint32_t localRowStart = 0;
640- uint32_t localRows = 0;
641- uint32_t rowSplitBase = 0;
642- uint32_t rowSplitRem = 0;
643- ReadSharedSegmentMetadata(
644- segmentIdx, segmentStartExpert, segmentEndExpert, segmentRowBase, segmentRows, cumsumRows, expertTokenRows,
645- rowSplitBase, rowSplitRem);
646- localRows = rowSplitBase + (coreIdx_ < rowSplitRem ? 1U : 0U);
647 const uint32_t prefixRows = coreIdx_ * rowSplitBase + (coreIdx_ < rowSplitRem ? coreIdx_ : rowSplitRem);578 const uint32_t prefixRows = coreIdx_ * rowSplitBase + (coreIdx_ < rowSplitRem ? coreIdx_ : rowSplitRem);
648- localRowStart = segmentRowBase + prefixRows;579+ RunFullRowEpilogue(groupBase + prefixRows, localRows);
649- RunFullRowEpilogue(localRowStart, localRows);
650 580 
651- pto::SYNCALL<pto::SyncCoreType::AIVOnly>();581+ NotifyGroupConsumersMte(
652- CrossCoreSetFlag<0x2, PIPE_MTE3>(MegaMoeV2CHardFlagId(segmentIdx));582+ workspaceGM_, tilingData_, sync.swigluArrivalBase, sync.gmm2ReadyBase, coreNum_,
583+ tilingData_->fixedGroupTiling.gmm2GroupSize, coreIdx_, 0U, groupIdx);
584+ groupBase += currentM;
653 }585 }
654}586}
655 587 
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/unpermute.h+178-20
@@ -20,6 +20,7 @@ See LICENSE in the root of the software repository for the full text of the Lice
20#include "utils/common_helpers.hpp"20#include "utils/common_helpers.hpp"
21#include "utils/const_args.hpp"21#include "utils/const_args.hpp"
22#include "utils/hccl_window.hpp"22#include "utils/hccl_window.hpp"
23+#include "utils/mega_expert_sync.hpp"
23#include "utils/pto_vector.hpp"24#include "utils/pto_vector.hpp"
24 25 
25constexpr uint32_t kUnpermuteVecTileElems = 2048U;26constexpr uint32_t kUnpermuteVecTileElems = 2048U;
@@ -27,13 +28,13 @@ constexpr uint32_t kUnpermuteMetadataBufferNum = 2U;
27constexpr uint32_t kUnpermuteTokenBufferNum = 2U;28constexpr uint32_t kUnpermuteTokenBufferNum = 2U;
28constexpr uint32_t kUnpermuteTaskSplitOutputToken = 1U;29constexpr uint32_t kUnpermuteTaskSplitOutputToken = 1U;
29constexpr uint32_t kUnpermuteKTileMode = 1U;30constexpr uint32_t kUnpermuteKTileMode = 1U;
30-constexpr int32_t kUnpermuteInvalidRow = -1;
31 31 
32template <typename OutputElement>32template <typename OutputElement>
33class Unpermute {33class Unpermute {
34public:34public:
35 AICORE inline void Init(35 AICORE inline void Init(
36- GM_ADDR workspaceGM, GM_ADDR probsGM, GM_ADDR outGM, const __gm__ MegaMoeTilingData* tilingData);36+ GM_ADDR workspaceGM, GM_ADDR expertIdGM, GM_ADDR probsGM, GM_ADDR outGM,
37+ const __gm__ MegaMoeTilingData* tilingData, uint32_t workerIdx, uint32_t workerCount);
37 AICORE inline void Process();38 AICORE inline void Process();
38 39 
39private:40private:
@@ -61,6 +62,10 @@ private:
61 uint32_t batch = tilingData_->unpermuteTiling.unpermuteTokenBatch;62 uint32_t batch = tilingData_->unpermuteTiling.unpermuteTokenBatch;
62 return batch == 0U ? 1U : batch;63 return batch == 0U ? 1U : batch;
63 }64 }
65+ AICORE inline bool RankStreamingEnabled() const
66+ {
67+ return tilingData_->unpermuteTiling.unpermuteImplMode == kMegaMoeUnpermuteImplRankStreaming;
68+ }
64 AICORE inline event_t LoadFreeEvent(uint32_t bufferId) const { return static_cast<event_t>(bufferId); }69 AICORE inline event_t LoadFreeEvent(uint32_t bufferId) const { return static_cast<event_t>(bufferId); }
65 AICORE inline event_t LoadReadyEvent(uint32_t bufferId) const { return static_cast<event_t>(bufferId + 2U); }70 AICORE inline event_t LoadReadyEvent(uint32_t bufferId) const { return static_cast<event_t>(bufferId + 2U); }
66 AICORE inline event_t StoreFreeEvent() const { return EVENT_ID4; }71 AICORE inline event_t StoreFreeEvent() const { return EVENT_ID4; }
@@ -80,12 +85,20 @@ private:
80 AICORE inline void AccumulateChunk(uint32_t bufferId, float prob, uint32_t cols);85 AICORE inline void AccumulateChunk(uint32_t bufferId, float prob, uint32_t cols);
81 AICORE inline void StoreOutputChunk(uint32_t token, uint32_t col, uint32_t cols);86 AICORE inline void StoreOutputChunk(uint32_t token, uint32_t col, uint32_t cols);
82 AICORE inline void ProcessToken(uint32_t metaBufferId, uint32_t batchStart, uint32_t localToken);87 AICORE inline void ProcessToken(uint32_t metaBufferId, uint32_t batchStart, uint32_t localToken);
88+ AICORE inline bool TokenReadyForExpertProgress(
89+ uint32_t metaBufferId, uint32_t batchStart, uint32_t localToken, const uint32_t* readyExpertCounts) const;
90+ AICORE inline void BuildTokenRange(
91+ uint32_t workerIdx, uint32_t workerCount, uint32_t& tokenStart, uint32_t& tokenCount) const;
92+ AICORE inline void ProcessRankStreamingRange(
93+ uint32_t tokenStart, uint32_t tokenCount, const uint32_t* phase1ReadyExpertCounts, bool processPhase1);
94+ AICORE inline void ProcessRankStreaming();
83 95 
84 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;96 const __gm__ MegaMoeTilingData* tilingData_ = nullptr;
85 97 
86 PtoRemoteWindow remoteWindow_;98 PtoRemoteWindow remoteWindow_;
87 MegaMoePeerMemoryLayout peerMemoryLayout_;99 MegaMoePeerMemoryLayout peerMemoryLayout_;
88 __gm__ OutputElement* offsetDPtr_ = nullptr;100 __gm__ OutputElement* offsetDPtr_ = nullptr;
101+ __gm__ int32_t* expertIdPtr_ = nullptr;
89 __gm__ int32_t* expandedRowIdxPtr_ = nullptr;102 __gm__ int32_t* expandedRowIdxPtr_ = nullptr;
90 __gm__ float* probsPtr_ = nullptr;103 __gm__ float* probsPtr_ = nullptr;
91 __gm__ OutputElement* outPtr_ = nullptr;104 __gm__ OutputElement* outPtr_ = nullptr;
@@ -94,6 +107,9 @@ private:
94 uint32_t problemK_ = 0;107 uint32_t problemK_ = 0;
95 uint32_t topK_ = 0;108 uint32_t topK_ = 0;
96 uint32_t maxOutputSize_ = 0;109 uint32_t maxOutputSize_ = 0;
110+ uint32_t expertPerRank_ = 0;
111+ uint32_t rankSize_ = 0;
112+ int32_t dataReadyEpoch_ = 0;
97 uint32_t expandedRowsValid_ = 0;113 uint32_t expandedRowsValid_ = 0;
98 uint32_t coreIdx_ = 0;114 uint32_t coreIdx_ = 0;
99 uint32_t coreNum_ = 1;115 uint32_t coreNum_ = 1;
@@ -113,7 +129,8 @@ private:
113 129 
114template <typename OutputElement>130template <typename OutputElement>
115AICORE inline void Unpermute<OutputElement>::Init(131AICORE inline void Unpermute<OutputElement>::Init(
116- GM_ADDR workspaceGM, GM_ADDR probsGM, GM_ADDR outGM, const __gm__ MegaMoeTilingData* tilingData)132+ GM_ADDR workspaceGM, GM_ADDR expertIdGM, GM_ADDR probsGM, GM_ADDR outGM, const __gm__ MegaMoeTilingData* tilingData,
133+ uint32_t workerIdx, uint32_t workerCount)
117{134{
118 tilingData_ = tilingData;135 tilingData_ = tilingData;
119 136 
@@ -121,14 +138,12 @@ AICORE inline void Unpermute<OutputElement>::Init(
121 problemK_ = tilingData_->megaMoeInfo.K;138 problemK_ = tilingData_->megaMoeInfo.K;
122 topK_ = tilingData_->megaMoeInfo.topK;139 topK_ = tilingData_->megaMoeInfo.topK;
123 maxOutputSize_ = tilingData_->megaMoeInfo.maxOutputSize;140 maxOutputSize_ = tilingData_->megaMoeInfo.maxOutputSize;
141+ expertPerRank_ = tilingData_->megaMoeInfo.expertPerRank;
142+ rankSize_ = tilingData_->runtimeInfo.rankSize;
124 const uint32_t expandedRows = problemM_ * topK_;143 const uint32_t expandedRows = problemM_ * topK_;
125 expandedRowsValid_ = expandedRows < maxOutputSize_ ? expandedRows : maxOutputSize_;144 expandedRowsValid_ = expandedRows < maxOutputSize_ ? expandedRows : maxOutputSize_;
126- coreIdx_ = get_block_idx();145+ coreIdx_ = workerIdx;
127- coreNum_ = get_block_num();146+ coreNum_ = workerCount;
128- if ASCEND_IS_AIV {
129- coreIdx_ = get_block_idx() + get_subblockid() * get_block_num();
130- coreNum_ = get_block_num() * get_subblockdim();
131- }
132 147 
133 splitBase_ = coreNum_ == 0U ? 0U : problemM_ / coreNum_;148 splitBase_ = coreNum_ == 0U ? 0U : problemM_ / coreNum_;
134 splitRem_ = coreNum_ == 0U ? 0U : problemM_ % coreNum_;149 splitRem_ = coreNum_ == 0U ? 0U : problemM_ % coreNum_;
@@ -138,10 +153,12 @@ AICORE inline void Unpermute<OutputElement>::Init(
138 remoteWindow_.Init(reinterpret_cast<GM_ADDR>(tilingData_->runtimeInfo.remoteWindowContext));153 remoteWindow_.Init(reinterpret_cast<GM_ADDR>(tilingData_->runtimeInfo.remoteWindowContext));
139 peerMemoryLayout_.Init(remoteWindow_);154 peerMemoryLayout_.Init(remoteWindow_);
140 offsetDPtr_ = reinterpret_cast<__gm__ OutputElement*>(remoteWindow_.LocalBase() + peerMemoryLayout_.offsetD);155 offsetDPtr_ = reinterpret_cast<__gm__ OutputElement*>(remoteWindow_.LocalBase() + peerMemoryLayout_.offsetD);
156+ expertIdPtr_ = reinterpret_cast<__gm__ int32_t*>(expertIdGM);
141 expandedRowIdxPtr_ =157 expandedRowIdxPtr_ =
142 reinterpret_cast<__gm__ int32_t*>(workspaceGM + tilingData_->frontReorderTiling.expandedRowIdxOffset);158 reinterpret_cast<__gm__ int32_t*>(workspaceGM + tilingData_->frontReorderTiling.expandedRowIdxOffset);
143 probsPtr_ = reinterpret_cast<__gm__ float*>(probsGM);159 probsPtr_ = reinterpret_cast<__gm__ float*>(probsGM);
144 outPtr_ = reinterpret_cast<__gm__ OutputElement*>(outGM);160 outPtr_ = reinterpret_cast<__gm__ OutputElement*>(outGM);
161+ dataReadyEpoch_ = RankStreamingEnabled() ? remoteWindow_.DataReadyEpoch() : 0;
145 162 
146 InitUbLayout();163 InitUbLayout();
147}164}
@@ -284,10 +301,6 @@ AICORE inline void Unpermute<OutputElement>::ProcessToken(
284 bool hasPending = false;301 bool hasPending = false;
285 uint32_t pendingBuffer = 0;302 uint32_t pendingBuffer = 0;
286 float pendingProb = 0.0f;303 float pendingProb = 0.0f;
287- uint32_t validTopk = 0;
288- uint32_t topkProcessed = 0;
289- int32_t firstExpandedRow = kUnpermuteInvalidRow;
290- float firstProb = 0.0f;
291 304 
292 for (uint32_t topkIdx = 0; topkIdx < topK_; ++topkIdx) {305 for (uint32_t topkIdx = 0; topkIdx < topK_; ++topkIdx) {
293 const int32_t expandedRow = ReadExpandedRow(metaBufferId, localToken, topkIdx);306 const int32_t expandedRow = ReadExpandedRow(metaBufferId, localToken, topkIdx);
@@ -302,24 +315,164 @@ AICORE inline void Unpermute<OutputElement>::ProcessToken(
302 LoadOffsetDChunk(bufferId, expandedRow, col, cols);315 LoadOffsetDChunk(bufferId, expandedRow, col, cols);
303 if (hasPending) {316 if (hasPending) {
304 AccumulateChunk(pendingBuffer, pendingProb, cols);317 AccumulateChunk(pendingBuffer, pendingProb, cols);
305- ++topkProcessed;
306 }318 }
307 pendingBuffer = bufferId;319 pendingBuffer = bufferId;
308 pendingProb = prob;320 pendingProb = prob;
309 hasPending = true;321 hasPending = true;
310- ++validTopk;
311- if (firstExpandedRow == kUnpermuteInvalidRow) {
312- firstExpandedRow = expandedRow;
313- firstProb = prob;
314- }
315 }322 }
316 if (hasPending) {323 if (hasPending) {
317 AccumulateChunk(pendingBuffer, pendingProb, cols);324 AccumulateChunk(pendingBuffer, pendingProb, cols);
318- ++topkProcessed;
319 }325 }
320 StoreOutputChunk(token, col, cols); // 写回GM326 StoreOutputChunk(token, col, cols); // 写回GM
321 }327 }
322}328}
329+ 
330+template <typename OutputElement>
331+AICORE inline bool Unpermute<OutputElement>::TokenReadyForExpertProgress(
332+ uint32_t metaBufferId, uint32_t batchStart, uint32_t localToken, const uint32_t* readyExpertCounts) const
333+{
334+ if (readyExpertCounts == nullptr || expertPerRank_ == 0U || rankSize_ == 0U ||
335+ rankSize_ > COMBINE_EXPERT_PROGRESS_MAX_RANKS) {
336+ return false;
337+ }
338+ bool hasValidRoute = false;
339+ bool allRoutesReady = true;
340+ for (uint32_t topkIdx = 0U; topkIdx < topK_; ++topkIdx) {
341+ const int32_t expandedRow = ReadExpandedRow(metaBufferId, localToken, topkIdx);
342+ if (expandedRow < 0 || static_cast<uint32_t>(expandedRow) >= expandedRowsValid_) {
343+ continue;
344+ }
345+ hasValidRoute = true;
346+ const int32_t expert = expertIdPtr_[static_cast<uint64_t>(batchStart + localToken) * topK_ + topkIdx];
347+ if (expert < 0) {
348+ allRoutesReady = false;
349+ continue;
350+ }
351+ const uint32_t globalExpert = static_cast<uint32_t>(expert);
352+ const uint32_t producerRank = globalExpert / expertPerRank_;
353+ const uint32_t localExpert = globalExpert - producerRank * expertPerRank_;
354+ if (producerRank >= rankSize_ || localExpert >= readyExpertCounts[producerRank]) {
355+ allRoutesReady = false;
356+ }
357+ }
358+ return hasValidRoute && allRoutesReady;
359+}
360+ 
361+template <typename OutputElement>
362+AICORE inline void Unpermute<OutputElement>::BuildTokenRange(
363+ uint32_t workerIdx, uint32_t workerCount, uint32_t& tokenStart, uint32_t& tokenCount) const
364+{
365+ tokenStart = 0U;
366+ tokenCount = 0U;
367+ if (workerCount == 0U || workerIdx >= workerCount) {
368+ return;
369+ }
370+ const uint32_t splitBase = problemM_ / workerCount;
371+ const uint32_t splitRem = problemM_ - splitBase * workerCount;
372+ tokenStart = workerIdx * splitBase + (workerIdx < splitRem ? workerIdx : splitRem);
373+ tokenCount = splitBase + (workerIdx < splitRem ? 1U : 0U);
374+}
375+ 
376+template <typename OutputElement>
377+AICORE inline void Unpermute<OutputElement>::ProcessRankStreamingRange(
378+ uint32_t tokenStart, uint32_t tokenCount, const uint32_t* phase1ReadyExpertCounts, bool processPhase1)
379+{
380+ if (tokenCount == 0U) {
381+ return;
382+ }
383+ 
384+ const uint32_t batchLimit = TokenBatch();
385+ uint32_t currentBatchStart = tokenStart;
386+ uint32_t currentBatchTokens = tokenCount < batchLimit ? tokenCount : batchLimit;
387+ uint32_t currentBuffer = 0U;
388+ PrefetchMetadata(currentBuffer, currentBatchStart, currentBatchTokens);
389+ 
390+ uint32_t consumedTokens = 0U;
391+ while (consumedTokens < tokenCount) {
392+ WaitMetadata(currentBuffer);
393+ const uint32_t nextConsumed = consumedTokens + currentBatchTokens;
394+ const bool hasNext = nextConsumed < tokenCount;
395+ const uint32_t nextBuffer = (currentBuffer + 1U) % kUnpermuteMetadataBufferNum;
396+ uint32_t nextBatchStart = 0U;
397+ uint32_t nextBatchTokens = 0U;
398+ if (hasNext) {
399+ nextBatchStart = tokenStart + nextConsumed;
400+ const uint32_t remaining = tokenCount - nextConsumed;
401+ nextBatchTokens = remaining < batchLimit ? remaining : batchLimit;
402+ PrefetchMetadata(nextBuffer, nextBatchStart, nextBatchTokens);
403+ }
404+ 
405+ for (uint32_t localToken = 0U; localToken < currentBatchTokens; ++localToken) {
406+ const bool phase1Task =
407+ TokenReadyForExpertProgress(currentBuffer, currentBatchStart, localToken, phase1ReadyExpertCounts);
408+ if (phase1Task == processPhase1) {
409+ ProcessToken(currentBuffer, currentBatchStart, localToken);
410+ }
411+ }
412+ 
413+ consumedTokens = nextConsumed;
414+ currentBatchStart = nextBatchStart;
415+ currentBatchTokens = nextBatchTokens;
416+ currentBuffer = nextBuffer;
417+ }
418+}
419+ 
420+template <typename OutputElement>
421+AICORE inline void Unpermute<OutputElement>::ProcessRankStreaming()
422+{
423+ const uint32_t rankCount = rankSize_;
424+ const uint32_t initialWorkerCount = coreNum_ < tilingData_->fixedGroupTiling.gmm1GroupSize * 2U ?
425+ coreNum_ :
426+ tilingData_->fixedGroupTiling.gmm1GroupSize * 2U;
427+ 
428+ WaitEpochAcquire(remoteWindow_.LocalUnpermutePhase1ProgressEpochSlot(), dataReadyEpoch_);
429+ uint32_t phase1ReadyExpertCounts[COMBINE_EXPERT_PROGRESS_MAX_RANKS] = {0U};
430+ remoteWindow_.ReadUnpermutePhase1Progress(phase1ReadyExpertCounts, rankCount);
431+ bool phase1AllReady = true;
432+ for (uint32_t producerRank = 0U; producerRank < rankCount; ++producerRank) {
433+ if (phase1ReadyExpertCounts[producerRank] < expertPerRank_) {
434+ phase1AllReady = false;
435+ }
436+ }
437+ 
438+ if (coreIdx_ < initialWorkerCount) {
439+ uint32_t phase1TokenStart = 0U;
440+ uint32_t phase1TokenCount = 0U;
441+ BuildTokenRange(coreIdx_, initialWorkerCount, phase1TokenStart, phase1TokenCount);
442+ ProcessRankStreamingRange(phase1TokenStart, phase1TokenCount, phase1ReadyExpertCounts, true);
443+ remoteWindow_.PublishPhase1Done(coreIdx_, dataReadyEpoch_);
444+ }
445+ 
446+ if (coreIdx_ == 0U) {
447+ remoteWindow_.WaitPhase1DoneMte(initialWorkerCount, dataReadyEpoch_);
448+ uint32_t liveReadyExpertCounts[COMBINE_EXPERT_PROGRESS_MAX_RANKS] = {0U};
449+ uint32_t minimumReady = 0U;
450+ while (minimumReady < expertPerRank_) {
451+ const uint32_t observedMinimum =
452+ remoteWindow_.ReadExpertProgressMte(dataReadyEpoch_, expertPerRank_, liveReadyExpertCounts);
453+ if (observedMinimum > minimumReady) {
454+ remoteWindow_.AcquireDataReady();
455+ minimumReady = observedMinimum;
456+ } else {
457+ EpochPollBackoff();
458+ }
459+ }
460+ if (coreNum_ > initialWorkerCount) {
461+ WaitEpochAcquire(remoteWindow_.LocalUnpermuteStartSlot(initialWorkerCount), dataReadyEpoch_);
462+ }
463+ remoteWindow_.PublishUnpermuteAllReady(coreNum_, dataReadyEpoch_);
464+ } else {
465+ WaitEpochAcquire(remoteWindow_.LocalUnpermuteAllReadySlot(coreIdx_), dataReadyEpoch_);
466+ }
467+ 
468+ if (!phase1AllReady) {
469+ uint32_t phase2TokenStart = 0U;
470+ uint32_t phase2TokenCount = 0U;
471+ BuildTokenRange(coreIdx_, coreNum_, phase2TokenStart, phase2TokenCount);
472+ ProcessRankStreamingRange(phase2TokenStart, phase2TokenCount, phase1ReadyExpertCounts, false);
473+ }
474+}
475+ 
323template <typename OutputElement>476template <typename OutputElement>
324AICORE inline void Unpermute<OutputElement>::Process()477AICORE inline void Unpermute<OutputElement>::Process()
325{478{
@@ -327,6 +480,11 @@ AICORE inline void Unpermute<OutputElement>::Process()
327 return;480 return;
328 }481 }
329 SetInitialFlags();482 SetInitialFlags();
483+ if (RankStreamingEnabled()) {
484+ ProcessRankStreaming();
485+ FinalizeLocalPipe();
486+ return;
487+ }
330 if (tokenCount_ == 0U) {488 if (tokenCount_ == 0U) {
331 FinalizeLocalPipe();489 FinalizeLocalPipe();
332 return;490 return;
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/utils/common_helpers.hpp+1-18
@@ -62,7 +62,7 @@ AICORE inline void V5DcciGmRangeNoFence(__gm__ void* ptr, uint64_t bytes)
62 const uint64_t end = (reinterpret_cast<uint64_t>(ptr) + bytes + cacheLineBytes - 1U) & ~(cacheLineBytes - 1U);62 const uint64_t end = (reinterpret_cast<uint64_t>(ptr) + bytes + cacheLineBytes - 1U) & ~(cacheLineBytes - 1U);
63 for (uint64_t addr = start; addr < end; addr += cacheLineBytes) {63 for (uint64_t addr = start; addr < end; addr += cacheLineBytes) {
64 __asm__ __volatile__("");64 __asm__ __volatile__("");
65- dcci(reinterpret_cast<__gm__ void*>(addr), cache_line_t::SINGLE_CACHE_LINE);65+ dcci(reinterpret_cast<__gm__ void*>(addr), SINGLE_CACHE_LINE);
66 __asm__ __volatile__("");66 __asm__ __volatile__("");
67 }67 }
68}68}
@@ -78,21 +78,4 @@ AICORE inline uint32_t MegaMoeActiveCopyCores(uint32_t rankSize, uint32_t coreNu
78 return rankSize < coreNum ? rankSize : coreNum;78 return rankSize < coreNum ? rankSize : coreNum;
79}79}
80 80 
81-AICORE inline uint16_t MegaMoeD2CHardFlagId(uint32_t logicalGroupEventIdx)
82-{
83- return static_cast<uint16_t>(MEGA_MOE_D2C_HARD_FLAG_BASE + logicalGroupEventIdx / CROSS_CORE_FLAG_MAX_SET_COUNT);
84-}
85- 
86-AICORE inline uint16_t MegaMoeC2VHardFlagId(uint32_t segmentIdx)
87-{
88- (void)segmentIdx;
89- return MEGA_MOE_C2V_HARD_FLAG_BASE;
90-}
91- 
92-AICORE inline uint16_t MegaMoeV2CHardFlagId(uint32_t segmentIdx)
93-{
94- (void)segmentIdx;
95- return MEGA_MOE_V2C_HARD_FLAG_BASE;
96-}
97- 
98#endif // COMMON_HELPERS_HPP81#endif // COMMON_HELPERS_HPP
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/utils/const_args.hpp+0-15
@@ -15,19 +15,6 @@ See LICENSE in the root of the software repository for the full text of the Lice
15 15 
16constexpr static uint64_t MB_SIZE = 1024 * 1024UL;16constexpr static uint64_t MB_SIZE = 1024 * 1024UL;
17constexpr static int32_t UB_ALIGN = 32;17constexpr static int32_t UB_ALIGN = 32;
18-constexpr uint16_t CROSS_CORE_FLAG_MAX_SET_COUNT = 15;
19-constexpr uint16_t MEGA_MOE_MAX_BUSINESS_HARD_FLAG_ID = 10;
20-constexpr uint16_t MEGA_MOE_D2C_HARD_FLAG_BASE = 0;
21-constexpr uint16_t MEGA_MOE_D2C_HARD_FLAG_COUNT = 9;
22-constexpr uint16_t MEGA_MOE_D2C_HARD_FLAG_LAST = MEGA_MOE_D2C_HARD_FLAG_BASE + MEGA_MOE_D2C_HARD_FLAG_COUNT - 1;
23-constexpr uint16_t MEGA_MOE_C2V_HARD_FLAG_BASE = 9;
24-constexpr uint16_t MEGA_MOE_V2C_HARD_FLAG_BASE = 10;
25-constexpr uint16_t MEGA_MOE_GMM2_TO_COMBINE_HARD_FLAG_BASE = 0;
26-constexpr uint32_t MEGA_MOE_D2C_MAX_LOGICAL_GROUP_EVENTS = MEGA_MOE_D2C_HARD_FLAG_COUNT * CROSS_CORE_FLAG_MAX_SET_COUNT;
27-constexpr uint32_t MEGA_MOE_GMM2_TO_COMBINE_MAX_LOGICAL_GROUP_EVENTS =
28- MEGA_MOE_D2C_HARD_FLAG_COUNT * CROSS_CORE_FLAG_MAX_SET_COUNT;
29-static_assert(MEGA_MOE_D2C_HARD_FLAG_LAST < MEGA_MOE_C2V_HARD_FLAG_BASE);
30-static_assert(MEGA_MOE_V2C_HARD_FLAG_BASE <= MEGA_MOE_MAX_BUSINESS_HARD_FLAG_ID);
31 18 
32struct AtlasA2 {19struct AtlasA2 {
33 static constexpr uint32_t BIAS_SIZE = 1024;20 static constexpr uint32_t BIAS_SIZE = 1024;
@@ -39,6 +26,4 @@ struct AtlasA2 {
39 static constexpr uint32_t L0C_SIZE = 128 * 1024;26 static constexpr uint32_t L0C_SIZE = 128 * 1024;
40};27};
41 28 
42-#include "moe_swiglu_segment.hpp"
43- 
44#endif // CONST_ARGS_HPP29#endif // CONST_ARGS_HPP
Mkernels/manual/a2a3/dispatch_mega_combine/op_kernel/utils/hccl_window.hpp+483-0
@@ -27,6 +27,88 @@ constexpr uint32_t START_AIV_BARRIER_COUNTER_BASE_INDEX = 14336;
27constexpr uint32_t START_AIV_BARRIER_EPOCH_INDEX = 18432;27constexpr uint32_t START_AIV_BARRIER_EPOCH_INDEX = 18432;
28constexpr uint32_t START_AIC_BARRIER_COUNTER_BASE_INDEX = 20480;28constexpr uint32_t START_AIC_BARRIER_COUNTER_BASE_INDEX = 20480;
29constexpr uint32_t START_AIC_BARRIER_EPOCH_INDEX = 24576;29constexpr uint32_t START_AIC_BARRIER_EPOCH_INDEX = 24576;
30+// Keep one cache line per producer slot. These slots are independent from the
31+// older all-rank barrier counters above and are reused with a launch epoch.
32+constexpr uint32_t COMBINE_DATA_READY_BASE_INDEX = 26624;
33+constexpr uint32_t COMBINE_DATA_READY_STRIDE = 16;
34+constexpr uint32_t COMBINE_EXPERT_PROGRESS_OFFSET = 1;
35+constexpr uint32_t COMBINE_EXPERT_PROGRESS_COUNT_BITS = 6;
36+constexpr uint32_t COMBINE_EXPERT_PROGRESS_COUNT_MASK = (1U << COMBINE_EXPERT_PROGRESS_COUNT_BITS) - 1U;
37+constexpr uint32_t COMBINE_EXPERT_PROGRESS_MAX_RANKS = kMegaMoeExpertProgressMaxRanks;
38+constexpr uint32_t COMBINE_DATA_READY_EPOCH_INDEX = 28672;
39+constexpr uint32_t UNPERMUTE_ALL_READY_BASE_INDEX = COMBINE_DATA_READY_EPOCH_INDEX + COMBINE_DATA_READY_STRIDE;
40+constexpr uint32_t UNPERMUTE_ALL_READY_STRIDE = 16;
41+constexpr uint32_t UNPERMUTE_ALL_READY_SLOT_COUNT = kMegaMoeFixedPhysicalAivNum;
42+constexpr uint32_t UNPERMUTE_START_BASE_INDEX =
43+ UNPERMUTE_ALL_READY_BASE_INDEX + UNPERMUTE_ALL_READY_SLOT_COUNT * UNPERMUTE_ALL_READY_STRIDE;
44+constexpr uint32_t UNPERMUTE_START_STRIDE = 16;
45+constexpr uint32_t UNPERMUTE_START_SLOT_COUNT = kMegaMoeFixedPhysicalAivNum;
46+constexpr uint32_t COMBINE_LOCAL_DONE_BASE_INDEX =
47+ UNPERMUTE_START_BASE_INDEX + UNPERMUTE_START_SLOT_COUNT * UNPERMUTE_START_STRIDE;
48+constexpr uint32_t COMBINE_LOCAL_DONE_STRIDE = 16;
49+constexpr uint32_t COMBINE_LOCAL_DONE_SLOT_COUNT = kMegaMoeFixedPhysicalAivNum;
50+constexpr uint32_t UNPERMUTE_PHASE1_PROGRESS_BASE_INDEX =
51+ COMBINE_LOCAL_DONE_BASE_INDEX + COMBINE_LOCAL_DONE_SLOT_COUNT * COMBINE_LOCAL_DONE_STRIDE;
52+constexpr uint32_t UNPERMUTE_PHASE1_PROGRESS_STRIDE = 32;
53+constexpr uint32_t UNPERMUTE_PHASE1_PROGRESS_EPOCH_OFFSET = 0;
54+constexpr uint32_t UNPERMUTE_PHASE1_PROGRESS_MASK_OFFSET = 1;
55+constexpr uint32_t UNPERMUTE_PHASE1_PROGRESS_COUNTS_OFFSET = 2;
56+constexpr uint32_t REMOTE_WINDOW_CACHE_LINE_VALUES = 64U / sizeof(int32_t);
57+constexpr uint32_t UNPERMUTE_PHASE1_PROGRESS_VALUE_COUNT =
58+ UNPERMUTE_PHASE1_PROGRESS_COUNTS_OFFSET + COMBINE_EXPERT_PROGRESS_MAX_RANKS;
59+constexpr uint32_t UNPERMUTE_PHASE1_PROGRESS_MAX_CACHE_LINE_COUNT =
60+ (UNPERMUTE_PHASE1_PROGRESS_VALUE_COUNT + REMOTE_WINDOW_CACHE_LINE_VALUES - 1U) / REMOTE_WINDOW_CACHE_LINE_VALUES;
61+constexpr uint32_t UNPERMUTE_DISPATCH_RELEASE_BASE_INDEX =
62+ UNPERMUTE_PHASE1_PROGRESS_BASE_INDEX + UNPERMUTE_PHASE1_PROGRESS_STRIDE;
63+constexpr uint32_t UNPERMUTE_DISPATCH_RELEASE_STRIDE = 16;
64+constexpr uint32_t UNPERMUTE_DISPATCH_RELEASE_SLOT_COUNT = kMegaMoeFixedPhysicalAicNum;
65+constexpr uint32_t UNPERMUTE_SWIGLU_RELEASE_BASE_INDEX =
66+ UNPERMUTE_DISPATCH_RELEASE_BASE_INDEX + UNPERMUTE_DISPATCH_RELEASE_SLOT_COUNT * UNPERMUTE_DISPATCH_RELEASE_STRIDE;
67+constexpr uint32_t UNPERMUTE_SWIGLU_RELEASE_STRIDE = 16;
68+constexpr uint32_t UNPERMUTE_SWIGLU_RELEASE_SLOT_COUNT = kMegaMoeFixedPhysicalAicNum;
69+constexpr uint32_t UNPERMUTE_PHASE1_DONE_BASE_INDEX =
70+ UNPERMUTE_SWIGLU_RELEASE_BASE_INDEX + UNPERMUTE_SWIGLU_RELEASE_SLOT_COUNT * UNPERMUTE_SWIGLU_RELEASE_STRIDE;
71+constexpr uint32_t UNPERMUTE_PHASE1_DONE_STRIDE = 16;
72+constexpr uint32_t UNPERMUTE_PHASE1_DONE_SLOT_COUNT = kMegaMoeFixedPhysicalAivNum;
73+constexpr uint32_t REMOTE_WINDOW_SYNC_MAX_SLOTS = kMegaMoeFixedPhysicalAivNum;
74+constexpr uint32_t REMOTE_WINDOW_SYNC_VALUES_PER_SLOT = 16;
75+constexpr uint32_t REMOTE_WINDOW_SYNC_MAX_VALUES = REMOTE_WINDOW_SYNC_MAX_SLOTS * REMOTE_WINDOW_SYNC_VALUES_PER_SLOT;
76+constexpr uint64_t REMOTE_WINDOW_SYNC_SNAPSHOT_BYTES =
77+ static_cast<uint64_t>(REMOTE_WINDOW_SYNC_MAX_VALUES) * sizeof(int32_t);
78+constexpr uint64_t REMOTE_WINDOW_SYNC_SNAPSHOT_UB_OFFSET = AtlasA2::UB_SIZE - REMOTE_WINDOW_SYNC_SNAPSHOT_BYTES;
79+constexpr event_t REMOTE_WINDOW_SYNC_SNAPSHOT_EVENT = EVENT_ID0;
80+static_assert(
81+ COMBINE_DATA_READY_BASE_INDEX + COMBINE_EXPERT_PROGRESS_MAX_RANKS * COMBINE_DATA_READY_STRIDE <=
82+ COMBINE_DATA_READY_EPOCH_INDEX);
83+static_assert(COMBINE_LOCAL_DONE_SLOT_COUNT >= COMBINE_EXPERT_PROGRESS_MAX_RANKS);
84+static_assert(UNPERMUTE_PHASE1_PROGRESS_VALUE_COUNT <= UNPERMUTE_PHASE1_PROGRESS_STRIDE);
85+static_assert(
86+ UNPERMUTE_PHASE1_PROGRESS_MAX_CACHE_LINE_COUNT * REMOTE_WINDOW_CACHE_LINE_VALUES <=
87+ UNPERMUTE_PHASE1_PROGRESS_STRIDE);
88+static_assert(UNPERMUTE_PHASE1_PROGRESS_BASE_INDEX % REMOTE_WINDOW_CACHE_LINE_VALUES == 0U);
89+static_assert(
90+ UNPERMUTE_PHASE1_DONE_BASE_INDEX + UNPERMUTE_PHASE1_DONE_SLOT_COUNT * UNPERMUTE_PHASE1_DONE_STRIDE <
91+ MB_SIZE / sizeof(int32_t));
92+static_assert(REMOTE_WINDOW_SYNC_SNAPSHOT_UB_OFFSET % UB_ALIGN == 0U);
93+ 
94+AICORE inline void RemoteWindowSyncPollBackoff()
95+{
96+ constexpr uint32_t kDelayTicks = 3U;
97+ const uint64_t deadline = get_sys_cnt() + kDelayTicks;
98+ while (get_sys_cnt() < deadline) {
99+ __asm__ __volatile__("");
100+ }
101+}
102+ 
103+AICORE inline void DcciUnpermutePhase1Progress(volatile __gm__ int32_t* base, uint32_t rankCount)
104+{
105+ const uint32_t valueCount = UNPERMUTE_PHASE1_PROGRESS_COUNTS_OFFSET + rankCount;
106+ const uint32_t cacheLineCount =
107+ (valueCount + REMOTE_WINDOW_CACHE_LINE_VALUES - 1U) / REMOTE_WINDOW_CACHE_LINE_VALUES;
108+ for (uint32_t line = 0U; line < cacheLineCount; ++line) {
109+ dcci((__gm__ void*)(base + line * REMOTE_WINDOW_CACHE_LINE_VALUES), SINGLE_CACHE_LINE);
110+ }
111+}
30 112 
31class PtoRemoteWindow {113class PtoRemoteWindow {
32public:114public:
@@ -74,6 +156,407 @@ public:
74 return RemotePtr(LocalSignalBase(), rankId);156 return RemotePtr(LocalSignalBase(), rankId);
75 }157 }
76 158 
159+ // The epoch is advanced by one AIV before the first stage starts. This
160+ // permits repeated launches even when the host does not clear the window.
161+ AICORE inline int32_t PrepareDataReadyEpoch() const
162+ {
163+ __gm__ int32_t* epochSlot = LocalSignalBase() + COMBINE_DATA_READY_EPOCH_INDEX;
164+ dcci((__gm__ void*)epochSlot, SINGLE_CACHE_LINE);
165+ __asm__ __volatile__("");
166+ int32_t epoch = *epochSlot + 1;
167+ if (epoch <= 0) {
168+ epoch = 1;
169+ }
170+ *epochSlot = epoch;
171+ dcci((__gm__ void*)epochSlot, SINGLE_CACHE_LINE);
172+ PublishUnpermuteAllReady(UNPERMUTE_ALL_READY_SLOT_COUNT, 0);
173+ dsb(DSB_DDR);
174+ return epoch;
175+ }
176+ 
177+ AICORE inline int32_t DataReadyEpoch() const
178+ {
179+ __gm__ int32_t* epochSlot = LocalSignalBase() + COMBINE_DATA_READY_EPOCH_INDEX;
180+ dcci((__gm__ void*)epochSlot, SINGLE_CACHE_LINE);
181+ __asm__ __volatile__("");
182+ return *epochSlot;
183+ }
184+ 
185+ AICORE inline volatile __gm__ int32_t* LocalDataReadySlot(int32_t producerRank) const
186+ {
187+ if (producerRank < 0 || producerRank >= rankSize_) {
188+ return nullptr;
189+ }
190+ return LocalSignalBase() + COMBINE_DATA_READY_BASE_INDEX +
191+ static_cast<uint32_t>(producerRank) * COMBINE_DATA_READY_STRIDE;
192+ }
193+ 
194+ AICORE inline void PublishDataReady(int32_t ownerRank, int32_t epoch) const
195+ {
196+ if (ownerRank < 0 || ownerRank >= rankSize_) {
197+ return;
198+ }
199+ if (ownerRank == rank_) {
200+ volatile __gm__ int32_t* localSlot = LocalDataReadySlot(rank_);
201+ *localSlot = epoch;
202+ dcci((__gm__ void*)localSlot, SINGLE_CACHE_LINE);
203+ __asm__ __volatile__("");
204+ return;
205+ }
206+ __gm__ int32_t* remoteSlot = RemoteSignalBase(ownerRank) + COMBINE_DATA_READY_BASE_INDEX +
207+ static_cast<uint32_t>(rank_) * COMBINE_DATA_READY_STRIDE;
208+ auto signal = pto::comm::Signal(remoteSlot);
209+ pto::comm::TNOTIFY(signal, epoch, pto::comm::NotifyOp::Set);
210+ }
211+ 
212+ AICORE inline int32_t EncodeExpertProgress(int32_t epoch, uint32_t readyExpertCount) const
213+ {
214+ const uint32_t count = readyExpertCount & COMBINE_EXPERT_PROGRESS_COUNT_MASK;
215+ return static_cast<int32_t>((static_cast<uint32_t>(epoch) << COMBINE_EXPERT_PROGRESS_COUNT_BITS) | count);
216+ }
217+ 
218+ AICORE inline void PublishExpertProgress(int32_t consumerRank, uint32_t readyExpertCount, int32_t epoch) const
219+ {
220+ if (consumerRank < 0 || consumerRank >= rankSize_ || readyExpertCount > COMBINE_EXPERT_PROGRESS_COUNT_MASK ||
221+ readyExpertCount > kMegaMoeFixedMaxExperts) {
222+ return;
223+ }
224+ const int32_t encoded = EncodeExpertProgress(epoch, readyExpertCount);
225+ if (consumerRank == rank_) {
226+ volatile __gm__ int32_t* localSlot = LocalDataReadySlot(rank_) + COMBINE_EXPERT_PROGRESS_OFFSET;
227+ *localSlot = encoded;
228+ dcci((__gm__ void*)localSlot, SINGLE_CACHE_LINE);
229+ dsb(DSB_DDR);
230+ return;
231+ }
232+ __gm__ int32_t* remoteSlot = RemoteSignalBase(consumerRank) + COMBINE_DATA_READY_BASE_INDEX +
233+ static_cast<uint32_t>(rank_) * COMBINE_DATA_READY_STRIDE +
234+ COMBINE_EXPERT_PROGRESS_OFFSET;
235+ auto signal = pto::comm::Signal(remoteSlot);
236+ pto::comm::TNOTIFY(signal, encoded, pto::comm::NotifyOp::Set);
237+ }
238+ 
239+ AICORE inline void AcquireDataReady() const
240+ {
241+ pipe_barrier(PIPE_ALL);
242+ dsb(DSB_DDR);
243+ }
244+ 
245+ // AIV-only snapshot helpers. Each signal owns one 64-byte cache line and
246+ // only the first int32_t carries the epoch or mask.
247+ AICORE inline uint32_t ReadEpochMaskMte(__gm__ int32_t* base, uint32_t count, int32_t epoch) const
248+ {
249+ if (base == nullptr || count == 0U || count > REMOTE_WINDOW_SYNC_MAX_SLOTS) {
250+ return 0U;
251+ }
252+ 
253+ using SnapshotShape = pto::Shape<1, 1, 1, 1, pto::DYNAMIC>;
254+ using SnapshotStride = pto::Stride<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, 1>;
255+ using SnapshotGlobal = pto::GlobalTensor<int32_t, SnapshotShape, SnapshotStride, pto::Layout::ND>;
256+ using SnapshotTile =
257+ pto::Tile<pto::TileType::Vec, int32_t, 1, REMOTE_WINDOW_SYNC_MAX_VALUES, pto::BLayout::RowMajor, -1, -1>;
258+ 
259+ const uint32_t snapshotValues = count * REMOTE_WINDOW_SYNC_VALUES_PER_SLOT;
260+ SnapshotShape snapshotShape(snapshotValues);
261+ SnapshotStride snapshotStride(snapshotValues, snapshotValues, snapshotValues, snapshotValues);
262+ SnapshotGlobal snapshotGlobal(base, snapshotShape, snapshotStride);
263+ SnapshotTile snapshotTile(1U, snapshotValues);
264+ pto::TASSIGN(snapshotTile, REMOTE_WINDOW_SYNC_SNAPSHOT_UB_OFFSET);
265+ pto::TLOAD(snapshotTile, snapshotGlobal);
266+ pto::PtoSetWaitFlag<PIPE_MTE2, PIPE_S>(REMOTE_WINDOW_SYNC_SNAPSHOT_EVENT, REMOTE_WINDOW_SYNC_SNAPSHOT_EVENT);
267+ 
268+ uint32_t readyMask = 0U;
269+ for (uint32_t slot = 0U; slot < count; ++slot) {
270+ if (snapshotTile.GetValue(slot * REMOTE_WINDOW_SYNC_VALUES_PER_SLOT) >= epoch) {
271+ readyMask |= 1U << slot;
272+ }
273+ }
274+ return readyMask;
275+ }
276+ 
277+ AICORE inline void PublishEpochRangeMte(__gm__ int32_t* base, uint32_t count, int32_t value) const
278+ {
279+ if (base == nullptr || count == 0U || count > REMOTE_WINDOW_SYNC_MAX_SLOTS) {
280+ return;
281+ }
282+ 
283+ using SnapshotShape = pto::Shape<1, 1, 1, 1, pto::DYNAMIC>;
284+ using SnapshotStride = pto::Stride<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, 1>;
285+ using SnapshotGlobal = pto::GlobalTensor<int32_t, SnapshotShape, SnapshotStride, pto::Layout::ND>;
286+ using SnapshotTile =
287+ pto::Tile<pto::TileType::Vec, int32_t, 1, REMOTE_WINDOW_SYNC_MAX_VALUES, pto::BLayout::RowMajor, -1, -1>;
288+ 
289+ const uint32_t snapshotValues = count * REMOTE_WINDOW_SYNC_VALUES_PER_SLOT;
290+ SnapshotShape snapshotShape(snapshotValues);
291+ SnapshotStride snapshotStride(snapshotValues, snapshotValues, snapshotValues, snapshotValues);
292+ SnapshotGlobal snapshotGlobal(base, snapshotShape, snapshotStride);
293+ SnapshotTile snapshotTile(1U, snapshotValues);
294+ pto::TASSIGN(snapshotTile, REMOTE_WINDOW_SYNC_SNAPSHOT_UB_OFFSET);
295+ for (uint32_t slot = 0U; slot < count; ++slot) {
296+ snapshotTile.SetValue(slot * REMOTE_WINDOW_SYNC_VALUES_PER_SLOT, value);
297+ }
298+ pto::PtoSetWaitFlag<PIPE_S, PIPE_MTE3>(REMOTE_WINDOW_SYNC_SNAPSHOT_EVENT, REMOTE_WINDOW_SYNC_SNAPSHOT_EVENT);
299+ pto::TSTORE(snapshotGlobal, snapshotTile);
300+ pto::PtoSetWaitFlag<PIPE_MTE3, PIPE_S>(REMOTE_WINDOW_SYNC_SNAPSHOT_EVENT, REMOTE_WINDOW_SYNC_SNAPSHOT_EVENT);
301+ }
302+ 
303+ AICORE inline uint32_t ReadDataReadyMaskMte(int32_t epoch) const
304+ {
305+ return ReadEpochMaskMte(
306+ const_cast<__gm__ int32_t*>(LocalDataReadySlot(0)), static_cast<uint32_t>(rankSize_), epoch);
307+ }
308+ 
309+ AICORE inline uint32_t ReadExpertProgressMte(
310+ int32_t epoch, uint32_t expertPerRank, uint32_t* readyExpertCounts) const
311+ {
312+ if (readyExpertCounts == nullptr || rankSize_ <= 0 ||
313+ static_cast<uint32_t>(rankSize_) > COMBINE_EXPERT_PROGRESS_MAX_RANKS) {
314+ return 0U;
315+ }
316+ 
317+ using SnapshotShape = pto::Shape<1, 1, 1, 1, pto::DYNAMIC>;
318+ using SnapshotStride = pto::Stride<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, 1>;
319+ using SnapshotGlobal = pto::GlobalTensor<int32_t, SnapshotShape, SnapshotStride, pto::Layout::ND>;
320+ using SnapshotTile =
321+ pto::Tile<pto::TileType::Vec, int32_t, 1, REMOTE_WINDOW_SYNC_MAX_VALUES, pto::BLayout::RowMajor, -1, -1>;
322+ 
323+ const uint32_t count = static_cast<uint32_t>(rankSize_);
324+ const uint32_t snapshotValues = count * REMOTE_WINDOW_SYNC_VALUES_PER_SLOT;
325+ SnapshotShape snapshotShape(snapshotValues);
326+ SnapshotStride snapshotStride(snapshotValues, snapshotValues, snapshotValues, snapshotValues);
327+ SnapshotGlobal snapshotGlobal(
328+ const_cast<__gm__ int32_t*>(LocalDataReadySlot(0)), snapshotShape, snapshotStride);
329+ SnapshotTile snapshotTile(1U, snapshotValues);
330+ pto::TASSIGN(snapshotTile, REMOTE_WINDOW_SYNC_SNAPSHOT_UB_OFFSET);
331+ pto::TLOAD(snapshotTile, snapshotGlobal);
332+ pto::PtoSetWaitFlag<PIPE_MTE2, PIPE_S>(REMOTE_WINDOW_SYNC_SNAPSHOT_EVENT, REMOTE_WINDOW_SYNC_SNAPSHOT_EVENT);
333+ 
334+ uint32_t minimumReady = expertPerRank;
335+ const uint32_t expectedEpoch = static_cast<uint32_t>(epoch);
336+ for (uint32_t producerRank = 0U; producerRank < count; ++producerRank) {
337+ const uint32_t encoded = static_cast<uint32_t>(snapshotTile.GetValue(
338+ producerRank * REMOTE_WINDOW_SYNC_VALUES_PER_SLOT + COMBINE_EXPERT_PROGRESS_OFFSET));
339+ const uint32_t observedEpoch = encoded >> COMBINE_EXPERT_PROGRESS_COUNT_BITS;
340+ uint32_t readyCount = observedEpoch == expectedEpoch ? encoded & COMBINE_EXPERT_PROGRESS_COUNT_MASK : 0U;
341+ if (readyCount > expertPerRank) {
342+ readyCount = expertPerRank;
343+ }
344+ readyExpertCounts[producerRank] = readyCount;
345+ minimumReady = readyCount < minimumReady ? readyCount : minimumReady;
346+ }
347+ return minimumReady;
348+ }
349+ 
350+ AICORE inline volatile __gm__ int32_t* LocalUnpermuteAllReadySlot(uint32_t workerIdx) const
351+ {
352+ if (workerIdx >= UNPERMUTE_ALL_READY_SLOT_COUNT) {
353+ return nullptr;
354+ }
355+ return LocalSignalBase() + UNPERMUTE_ALL_READY_BASE_INDEX + workerIdx * UNPERMUTE_ALL_READY_STRIDE;
356+ }
357+ 
358+ AICORE inline void PublishUnpermuteAllReady(uint32_t workerCount, int32_t epoch) const
359+ {
360+ const uint32_t count =
361+ workerCount < UNPERMUTE_ALL_READY_SLOT_COUNT ? workerCount : UNPERMUTE_ALL_READY_SLOT_COUNT;
362+ PublishEpochRangeMte(const_cast<__gm__ int32_t*>(LocalUnpermuteAllReadySlot(0)), count, epoch);
363+ }
364+ 
365+ AICORE inline volatile __gm__ int32_t* LocalUnpermuteStartSlot(uint32_t workerIdx) const
366+ {
367+ if (workerIdx >= UNPERMUTE_START_SLOT_COUNT) {
368+ return nullptr;
369+ }
370+ return LocalSignalBase() + UNPERMUTE_START_BASE_INDEX + workerIdx * UNPERMUTE_START_STRIDE;
371+ }
372+ 
373+ AICORE inline void PublishUnpermuteStart(uint32_t workerCount, int32_t epoch) const
374+ {
375+ const uint32_t count = workerCount < UNPERMUTE_START_SLOT_COUNT ? workerCount : UNPERMUTE_START_SLOT_COUNT;
376+ for (uint32_t workerIdx = 0U; workerIdx < count; ++workerIdx) {
377+ *LocalUnpermuteStartSlot(workerIdx) = epoch;
378+ }
379+ for (uint32_t workerIdx = 0U; workerIdx < count; ++workerIdx) {
380+ dcci((__gm__ void*)LocalUnpermuteStartSlot(workerIdx), SINGLE_CACHE_LINE);
381+ }
382+ dsb(DSB_DDR);
383+ __asm__ __volatile__("");
384+ }
385+ 
386+ AICORE inline void PublishUnpermuteStartRangeMte(uint32_t firstWorker, uint32_t workerCount, int32_t epoch) const
387+ {
388+ if (firstWorker >= UNPERMUTE_START_SLOT_COUNT) {
389+ return;
390+ }
391+ const uint32_t remaining = UNPERMUTE_START_SLOT_COUNT - firstWorker;
392+ const uint32_t count = workerCount < remaining ? workerCount : remaining;
393+ PublishEpochRangeMte(const_cast<__gm__ int32_t*>(LocalUnpermuteStartSlot(firstWorker)), count, epoch);
394+ }
395+ 
396+ AICORE inline volatile __gm__ int32_t* LocalUnpermutePhase1ProgressEpochSlot() const
397+ {
398+ return LocalSignalBase() + UNPERMUTE_PHASE1_PROGRESS_BASE_INDEX + UNPERMUTE_PHASE1_PROGRESS_EPOCH_OFFSET;
399+ }
400+ 
401+ AICORE inline void PublishUnpermutePhase1Progress(
402+ const uint32_t* readyExpertCounts, uint32_t rankCount, uint32_t readyRankMask, int32_t epoch) const
403+ {
404+ if (readyExpertCounts == nullptr || rankCount == 0U || rankCount > COMBINE_EXPERT_PROGRESS_MAX_RANKS) {
405+ return;
406+ }
407+ volatile __gm__ int32_t* base = LocalSignalBase() + UNPERMUTE_PHASE1_PROGRESS_BASE_INDEX;
408+ base[UNPERMUTE_PHASE1_PROGRESS_MASK_OFFSET] = static_cast<int32_t>(readyRankMask);
409+ for (uint32_t producerRank = 0U; producerRank < rankCount; ++producerRank) {
410+ base[UNPERMUTE_PHASE1_PROGRESS_COUNTS_OFFSET + producerRank] =
411+ static_cast<int32_t>(readyExpertCounts[producerRank]);
412+ }
413+ DcciUnpermutePhase1Progress(base, rankCount);
414+ dsb(DSB_DDR);
415+ base[UNPERMUTE_PHASE1_PROGRESS_EPOCH_OFFSET] = epoch;
416+ dcci((__gm__ void*)base, SINGLE_CACHE_LINE);
417+ dsb(DSB_DDR);
418+ }
419+ 
420+ AICORE inline uint32_t ReadUnpermutePhase1Progress(uint32_t* readyExpertCounts, uint32_t rankCount) const
421+ {
422+ if (readyExpertCounts == nullptr || rankCount == 0U || rankCount > COMBINE_EXPERT_PROGRESS_MAX_RANKS) {
423+ return 0U;
424+ }
425+ volatile __gm__ int32_t* base = LocalSignalBase() + UNPERMUTE_PHASE1_PROGRESS_BASE_INDEX;
426+ DcciUnpermutePhase1Progress(base, rankCount);
427+ dsb(DSB_DDR);
428+ for (uint32_t producerRank = 0U; producerRank < rankCount; ++producerRank) {
429+ readyExpertCounts[producerRank] =
430+ static_cast<uint32_t>(base[UNPERMUTE_PHASE1_PROGRESS_COUNTS_OFFSET + producerRank]);
431+ }
432+ return static_cast<uint32_t>(base[UNPERMUTE_PHASE1_PROGRESS_MASK_OFFSET]);
433+ }
434+ 
435+ AICORE inline volatile __gm__ int32_t* LocalDispatchReleaseSlot(uint32_t workerIdx) const
436+ {
437+ if (workerIdx >= UNPERMUTE_DISPATCH_RELEASE_SLOT_COUNT) {
438+ return nullptr;
439+ }
440+ return LocalSignalBase() + UNPERMUTE_DISPATCH_RELEASE_BASE_INDEX +
441+ workerIdx * UNPERMUTE_DISPATCH_RELEASE_STRIDE;
442+ }
443+ 
444+ AICORE inline volatile __gm__ int32_t* LocalSwigluReleaseSlot(uint32_t workerIdx) const
445+ {
446+ if (workerIdx >= UNPERMUTE_SWIGLU_RELEASE_SLOT_COUNT) {
447+ return nullptr;
448+ }
449+ return LocalSignalBase() + UNPERMUTE_SWIGLU_RELEASE_BASE_INDEX + workerIdx * UNPERMUTE_SWIGLU_RELEASE_STRIDE;
450+ }
451+ 
452+ AICORE inline void PublishDispatchRelease(uint32_t workerIdx, int32_t epoch) const
453+ {
454+ volatile __gm__ int32_t* slot = LocalDispatchReleaseSlot(workerIdx);
455+ if (slot == nullptr) {
456+ return;
457+ }
458+ *slot = epoch;
459+ dcci((__gm__ void*)slot, SINGLE_CACHE_LINE);
460+ dsb(DSB_DDR);
461+ }
462+ 
463+ AICORE inline void PublishSwigluRelease(uint32_t workerIdx, int32_t epoch) const
464+ {
465+ volatile __gm__ int32_t* slot = LocalSwigluReleaseSlot(workerIdx);
466+ if (slot == nullptr) {
467+ return;
468+ }
469+ *slot = epoch;
470+ dcci((__gm__ void*)slot, SINGLE_CACHE_LINE);
471+ dsb(DSB_DDR);
472+ }
473+ 
474+ AICORE inline void WaitDispatchReleaseMte(uint32_t workerCount, int32_t epoch) const
475+ {
476+ const uint32_t count =
477+ workerCount < UNPERMUTE_DISPATCH_RELEASE_SLOT_COUNT ? workerCount : UNPERMUTE_DISPATCH_RELEASE_SLOT_COUNT;
478+ const uint32_t allReadyMask = count >= 32U ? 0xFFFFFFFFU : ((1U << count) - 1U);
479+ __gm__ int32_t* base = const_cast<__gm__ int32_t*>(LocalDispatchReleaseSlot(0U));
480+ while (ReadEpochMaskMte(base, count, epoch) != allReadyMask) {
481+ RemoteWindowSyncPollBackoff();
482+ }
483+ AcquireDataReady();
484+ }
485+ 
486+ AICORE inline void WaitSwigluReleaseMte(uint32_t workerCount, int32_t epoch) const
487+ {
488+ const uint32_t count =
489+ workerCount < UNPERMUTE_SWIGLU_RELEASE_SLOT_COUNT ? workerCount : UNPERMUTE_SWIGLU_RELEASE_SLOT_COUNT;
490+ const uint32_t allReadyMask = count >= 32U ? 0xFFFFFFFFU : ((1U << count) - 1U);
491+ __gm__ int32_t* base = const_cast<__gm__ int32_t*>(LocalSwigluReleaseSlot(0U));
492+ while (ReadEpochMaskMte(base, count, epoch) != allReadyMask) {
493+ RemoteWindowSyncPollBackoff();
494+ }
495+ AcquireDataReady();
496+ }
497+ 
498+ AICORE inline volatile __gm__ int32_t* LocalPhase1DoneSlot(uint32_t workerIdx) const
499+ {
500+ if (workerIdx >= UNPERMUTE_PHASE1_DONE_SLOT_COUNT) {
501+ return nullptr;
502+ }
503+ return LocalSignalBase() + UNPERMUTE_PHASE1_DONE_BASE_INDEX + workerIdx * UNPERMUTE_PHASE1_DONE_STRIDE;
504+ }
505+ 
506+ AICORE inline void PublishPhase1Done(uint32_t workerIdx, int32_t epoch) const
507+ {
508+ volatile __gm__ int32_t* slot = LocalPhase1DoneSlot(workerIdx);
509+ if (slot == nullptr) {
510+ return;
511+ }
512+ *slot = epoch;
513+ dcci((__gm__ void*)slot, SINGLE_CACHE_LINE);
514+ dsb(DSB_DDR);
515+ }
516+ 
517+ AICORE inline void WaitPhase1DoneMte(uint32_t workerCount, int32_t epoch) const
518+ {
519+ const uint32_t count =
520+ workerCount < UNPERMUTE_PHASE1_DONE_SLOT_COUNT ? workerCount : UNPERMUTE_PHASE1_DONE_SLOT_COUNT;
521+ const uint32_t allReadyMask = count == 32U ? 0xFFFFFFFFU : ((1U << count) - 1U);
522+ __gm__ int32_t* base = const_cast<__gm__ int32_t*>(LocalPhase1DoneSlot(0U));
523+ while (ReadEpochMaskMte(base, count, epoch) != allReadyMask) {
524+ RemoteWindowSyncPollBackoff();
525+ }
526+ AcquireDataReady();
527+ }
528+ 
529+ AICORE inline volatile __gm__ int32_t* LocalCombineDoneSlot(uint32_t laneIdx) const
530+ {
531+ if (laneIdx >= COMBINE_LOCAL_DONE_SLOT_COUNT) {
532+ return nullptr;
533+ }
534+ return LocalSignalBase() + COMBINE_LOCAL_DONE_BASE_INDEX + laneIdx * COMBINE_LOCAL_DONE_STRIDE;
535+ }
536+ 
537+ AICORE inline void PublishLocalCombineDone(uint32_t laneIdx, int32_t epoch) const
538+ {
539+ volatile __gm__ int32_t* slot = LocalCombineDoneSlot(laneIdx);
540+ if (slot == nullptr) {
541+ return;
542+ }
543+ *slot = epoch;
544+ dcci((__gm__ void*)slot, SINGLE_CACHE_LINE);
545+ __asm__ __volatile__("");
546+ }
547+ 
548+ AICORE inline void WaitLocalCombineDoneMte(uint32_t laneCount, int32_t epoch) const
549+ {
550+ const uint32_t count = laneCount < COMBINE_LOCAL_DONE_SLOT_COUNT ? laneCount : COMBINE_LOCAL_DONE_SLOT_COUNT;
551+ const uint32_t allReadyMask = count >= 32U ? 0xFFFFFFFFU : ((1U << count) - 1U);
552+ __gm__ int32_t* base = const_cast<__gm__ int32_t*>(LocalCombineDoneSlot(0));
553+ while (ReadEpochMaskMte(base, count, epoch) != allReadyMask) {
554+ RemoteWindowSyncPollBackoff();
555+ }
556+ pipe_barrier(PIPE_ALL);
557+ dsb(DSB_DDR);
558+ }
559+ 
77 AICORE inline void CrossRankSync() const560 AICORE inline void CrossRankSync() const
78 {561 {
79 __gm__ int32_t* localSignalBase = LocalSignalBase();562 __gm__ int32_t* localSignalBase = LocalSignalBase();
Akernels/manual/a2a3/dispatch_mega_combine/op_kernel/utils/mega_expert_sync.hpp+315-0
@@ -0,0 +1,315 @@
1+/**
2+Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+CANN Open Software License Agreement Version 2.0 (the "License").
5+Please refer to the License for details. You may not use this file except in compliance with the License.
6+THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+See LICENSE in the root of the software repository for the full text of the License.
9+*/
10+ 
11+#ifndef DISPATCH_MEGA_COMBINE_MEGA_EXPERT_SYNC_HPP
12+#define DISPATCH_MEGA_COMBINE_MEGA_EXPERT_SYNC_HPP
13+ 
14+#include "kernel_operator.h"
15+ 
16+#include <pto/common/pto_tile.hpp>
17+#include <pto/pto-inst.hpp>
18+ 
19+#include "dispatch_mega_combine_tiling.h"
20+#include "const_args.hpp"
21+ 
22+struct MegaMoeFixedCoreRoleInfo {
23+ uint32_t role = 0U;
24+ uint32_t physicalBlockId = 0U;
25+ uint32_t subblockId = 0U;
26+ uint32_t flatAivId = 0U;
27+ uint32_t groupLocalId = 0U;
28+ uint32_t groupSize = 0U;
29+};
30+ 
31+AICORE inline MegaMoeFixedCoreRoleInfo FixedCoreRole(const __gm__ MegaMoeTilingData* tilingData)
32+{
33+ MegaMoeFixedCoreRoleInfo info;
34+ const __gm__ MegaMoeFixedGroupTiling& fixed = tilingData->fixedGroupTiling;
35+ info.physicalBlockId = get_block_idx();
36+ const bool firstGroup = info.physicalBlockId < fixed.gmm1GroupSize;
37+ if ASCEND_IS_AIC {
38+ if (firstGroup) {
39+ info.role = kMegaMoeFixedRoleGmm1;
40+ info.groupLocalId = info.physicalBlockId;
41+ info.groupSize = fixed.gmm1GroupSize;
42+ } else {
43+ info.role = kMegaMoeFixedRoleGmm2;
44+ info.groupLocalId = info.physicalBlockId - fixed.gmm1GroupSize;
45+ info.groupSize = fixed.gmm2GroupSize;
46+ }
47+ return info;
48+ }
49+ 
50+ info.subblockId = get_subblockid();
51+ info.flatAivId = info.physicalBlockId + info.subblockId * fixed.physicalAicNum;
52+ if (firstGroup) {
53+ info.groupLocalId = info.physicalBlockId;
54+ if (info.subblockId == 0U) {
55+ info.role = kMegaMoeFixedRoleDispatch;
56+ info.groupSize = fixed.dispatchGroupSize;
57+ } else {
58+ info.role = kMegaMoeFixedRoleSwiglu;
59+ info.groupSize = fixed.swigluGroupSize;
60+ }
61+ } else {
62+ info.role = kMegaMoeFixedRoleCombine;
63+ const uint32_t pairedGmm2LocalId = info.physicalBlockId - fixed.gmm1GroupSize;
64+ info.groupLocalId = pairedGmm2LocalId + info.subblockId * fixed.gmm2GroupSize;
65+ info.groupSize = fixed.combineGroupSize;
66+ }
67+ return info;
68+}
69+ 
70+AICORE inline bool FixedRoleActiveForStage(const MegaMoeFixedCoreRoleInfo& info, uint32_t stageNum)
71+{
72+ return (info.role == kMegaMoeFixedRoleDispatch && stageNum >= 9U) ||
73+ (info.role == kMegaMoeFixedRoleGmm1 && stageNum >= 10U) ||
74+ (info.role == kMegaMoeFixedRoleSwiglu && stageNum >= 11U) ||
75+ (info.role == kMegaMoeFixedRoleGmm2 && stageNum >= 12U) ||
76+ (info.role == kMegaMoeFixedRoleCombine && info.groupLocalId < info.groupSize && stageNum >= 13U);
77+}
78+ 
79+AICORE inline MegaMoeSyncLayout FixedSyncLayout(const __gm__ MegaMoeTilingData* tilingData)
80+{
81+ const __gm__ MegaMoeFixedGroupTiling& fixed = tilingData->fixedGroupTiling;
82+ MegaMoeSyncLayout layout;
83+ layout.dispatchArrivalBase = kMegaMoeFixedSyncHeadCanarySlot + 1U;
84+ layout.dispatchReadyBase = layout.dispatchArrivalBase + fixed.dispatchGroupSize;
85+ layout.gmm1ArrivalBase = layout.dispatchReadyBase + fixed.physicalAicNum;
86+ layout.swigluReadyBase = layout.gmm1ArrivalBase + fixed.gmm1GroupSize;
87+ layout.swigluArrivalBase = layout.swigluReadyBase + fixed.swigluGroupSize;
88+ layout.gmm2ReadyBase = layout.swigluArrivalBase + fixed.swigluGroupSize;
89+ layout.gmm2ArrivalBase = layout.gmm2ReadyBase + fixed.gmm2GroupSize;
90+ layout.combineReadyBase = layout.gmm2ArrivalBase + fixed.physicalAicNum;
91+ layout.gmm1DoneSlot = layout.combineReadyBase + fixed.gmm2GroupSize;
92+ layout.gmm2JoinSlot = layout.gmm1DoneSlot + 1U;
93+ layout.tailCanarySlot = layout.gmm2JoinSlot + 1U;
94+ layout.slotCount = layout.tailCanarySlot + 1U;
95+ return layout;
96+}
97+ 
98+AICORE inline volatile __gm__ int32_t* FixedSyncSlot(
99+ GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData, uint32_t slot)
100+{
101+ const __gm__ MegaMoeFixedGroupTiling& fixed = tilingData->fixedGroupTiling;
102+ return reinterpret_cast<volatile __gm__ int32_t*>(
103+ workspaceGM + fixed.syncOffset + static_cast<uint64_t>(slot) * fixed.syncSlotBytes);
104+}
105+ 
106+// The doorbell is a scalar GM store. The caller drains payload pipelines before publishing.
107+AICORE inline void PublishScalarEpoch(volatile __gm__ int32_t* slot, int32_t epoch)
108+{
109+ *slot = epoch;
110+ dcci((__gm__ void*)slot, SINGLE_CACHE_LINE);
111+ __asm__ __volatile__("");
112+}
113+ 
114+AICORE inline void PublishScalarEpochRange(
115+ GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData, uint32_t baseSlot, uint32_t count, int32_t epoch)
116+{
117+ for (uint32_t localId = 0U; localId < count; ++localId) {
118+ *FixedSyncSlot(workspaceGM, tilingData, baseSlot + localId) = epoch;
119+ }
120+ for (uint32_t localId = 0U; localId < count; ++localId) {
121+ dcci((__gm__ void*)FixedSyncSlot(workspaceGM, tilingData, baseSlot + localId), SINGLE_CACHE_LINE);
122+ }
123+ __asm__ __volatile__("");
124+}
125+ 
126+AICORE inline void ResetFixedSyncWorkspace(GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData)
127+{
128+ const MegaMoeSyncLayout sync = FixedSyncLayout(tilingData);
129+ for (uint32_t slotId = 0U; slotId < sync.slotCount; ++slotId) {
130+ volatile __gm__ int32_t* slot = FixedSyncSlot(workspaceGM, tilingData, slotId);
131+ dcci((__gm__ void*)slot, SINGLE_CACHE_LINE);
132+ __asm__ __volatile__("");
133+ uint32_t value = 0U;
134+ if (slotId == kMegaMoeFixedSyncHeadCanarySlot) {
135+ value = kMegaMoeFixedHeadCanary;
136+ } else if (slotId == sync.tailCanarySlot) {
137+ value = kMegaMoeFixedTailCanary;
138+ }
139+ *slot = static_cast<int32_t>(value);
140+ }
141+ for (uint32_t slotId = 0U; slotId < sync.slotCount; ++slotId) {
142+ dcci((__gm__ void*)FixedSyncSlot(workspaceGM, tilingData, slotId), SINGLE_CACHE_LINE);
143+ }
144+ dsb(DSB_DDR);
145+ pipe_barrier(PIPE_ALL);
146+}
147+ 
148+AICORE inline int32_t ReadScalarEpoch(volatile __gm__ int32_t* slot)
149+{
150+ dcci((__gm__ void*)slot, SINGLE_CACHE_LINE);
151+ __asm__ __volatile__("");
152+ return *slot;
153+}
154+ 
155+AICORE inline bool TestEpoch(volatile __gm__ int32_t* slot, int32_t epoch, int32_t& observed)
156+{
157+ observed = ReadScalarEpoch(slot);
158+ return observed >= epoch;
159+}
160+ 
161+AICORE inline void EpochPollBackoff()
162+{
163+ constexpr uint32_t kDelayTicks = 3U;
164+ const uint64_t deadline = get_sys_cnt() + kDelayTicks;
165+ while (get_sys_cnt() < deadline) {
166+ __asm__ __volatile__("");
167+ }
168+}
169+ 
170+AICORE inline int32_t WaitEpochRaw(volatile __gm__ int32_t* slot, int32_t epoch)
171+{
172+ while (true) {
173+ int32_t observed = 0;
174+ if (TestEpoch(slot, epoch, observed)) {
175+ return observed;
176+ }
177+ EpochPollBackoff();
178+ }
179+}
180+ 
181+AICORE inline int32_t WaitEpochAcquire(volatile __gm__ int32_t* slot, int32_t epoch)
182+{
183+ const int32_t observed = WaitEpochRaw(slot, epoch);
184+ pipe_barrier(PIPE_ALL);
185+ dsb(DSB_DDR);
186+ return observed;
187+}
188+ 
189+// These snapshot helpers require AIV UB and MTE2/MTE3; AIC callers publish scalar arrivals only.
190+constexpr uint32_t kMegaMoeSyncSnapshotValuesPerSlot = kMegaMoeFixedSyncSlotBytes / sizeof(int32_t);
191+constexpr uint32_t kMegaMoeSyncSnapshotMaxValues = kMegaMoeFixedPhysicalAicNum * kMegaMoeSyncSnapshotValuesPerSlot;
192+constexpr uint64_t kMegaMoeSyncSnapshotUbBytes = static_cast<uint64_t>(kMegaMoeSyncSnapshotMaxValues) * sizeof(int32_t);
193+constexpr uint64_t kMegaMoeSyncSnapshotUbOffset = AtlasA2::UB_SIZE - kMegaMoeSyncSnapshotUbBytes;
194+constexpr event_t kMegaMoeSyncSnapshotEvent = EVENT_ID0;
195+ 
196+static_assert(kMegaMoeFixedSyncSlotBytes % sizeof(int32_t) == 0U);
197+static_assert(kMegaMoeSyncSnapshotUbOffset % UB_ALIGN == 0U);
198+ 
199+AICORE inline int32_t ReadArrivalMinMte(
200+ GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData, uint32_t arrivalBaseSlot, uint32_t producerCount)
201+{
202+ if (producerCount == 0U || producerCount > kMegaMoeFixedPhysicalAicNum) {
203+ return 0;
204+ }
205+ 
206+ using SnapshotShape = pto::Shape<1, 1, 1, 1, pto::DYNAMIC>;
207+ using SnapshotStride = pto::Stride<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, 1>;
208+ using SnapshotGlobal = pto::GlobalTensor<int32_t, SnapshotShape, SnapshotStride, pto::Layout::ND>;
209+ using SnapshotTile =
210+ pto::Tile<pto::TileType::Vec, int32_t, 1, kMegaMoeSyncSnapshotMaxValues, pto::BLayout::RowMajor, -1, -1>;
211+ 
212+ const __gm__ MegaMoeFixedGroupTiling& fixed = tilingData->fixedGroupTiling;
213+ __gm__ int32_t* arrivalBase = reinterpret_cast<__gm__ int32_t*>(
214+ workspaceGM + fixed.syncOffset + static_cast<uint64_t>(arrivalBaseSlot) * fixed.syncSlotBytes);
215+ const uint32_t snapshotValues = producerCount * kMegaMoeSyncSnapshotValuesPerSlot;
216+ SnapshotShape snapshotShape(snapshotValues);
217+ SnapshotStride snapshotStride(snapshotValues, snapshotValues, snapshotValues, snapshotValues);
218+ SnapshotGlobal snapshotGlobal(arrivalBase, snapshotShape, snapshotStride);
219+ SnapshotTile snapshotTile(1U, snapshotValues);
220+ pto::TASSIGN(snapshotTile, kMegaMoeSyncSnapshotUbOffset);
221+ pto::TLOAD(snapshotTile, snapshotGlobal);
222+ pto::PtoSetWaitFlag<PIPE_MTE2, PIPE_S>(kMegaMoeSyncSnapshotEvent, kMegaMoeSyncSnapshotEvent);
223+ 
224+ int32_t observedMin = snapshotTile.GetValue(0U);
225+ for (uint32_t peer = 1U; peer < producerCount; ++peer) {
226+ const int32_t observed = snapshotTile.GetValue(peer * kMegaMoeSyncSnapshotValuesPerSlot);
227+ observedMin = observed < observedMin ? observed : observedMin;
228+ }
229+ return observedMin;
230+}
231+ 
232+AICORE inline int32_t WaitArrivalMinMte(
233+ GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData, uint32_t arrivalBaseSlot, uint32_t producerCount,
234+ int32_t epoch)
235+{
236+ while (true) {
237+ const int32_t observedMin = ReadArrivalMinMte(workspaceGM, tilingData, arrivalBaseSlot, producerCount);
238+ if (observedMin >= epoch) {
239+ return observedMin;
240+ }
241+ EpochPollBackoff();
242+ }
243+}
244+ 
245+AICORE inline void PublishEpochRangeMte(
246+ GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData, uint32_t baseSlot, uint32_t count, int32_t epoch)
247+{
248+ if (count == 0U || count > kMegaMoeFixedPhysicalAicNum) {
249+ return;
250+ }
251+ 
252+ using SnapshotShape = pto::Shape<1, 1, 1, 1, pto::DYNAMIC>;
253+ using SnapshotStride = pto::Stride<pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, pto::DYNAMIC, 1>;
254+ using SnapshotGlobal = pto::GlobalTensor<int32_t, SnapshotShape, SnapshotStride, pto::Layout::ND>;
255+ using SnapshotTile =
256+ pto::Tile<pto::TileType::Vec, int32_t, 1, kMegaMoeSyncSnapshotMaxValues, pto::BLayout::RowMajor, -1, -1>;
257+ 
258+ const __gm__ MegaMoeFixedGroupTiling& fixed = tilingData->fixedGroupTiling;
259+ __gm__ int32_t* readyBase = reinterpret_cast<__gm__ int32_t*>(
260+ workspaceGM + fixed.syncOffset + static_cast<uint64_t>(baseSlot) * fixed.syncSlotBytes);
261+ const uint32_t snapshotValues = count * kMegaMoeSyncSnapshotValuesPerSlot;
262+ SnapshotShape snapshotShape(snapshotValues);
263+ SnapshotStride snapshotStride(snapshotValues, snapshotValues, snapshotValues, snapshotValues);
264+ SnapshotGlobal snapshotGlobal(readyBase, snapshotShape, snapshotStride);
265+ SnapshotTile snapshotTile(1U, snapshotValues);
266+ pto::TASSIGN(snapshotTile, kMegaMoeSyncSnapshotUbOffset);
267+ for (uint32_t localId = 0U; localId < count; ++localId) {
268+ snapshotTile.SetValue(localId * kMegaMoeSyncSnapshotValuesPerSlot, epoch);
269+ }
270+ pto::PtoSetWaitFlag<PIPE_S, PIPE_MTE3>(kMegaMoeSyncSnapshotEvent, kMegaMoeSyncSnapshotEvent);
271+ pto::TSTORE(snapshotGlobal, snapshotTile);
272+ pto::PtoSetWaitFlag<PIPE_MTE3, PIPE_S>(kMegaMoeSyncSnapshotEvent, kMegaMoeSyncSnapshotEvent);
273+}
274+ 
275+AICORE inline void PublishGroupArrival(
276+ GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData, uint32_t arrivalBaseSlot, uint32_t localId,
277+ uint32_t notifyCall)
278+{
279+ const int32_t epoch = static_cast<int32_t>(notifyCall * 2U + 1U);
280+ pipe_barrier(PIPE_ALL);
281+ dsb(DSB_DDR);
282+ PublishScalarEpoch(FixedSyncSlot(workspaceGM, tilingData, arrivalBaseSlot + localId), epoch);
283+}
284+ 
285+AICORE inline void CoordinateGroupConsumersMte(
286+ GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData, uint32_t arrivalBaseSlot, uint32_t readyBaseSlot,
287+ uint32_t producerCount, uint32_t consumerCount, uint32_t notifyCall)
288+{
289+ const int32_t arriveEpoch = static_cast<int32_t>(notifyCall * 2U + 1U);
290+ const int32_t releaseEpoch = arriveEpoch + 1;
291+ WaitArrivalMinMte(workspaceGM, tilingData, arrivalBaseSlot, producerCount, arriveEpoch);
292+ pipe_barrier(PIPE_ALL);
293+ dsb(DSB_DDR);
294+ PublishEpochRangeMte(workspaceGM, tilingData, readyBaseSlot, consumerCount, releaseEpoch);
295+}
296+ 
297+AICORE inline void NotifyGroupConsumersMte(
298+ GM_ADDR workspaceGM, const __gm__ MegaMoeTilingData* tilingData, uint32_t arrivalBaseSlot, uint32_t readyBaseSlot,
299+ uint32_t producerCount, uint32_t consumerCount, uint32_t localId, uint32_t coordinatorLocalId, uint32_t notifyCall)
300+{
301+ const int32_t arriveEpoch = static_cast<int32_t>(notifyCall * 2U + 1U);
302+ const int32_t releaseEpoch = arriveEpoch + 1;
303+ pipe_barrier(PIPE_ALL);
304+ dsb(DSB_DDR);
305+ PublishScalarEpoch(FixedSyncSlot(workspaceGM, tilingData, arrivalBaseSlot + localId), arriveEpoch);
306+ 
307+ if (localId == coordinatorLocalId) {
308+ WaitArrivalMinMte(workspaceGM, tilingData, arrivalBaseSlot, producerCount, arriveEpoch);
309+ pipe_barrier(PIPE_ALL);
310+ dsb(DSB_DDR);
311+ PublishEpochRangeMte(workspaceGM, tilingData, readyBaseSlot, consumerCount, releaseEpoch);
312+ }
313+}
314+ 
315+#endif // DISPATCH_MEGA_COMBINE_MEGA_EXPERT_SYNC_HPP