Optimize the mhc kernels#14
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes and simplifies several MHC kernels. In norm_fn_kernel.py, the sum of squares computation is refactored to use T.reduce_sum on a 1D fragment instead of a 2D fragment. In pre_apply_mix_kernel.py, intermediate shared memory buffers are removed to simplify data copying. In sinkhorn.py, the forward kernel now dynamically determines the token_block_size based on the input size. Feedback suggests that the backward kernel in sinkhorn.py should also use this dynamic block size to prevent potential out-of-bounds accesses, and the pipelined loop in pre_apply_mix_kernel.py should be changed to a serial loop since shared memory is no longer used.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| else: | ||
| token_block_size = 1 | ||
| fwd_kernel = _mhc_sinkhorn_fwd(hidden_size, token_block_size, repeat, eps) | ||
| bwd_kernel = _mhc_sinkhorn_bwd(hidden_size, 32, repeat, eps) |
There was a problem hiding this comment.
The backward kernel _mhc_sinkhorn_bwd is hardcoded to use a token_block_size of 32. If n is not a multiple of 32 (e.g., when token_block_size is dynamically chosen as 16, 4, or 1 for the forward kernel), the backward kernel will perform out-of-bounds memory accesses on the last block. The backward kernel should use the same dynamically chosen token_block_size as the forward kernel to ensure correctness.
| bwd_kernel = _mhc_sinkhorn_bwd(hidden_size, 32, repeat, eps) | |
| bwd_kernel = _mhc_sinkhorn_bwd(hidden_size, token_block_size, repeat, eps) |
| @@ -35,21 +35,17 @@ def _mhc_pre_apply_mix_fwd_kernel( | |||
| T.copy(mix[pid_n, 0], mixl) | |||
|
|
|||
| for i0_h in T.Pipelined(h // h_blk, num_stages=2): | |||
There was a problem hiding this comment.
Since the shared memory buffers (xs and os) have been removed, software pipelining (T.Pipelined) is no longer effective or supported by the compiler for direct global-to-register copies. It should be changed to a standard serial loop (T.serial) to avoid compilation issues or overhead.
| for i0_h in T.Pipelined(h // h_blk, num_stages=2): | |
| for i0_h in T.serial(h // h_blk): |
Tune MHC forward and backward kernels across norm, pre/post processing, expand, sinkhorn, and multilayer recompute paths. Use more efficient fragment dataflow, specialized zero-gradient paths, shape-aware launch configuration, and C500 104-AP persistent CTA dispatch to improve throughput while preserving correctness.
Pack routing comparisons into stable uint64 keys, use 32-lane subgroup execution for top-k and grouped routing, and specialize short top-k fused mapping workloads. Improve fused expansion scheduling and mapping correctness for the 32-bit match-any lowering. Add cleared-L2 MoE benchmark timing and coverage for stable ties and special floating-point values.
1 participant
No description provided.