[Question] Question about CU pressure from mori dispatch/combine

[dongmin-ra]

Hello, team.

I have a question about CU pressure from mori dispatch/combine kernels.

When setting up a mori configuration, I noticed that the grid and block dimensions are adjusted using the block_num and warp_num_per_block parameters.

I'm wondering if these values can give any indication of roughly how many CUs (Compute Units) Mori is utilizing.

Additionally, is there any guidance on how to adjust the warp_num_per_block and block_num parameters in different scenarios to achieve better performance?

[whitememory]

FYI,
We created this issue while answering the question at vllm mori PR.

[jhchouuu]

Hi @dongmin-ra
Currently, on the internode, we are using block_num = 64 and warp_num_per_block = 16. On the intranode, we are using block_num = 80, and dispatch uses warp_num_per_block = 16, while combine uses warp_num_per_block = 8.
But this is not the final version. We are working to optimize the kernel, and the number of cu usages is also one of the optimization targets. However, for now, you can use it like this.

[TianDi101]

number of CU = block_num

For best configurations, my suggestion is to write a script to tune different parameters under different token number.

I believe we can pre-tune these parameters and turn the tuning results into a heuristic algorithm, but team has not got enough bandwidth to do this recently, we may schedule this feature after recent performance optimization on inter-node kernel.

[whitememory]

@jhchouuu @TianDi101
Thank you for swift answers!

When we use computation-communication overlap, number of cu usage is crucial for the optimal performance. (We do not know details of mori-ccl comp-comm overlap. Currently, we use stream and event to overlap microbatches)

A guideline to tune parameters to optimize the performance for both normal case and overlap-case may help everyone to get greater number from mori.

Anyways, we will test below configs in our performance testings.

Currently, on the internode, we are using block_num = 64 and warp_num_per_block = 16. On the intranode, we are using block_num = 80, and dispatch uses warp_num_per_block = 16, while combine uses warp_num_per_block = 8.

[whitememory]

@jhchouuu @TianDi101

On the intranode, we are using block_num = 80, and dispatch uses warp_num_per_block = 16, while combine uses warp_num_per_block = 8.

I have one more question about above comment. (I realized your warp_num_per_block value is different for dispatch and combine)

is it possible to have different configs for dispatch and combine in a same moe layer?
we'd been using one config to get one EpDispatchCombineOp, and we executed dispatch and combine from the EpDispatchCombineOp.

like

op1 = mori.ops.EpDispatchCombineOp(config1)
...
xxx = op1.dispatch()
...
yyy = op1.combine()

vs

op1 = mori.ops.EpDispatchCombineOp(config1)
op2 = mori.ops.EpDispatchCombineOp(config2)
...
xxx = op1.dispatch()
...
yyy = op2.combine()

[jhchouuu]

Hi @whitememory
you can change/adjust block_num and warp_num_per_block when calling dispatch and combine,
like

op = mori.ops.EpDispatchCombineOp(config)
xxx = op.dispatch(..., block_num=80, warp_per_block=16)
...
yyy = op.combine(..., block_num=80, warp_per_block=8)

[whitememory]

Thank you, @jhchouuu

we confirmed that tuning warp_per_block to 16 only at dispatch in intranode setup increases overall performance.
(we did not measure dispatch performance, we used vllm benchmark result)