398 lines
16 KiB
Python
398 lines
16 KiB
Python
|
import torch
|
||
|
|
import warnings
|
||
|
|
from wall_x.fusions import backend
|
||
|
|
|
||
|
|
|
||
|
|
class AsymmetricDualExpertGemm(torch.autograd.Function):
|
||
|
|
@staticmethod
|
||
|
|
def forward(ctx, input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False):
|
||
|
|
"""
|
||
|
|
Forward pass for asymmetric dual expert GEMM.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
input_expert0: Expert 0 input [m0, k]
|
||
|
|
input_expert1: Expert 1 input [m1, k]
|
||
|
|
weight_expert0: Expert 0 weight [k, n0] or [n0, k] if trans_b=True
|
||
|
|
weight_expert1: Expert 1 weight [k, n1] or [n1, k] if trans_b=True
|
||
|
|
trans_b: Whether to transpose the weight matrices
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Tuple of (output_expert0, output_expert1)
|
||
|
|
"""
|
||
|
|
# Validate inputs
|
||
|
|
assert input_expert0.dim() == 2, "input_expert0 must be 2D"
|
||
|
|
assert input_expert1.dim() == 2, "input_expert1 must be 2D"
|
||
|
|
assert weight_expert0.dim() == 2, "weight_expert0 must be 2D"
|
||
|
|
assert weight_expert1.dim() == 2, "weight_expert1 must be 2D"
|
||
|
|
|
||
|
|
# Dimension validation depends on trans_b
|
||
|
|
if trans_b:
|
||
|
|
assert input_expert0.size(1) == weight_expert0.size(1), "Expert 0 dimension mismatch (trans_b=True)"
|
||
|
|
assert input_expert1.size(1) == weight_expert1.size(1), "Expert 1 dimension mismatch (trans_b=True)"
|
||
|
|
else:
|
||
|
|
assert input_expert0.size(1) == weight_expert0.size(0), "Expert 0 dimension mismatch (trans_b=False)"
|
||
|
|
assert input_expert1.size(1) == weight_expert1.size(0), "Expert 1 dimension mismatch (trans_b=False)"
|
||
|
|
|
||
|
|
# Save tensors and trans_b for backward pass
|
||
|
|
ctx.save_for_backward(input_expert0, input_expert1, weight_expert0, weight_expert1)
|
||
|
|
ctx.trans_b = trans_b
|
||
|
|
|
||
|
|
# Allocate output tensors
|
||
|
|
m0 = input_expert0.size(0)
|
||
|
|
m1 = input_expert1.size(0)
|
||
|
|
n0 = weight_expert0.size(0) if trans_b else weight_expert0.size(1)
|
||
|
|
n1 = weight_expert1.size(0) if trans_b else weight_expert1.size(1)
|
||
|
|
|
||
|
|
output_expert0 = torch.empty(m0, n0, device=input_expert0.device, dtype=input_expert0.dtype)
|
||
|
|
output_expert1 = torch.empty(m1, n1, device=input_expert1.device, dtype=input_expert1.dtype)
|
||
|
|
|
||
|
|
# Call the backend C++ function
|
||
|
|
backend.asym_dual_gmm_separated(input_expert0, input_expert1, weight_expert0, weight_expert1, output_expert0, output_expert1, trans_b=trans_b)
|
||
|
|
|
||
|
|
return output_expert0, output_expert1
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def backward(ctx, grad_output_expert0, grad_output_expert1):
|
||
|
|
"""
|
||
|
|
Optimized backward pass using specialized kernels.
|
||
|
|
Always computes all gradients to minimize kernel calls.
|
||
|
|
"""
|
||
|
|
grad_output_expert0 = grad_output_expert0.contiguous()
|
||
|
|
grad_output_expert1 = grad_output_expert1.contiguous()
|
||
|
|
|
||
|
|
input_expert0, input_expert1, weight_expert0, weight_expert1 = ctx.saved_tensors
|
||
|
|
trans_b = ctx.trans_b
|
||
|
|
|
||
|
|
# Always allocate all gradient tensors (no conditional computation)
|
||
|
|
grad_input_expert0 = torch.empty_like(input_expert0)
|
||
|
|
grad_input_expert1 = torch.empty_like(input_expert1)
|
||
|
|
grad_weight_expert0 = torch.empty_like(weight_expert0)
|
||
|
|
grad_weight_expert1 = torch.empty_like(weight_expert1)
|
||
|
|
|
||
|
|
# Compute input gradients: grad_input = grad_output @ weight^T (if trans_b=False)
|
||
|
|
# = grad_output @ weight (if trans_b=True)
|
||
|
|
backend.asym_dual_gmm_separated(
|
||
|
|
grad_output_expert0,
|
||
|
|
grad_output_expert1,
|
||
|
|
weight_expert0,
|
||
|
|
weight_expert1,
|
||
|
|
grad_input_expert0,
|
||
|
|
grad_input_expert1,
|
||
|
|
trans_a=False,
|
||
|
|
trans_b=not trans_b,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Compute weight gradients
|
||
|
|
if trans_b:
|
||
|
|
# When trans_b=True in forward: output = input @ weight^T
|
||
|
|
# So grad_weight^T = input^T @ grad_output
|
||
|
|
# Which means grad_weight = grad_output^T @ input
|
||
|
|
backend.asym_dual_gmm_separated(
|
||
|
|
grad_output_expert0,
|
||
|
|
grad_output_expert1,
|
||
|
|
input_expert0,
|
||
|
|
input_expert1,
|
||
|
|
grad_weight_expert0,
|
||
|
|
grad_weight_expert1,
|
||
|
|
trans_a=True,
|
||
|
|
trans_b=False,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
# When trans_b=False in forward: output = input @ weight
|
||
|
|
# So grad_weight = input^T @ grad_output
|
||
|
|
backend.asym_dual_gmm_separated(
|
||
|
|
input_expert0,
|
||
|
|
input_expert1,
|
||
|
|
grad_output_expert0,
|
||
|
|
grad_output_expert1,
|
||
|
|
grad_weight_expert0,
|
||
|
|
grad_weight_expert1,
|
||
|
|
trans_a=True,
|
||
|
|
trans_b=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
return grad_input_expert0, grad_input_expert1, grad_weight_expert0, grad_weight_expert1, None
|
||
|
|
|
||
|
|
|
||
|
|
def asym_dual_gmm(input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b=False):
|
||
|
|
"""
|
||
|
|
Convenience function for asymmetric dual expert GEMM.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
input_expert0: Expert 0 input [m0, k]
|
||
|
|
input_expert1: Expert 1 input [m1, k]
|
||
|
|
weight_expert0: Expert 0 weight [k, n0] or [n0, k] if trans_b=True
|
||
|
|
weight_expert1: Expert 1 weight [k, n1] or [n1, k] if trans_b=True
|
||
|
|
trans_b: Whether to transpose the weight matrices
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Tuple of (output_expert0, output_expert1)
|
||
|
|
"""
|
||
|
|
return AsymmetricDualExpertGemm.apply(input_expert0, input_expert1, weight_expert0, weight_expert1, trans_b)
|
||
|
|
|
||
|
|
|
||
|
|
################################################################################################
|
||
|
|
##
|
||
|
|
## PermuteMoE topK
|
||
|
|
##
|
||
|
|
################################################################################################
|
||
|
|
|
||
|
|
|
||
|
|
class PermuteMoE_topK(torch.autograd.Function):
|
||
|
|
|
||
|
|
workspace_fw = None
|
||
|
|
dtype = None
|
||
|
|
max_expanded_token_num = 0
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def forward(ctx, input_act: torch.Tensor, indices: torch.Tensor, num_out_tokens: int, max_token_num: int):
|
||
|
|
"""
|
||
|
|
indices: for topK=1, indices in a 1-d tensor of shape [num_tokens],
|
||
|
|
otherwise, it's a 2-d tensor of shape [num_tokens, topK]
|
||
|
|
"""
|
||
|
|
# Empty input check
|
||
|
|
if not input_act.numel():
|
||
|
|
return input_act, None
|
||
|
|
|
||
|
|
# For top1 case, view the indices as 2D tensor to unify the shape for topk>=2 cases.
|
||
|
|
if indices.dim() == 1:
|
||
|
|
indices = indices.view(-1, 1)
|
||
|
|
|
||
|
|
# Device check
|
||
|
|
if input_act.is_cpu:
|
||
|
|
raise RuntimeError("[Error] The input `input_act` of permute_topK op is on the device: CPU!")
|
||
|
|
if indices.is_cpu:
|
||
|
|
warnings.warn("The input `indices` of permute_topK op is on the device: CPU!")
|
||
|
|
expert_for_rows = expert_for_rows.cuda()
|
||
|
|
|
||
|
|
# Shape check
|
||
|
|
if input_act.size(0) != indices.size(0):
|
||
|
|
raise RuntimeError(f"[Error] permute_topK op input `indices` shape mismatch! " f"Expect {input_act.size(0)}, but got {indices.size(0)}.")
|
||
|
|
|
||
|
|
# Data type check
|
||
|
|
if indices.dtype != torch.int32:
|
||
|
|
warnings.warn(f"The data type of the input `indices` of permute_topK op is {indices.dtype}! " "The recommended type is torch.int32.")
|
||
|
|
indices = indices.to(torch.int32)
|
||
|
|
|
||
|
|
# Contiguous check
|
||
|
|
if not input_act.is_contiguous():
|
||
|
|
warnings.warn("The input `input_act` of permute_topK op is discontiguous!")
|
||
|
|
input_act = input_act.contiguous()
|
||
|
|
if not indices.is_contiguous():
|
||
|
|
warnings.warn("The input `indices` of permute_topK op is discontiguous!")
|
||
|
|
indices = indices.contiguous()
|
||
|
|
|
||
|
|
num_topK = indices.size(1)
|
||
|
|
|
||
|
|
input_max_expanded_token_num = max(max_token_num, input_act.size(0)) * num_topK
|
||
|
|
if PermuteMoE_topK.max_expanded_token_num < input_max_expanded_token_num:
|
||
|
|
PermuteMoE_topK.max_expanded_token_num = input_max_expanded_token_num
|
||
|
|
PermuteMoE_topK.workspace_fw = []
|
||
|
|
|
||
|
|
if PermuteMoE_topK.dtype != input_act.dtype:
|
||
|
|
PermuteMoE_topK.dtype = input_act.dtype
|
||
|
|
PermuteMoE_topK.workspace_fw = []
|
||
|
|
|
||
|
|
permuted_act, row_id_map, PermuteMoE_topK.workspace_fw = backend.permute(
|
||
|
|
input_act, indices, num_out_tokens, PermuteMoE_topK.workspace_fw, PermuteMoE_topK.max_expanded_token_num
|
||
|
|
)
|
||
|
|
|
||
|
|
ctx.row_id_map = row_id_map
|
||
|
|
ctx.num_tokens = indices.size(0)
|
||
|
|
ctx.num_topK = num_topK
|
||
|
|
return permuted_act, row_id_map
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def backward(ctx, permuted_act_grad, _):
|
||
|
|
# Empty input check
|
||
|
|
if not permuted_act_grad.numel():
|
||
|
|
return permuted_act_grad, None, None, None
|
||
|
|
|
||
|
|
if not permuted_act_grad.is_contiguous():
|
||
|
|
permuted_act_grad = permuted_act_grad.contiguous()
|
||
|
|
|
||
|
|
row_id_map = ctx.row_id_map
|
||
|
|
num_tokens = ctx.num_tokens
|
||
|
|
num_topK = ctx.num_topK
|
||
|
|
|
||
|
|
unpermuted_act_grad = backend.unpermute(permuted_act_grad, row_id_map, torch.tensor([]), num_tokens, num_topK)
|
||
|
|
return unpermuted_act_grad, None, None, None
|
||
|
|
|
||
|
|
|
||
|
|
################################################################################################
|
||
|
|
##
|
||
|
|
## UnpermuteMoE topK
|
||
|
|
##
|
||
|
|
################################################################################################
|
||
|
|
|
||
|
|
|
||
|
|
class UnpermuteMoE_topK(torch.autograd.Function):
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def forward(ctx, input_act: torch.Tensor, row_id_map: torch.Tensor, probs: torch.Tensor = None):
|
||
|
|
# Empty input check
|
||
|
|
if not input_act.numel():
|
||
|
|
ctx.probs = probs
|
||
|
|
return input_act
|
||
|
|
|
||
|
|
# Device check
|
||
|
|
if input_act.is_cpu:
|
||
|
|
raise RuntimeError("[Error] The input `input_act` of unpermute_topK op is on the device: CPU!")
|
||
|
|
if row_id_map.is_cpu:
|
||
|
|
warnings.warn("The input `row_id_map` of unpermute_topK op is on the device: CPU!")
|
||
|
|
row_id_map = row_id_map.cuda()
|
||
|
|
if probs is not None and probs.is_cpu:
|
||
|
|
warnings.warn("The input `probs` of unpermute_topK op is on the device: CPU!")
|
||
|
|
probs = probs.cuda()
|
||
|
|
|
||
|
|
# Shape check
|
||
|
|
if probs is not None and row_id_map.size(0) != probs.size(0) * probs.size(1):
|
||
|
|
raise RuntimeError(
|
||
|
|
f"[Error] unpermute_topK op input `probs` shape mismatch! " f"Expect {row_id_map.size(0)}, but got {probs.size(0) * probs.size(1)}."
|
||
|
|
)
|
||
|
|
|
||
|
|
# Data type check
|
||
|
|
if row_id_map.dtype != torch.int32:
|
||
|
|
warnings.warn(
|
||
|
|
f"The data type of the input `row_id_map` of unpermute_topK op is {row_id_map.dtype}! " "The recommended type is torch.int32."
|
||
|
|
)
|
||
|
|
row_id_map = row_id_map.to(torch.int32)
|
||
|
|
if probs is not None and probs.dtype != torch.float32:
|
||
|
|
warnings.warn(f"The data type of the input `probs` of unpermute_topK op is {probs.dtype}! " "The recommended type is torch.float32.")
|
||
|
|
probs = probs.to(torch.float32)
|
||
|
|
|
||
|
|
# Contiguous check
|
||
|
|
if not input_act.is_contiguous():
|
||
|
|
warnings.warn("The input `input_act` of unpermute_topK op is discontiguous!")
|
||
|
|
input_act = input_act.contiguous()
|
||
|
|
if not row_id_map.is_contiguous():
|
||
|
|
warnings.warn("The input `row_id_map` of unpermute_topK op is discontiguous!")
|
||
|
|
row_id_map = row_id_map.contiguous()
|
||
|
|
if probs is not None and not probs.is_contiguous():
|
||
|
|
warnings.warn("The input `probs` of unpermute_topK op is discontiguous!")
|
||
|
|
probs = probs.contiguous()
|
||
|
|
|
||
|
|
num_tokens = probs.size(0) if probs is not None else input_act.size(0)
|
||
|
|
num_topK = probs.size(1) if probs is not None else 1
|
||
|
|
|
||
|
|
unpermuted_output = backend.unpermute(input_act, row_id_map, probs if probs is not None else torch.tensor([]), num_tokens, num_topK)
|
||
|
|
|
||
|
|
ctx.save_for_backward(input_act, row_id_map, probs)
|
||
|
|
return unpermuted_output
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def backward(ctx, unpermuted_act_grad):
|
||
|
|
# Empty input check
|
||
|
|
if not unpermuted_act_grad.numel():
|
||
|
|
return unpermuted_act_grad, None, ctx.probs
|
||
|
|
|
||
|
|
if not unpermuted_act_grad.is_contiguous():
|
||
|
|
unpermuted_act_grad = unpermuted_act_grad.contiguous()
|
||
|
|
|
||
|
|
input_act, row_id_map, probs = ctx.saved_tensors
|
||
|
|
|
||
|
|
act_grad = None
|
||
|
|
if ctx.needs_input_grad[0]:
|
||
|
|
act_grad, prob_grad = backend.unpermute_bwd(unpermuted_act_grad, input_act, row_id_map, probs)
|
||
|
|
|
||
|
|
if not ctx.needs_input_grad[2]:
|
||
|
|
prob_grad = None
|
||
|
|
return act_grad, None, prob_grad
|
||
|
|
|
||
|
|
|
||
|
|
def permute(input_act, indices, num_out_tokens=None, max_token_num=0):
|
||
|
|
num_out_tokens = 0 if num_out_tokens is None else num_out_tokens
|
||
|
|
return PermuteMoE_topK.apply(input_act, indices, num_out_tokens, max_token_num)
|
||
|
|
|
||
|
|
|
||
|
|
def unpermute(input_act, row_id_map, probs=None):
|
||
|
|
return UnpermuteMoE_topK.apply(input_act, row_id_map, probs)
|
||
|
|
|
||
|
|
|
||
|
|
################################################################################################
|
||
|
|
##
|
||
|
|
## mutlimodal RoPE
|
||
|
|
##
|
||
|
|
################################################################################################
|
||
|
|
|
||
|
|
|
||
|
|
class MultimodalRoPE(torch.autograd.Function):
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def forward(ctx, q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, mrope_section: list):
|
||
|
|
|
||
|
|
# Device check
|
||
|
|
if q.is_cpu:
|
||
|
|
raise RuntimeError("[Error] The input `q` of multimodal_rope op is on the device: CPU!")
|
||
|
|
if k.is_cpu:
|
||
|
|
raise RuntimeError("[Error] The input `k` of multimodal_rope op is on the device: CPU!")
|
||
|
|
if cos.is_cpu:
|
||
|
|
raise RuntimeError("[Error] The input `cos` of multimodal_rope op is on the device: CPU!")
|
||
|
|
if sin.is_cpu:
|
||
|
|
raise RuntimeError("[Error] The input `sin` of multimodal_rope op is on the device: CPU!")
|
||
|
|
if len(mrope_section) != 3:
|
||
|
|
raise RuntimeError("[Error] The input `mrope_section` of multimodal_rope op must be a list of 3 integers!")
|
||
|
|
|
||
|
|
# Contiguous check
|
||
|
|
if not q.is_contiguous():
|
||
|
|
warnings.warn("The input `q` of multimodal_rope op is discontiguous!")
|
||
|
|
q = q.contiguous()
|
||
|
|
if not k.is_contiguous():
|
||
|
|
warnings.warn("The input `k` of multimodal_rope op is discontiguous!")
|
||
|
|
k = k.contiguous()
|
||
|
|
if not cos.is_contiguous():
|
||
|
|
warnings.warn("The input `cos` of multimodal_rope op is discontiguous!")
|
||
|
|
cos = cos.contiguous()
|
||
|
|
if not sin.is_contiguous():
|
||
|
|
warnings.warn("The input `sin` of multimodal_rope op is discontiguous!")
|
||
|
|
sin = sin.contiguous()
|
||
|
|
|
||
|
|
# Prepare mrope_section_doubled
|
||
|
|
mrope_section_doubled = [x * 2 for x in mrope_section]
|
||
|
|
|
||
|
|
# Create output tensors
|
||
|
|
q_out = torch.empty_like(q)
|
||
|
|
k_out = torch.empty_like(k)
|
||
|
|
backend.rope(q, k, cos, sin, q_out, k_out, mrope_section_doubled)
|
||
|
|
|
||
|
|
ctx.save_for_backward(q, k, cos, sin)
|
||
|
|
ctx.mrope_section_doubled = mrope_section_doubled
|
||
|
|
|
||
|
|
return q_out, k_out
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def backward(ctx, grad_q_out, grad_k_out):
|
||
|
|
|
||
|
|
if not grad_q_out.is_contiguous():
|
||
|
|
grad_q_out = grad_q_out.contiguous()
|
||
|
|
if not grad_k_out.is_contiguous():
|
||
|
|
grad_k_out = grad_k_out.contiguous()
|
||
|
|
|
||
|
|
q, k, cos, sin = ctx.saved_tensors
|
||
|
|
|
||
|
|
grad_q = None
|
||
|
|
grad_k = None
|
||
|
|
if ctx.needs_input_grad[0]:
|
||
|
|
grad_q = torch.empty_like(q)
|
||
|
|
if ctx.needs_input_grad[1]:
|
||
|
|
grad_k = torch.empty_like(k)
|
||
|
|
|
||
|
|
if grad_q is not None or grad_k is not None:
|
||
|
|
backend.rope_bwd(
|
||
|
|
grad_q_out,
|
||
|
|
grad_k_out,
|
||
|
|
q,
|
||
|
|
k,
|
||
|
|
cos,
|
||
|
|
sin,
|
||
|
|
grad_q if grad_q is not None else torch.empty_like(q),
|
||
|
|
grad_k if grad_k is not None else torch.empty_like(k),
|
||
|
|
ctx.mrope_section_doubled,
|
||
|
|
)
|
||
|
|
|
||
|
|
return grad_q, grad_k, None, None, None
|
||
|
|
|
||
|
|
|
||
|
|
def multimodal_rope(q, k, cos, sin, mrope_section):
|
||
|
|
return MultimodalRoPE.apply(q, k, cos, sin, mrope_section)
|