diff --git a/megatron/core/transformer/moe/sgl_fused_moe/fused_moe_triton_backward_kernels.py b/megatron/core/transformer/moe/sgl_fused_moe/fused_moe_triton_backward_kernels.py index bef19e4cf2c..3c8b5e66750 100644 --- a/megatron/core/transformer/moe/sgl_fused_moe/fused_moe_triton_backward_kernels.py +++ b/megatron/core/transformer/moe/sgl_fused_moe/fused_moe_triton_backward_kernels.py @@ -121,3 +121,144 @@ def fused_moe_backward_input_kernel( ) grad_input_mask = token_mask[:, None] & (curr_offs_k[None, :] < K) tl.atomic_add(grad_input_ptrs, contribution.to(compute_type), mask=grad_input_mask) + + +@triton.jit +def fused_moe_backward_weight_kernel( + # Pointers to matrices + grad_output_ptr, + input_ptr, + grad_weight_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + # Matrix dimensions + N, + K, + EM, + num_valid_tokens, + # Strides + stride_gom, + stride_gon, + stride_im, + stride_ik, + stride_gwe, + stride_gwn, + stride_gwk, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, + top_k: tl.constexpr, + compute_type: tl.constexpr, +): + """ + Backward kernel for computing grad_weight. + + Forward: output = input @ weight.T, optionally scaled by top-k weights. + Backward accumulates each expert's weight gradient across routed tokens. + + Parallelization: Parallel over M and N dimensions with grouping, loop over K. + """ + # Map program ids to blocks, similar to forward and backward_input. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # Check bounds + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + + # Only process if this block is valid + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + + # Get expert ID for this M block + expert_id = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + + # Only process if expert is valid + if expert_id == -1: + return + + # Load token information for this M block + offs_m = tl.arange(0, BLOCK_SIZE_M) + offs_token_id = pid_m * BLOCK_SIZE_M + offs_m.to(tl.int64) + offs_token = tl.load( + sorted_token_ids_ptr + offs_token_id, + mask=offs_token_id < num_tokens_post_padded, + other=num_valid_tokens, + ) + offs_token = offs_token.to(tl.int64) + token_mask = (offs_token_id < num_tokens_post_padded) & (offs_token < num_valid_tokens) + + # Clamp offs_token to valid range + offs_token_clamped = tl.where(token_mask, offs_token, 0) + + # Determine input token indices based on MUL_ROUTED_WEIGHT + if MUL_ROUTED_WEIGHT: + input_token_idx = offs_token_clamped + input_mask = token_mask + else: + input_token_idx = offs_token_clamped // top_k + num_input_tokens = num_valid_tokens // top_k + input_mask = token_mask & (input_token_idx < num_input_tokens) + + # Load topk_weights if needed + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load(topk_weights_ptr + offs_token_clamped, mask=token_mask, other=0.0) + + # Current N offset for this program + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64) + + # Load grad_output for this N block: shape (M, BLOCK_SIZE_N) + # grad_output is always indexed by sorted_token_ids (offs_token_clamped) + # because it has shape (num_tokens * topk, N) + grad_output_ptrs = grad_output_ptr + ( + offs_token_clamped[:, None] * stride_gom + offs_n[None, :] * stride_gon + ) + grad_out = tl.load( + grad_output_ptrs, mask=token_mask[:, None] & (offs_n[None, :] < N), other=0.0 + ) + + # Apply topk_weights if needed + if MUL_ROUTED_WEIGHT: + grad_out = grad_out * moe_weight[:, None] + + # Zero out padding tokens + token_mask_col = token_mask[:, None] + grad_out = grad_out * token_mask_col + + # Iterate over K blocks and accumulate + for k_block in range(tl.cdiv(K, BLOCK_SIZE_K)): + offs_k = k_block * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64) + + # Load input for this K block + input_ptrs = input_ptr + ( + input_token_idx[:, None] * stride_im + offs_k[None, :] * stride_ik + ) + inp = tl.load(input_ptrs, mask=input_mask[:, None] & (offs_k[None, :] < K), other=0.0) + + # Zero out padding tokens - use input_mask for input, token_mask for grad_output + input_mask_col = input_mask[:, None] + inp = inp * input_mask_col + + # Compute grad_weight contribution: grad_out.T @ inp + grad_w_contribution = tl.dot(grad_out.to(compute_type).T, inp.to(compute_type)) + + # Write back using atomic add + grad_weight_ptrs = ( + grad_weight_ptr + + expert_id * stride_gwe + + offs_n[:, None] * stride_gwn + + offs_k[None, :] * stride_gwk + ) + grad_weight_mask = (offs_n[:, None] < N) & (offs_k[None, :] < K) + tl.atomic_add(grad_weight_ptrs, grad_w_contribution.to(compute_type), mask=grad_weight_mask)