huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Init

Browse Source
This commit is contained in:
Starrick
2025-09-07 14:59:17 +08:00
commit 24dbdbd24b
40 changed files with 10754 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/fusions/__init__.py
View File
wall_x/fusions/backend.py
+270
View File
@@ -0,0 +1,270 @@
"""
High-performance C++ backend interface for optimized matrix operations.
This module provides Python bindings for custom CUDA kernels optimized for
transformer and MoE (Mixture of Experts) operations, including:
- Asymmetric dual expert operations
- Token permutation/unpermutation for MoE routing
- RoPE (Rotary Position Embedding) operations
"""
import torch
from typing import Tuple, Optional
import wallx_csrc as backend
def _allocate_asymmetric_dual_outputs(input_expert0: torch.Tensor,
input_expert1: torch.Tensor,
weight_expert0: torch.Tensor,
weight_expert1: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Allocate output tensors for asymmetric dual expert GEMM operations.
This function handles the case where two experts may have different output
dimensions, which is common in heterogeneous MoE architectures.
Args:
input_expert0 (torch.Tensor): Expert 0 input tensor of shape [m0, k]
input_expert1 (torch.Tensor): Expert 1 input tensor of shape [m1, k]
weight_expert0 (torch.Tensor): Expert 0 weight tensor of shape [k, n0]
weight_expert1 (torch.Tensor): Expert 1 weight tensor of shape [k, n1]
Returns:
Tuple[torch.Tensor, torch.Tensor]: Pre-allocated output tensors
- output_expert0: Shape [m0, n0]
- output_expert1: Shape [m1, n1]
Raises:
AssertionError: If tensor dimensions are incompatible
"""
# Validate input tensor dimensions
assert input_expert0.ndim == 2, "Expected 2D tensor for input_expert0"
assert input_expert1.ndim == 2, "Expected 2D tensor for input_expert1"
assert weight_expert0.ndim == 2, "Expected 2D tensor for weight_expert0"
assert weight_expert1.ndim == 2, "Expected 2D tensor for weight_expert1"
# Verify dimension compatibility for matrix multiplication
assert input_expert0.size(1) == weight_expert0.size(0), \
f"Input expert0 K dimension {input_expert0.size(1)} != weight expert0 K dimension {weight_expert0.size(0)}"
assert input_expert1.size(1) == weight_expert1.size(0), \
f"Input expert1 K dimension {input_expert1.size(1)} != weight expert1 K dimension {weight_expert1.size(0)}"
# Calculate output shapes: [m, k] × [k, n] = [m, n]
m0, n0 = input_expert0.size(0), weight_expert0.size(1)
m1, n1 = input_expert1.size(0), weight_expert1.size(1)
# Allocate output tensors with matching device and dtype
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)
return output_expert0, output_expert1
def asym_dual_gmm_separated(input_expert0: torch.Tensor,
input_expert1: torch.Tensor,
weight_expert0: torch.Tensor,
weight_expert1: torch.Tensor,
output_expert0: Optional[torch.Tensor] = None,
output_expert1: Optional[torch.Tensor] = None,
trans_a: bool = False,
trans_b: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Asymmetric dual expert grouped GEMM with separated inputs and outputs.
This is the recommended interface for maximum flexibility and performance when
dealing with two experts that may have different intermediate dimensions.
The operation is equivalent to:
output_expert0 = input_expert0 @ weight_expert0
output_expert1 = input_expert1 @ weight_expert1
But optimized as a single fused kernel call.
Args:
input_expert0 (torch.Tensor): Expert 0 input tensor of shape [m0, k]
input_expert1 (torch.Tensor): Expert 1 input tensor of shape [m1, k]
weight_expert0 (torch.Tensor): Expert 0 weight tensor of shape [k, n0]
weight_expert1 (torch.Tensor): Expert 1 weight tensor of shape [k, n1]
Note: n0 can be different from n1
output_expert0 (torch.Tensor, optional): Pre-allocated output for expert 0 [m0, n0]
output_expert1 (torch.Tensor, optional): Pre-allocated output for expert 1 [m1, n1]
trans_a (bool, optional): Whether to transpose input tensors. Defaults to False.
trans_b (bool, optional): Whether to transpose weight tensors. Defaults to False.
Returns:
Tuple[torch.Tensor, torch.Tensor]: Output tensors (output_expert0, output_expert1)
Example:
>>> # Two experts with different output dimensions
>>> input0 = torch.randn(512, 1024, device='cuda') # 512 tokens for expert 0
>>> input1 = torch.randn(256, 1024, device='cuda') # 256 tokens for expert 1
>>> weight0 = torch.randn(1024, 2048, device='cuda') # Expert 0: 1024->2048
>>> weight1 = torch.randn(1024, 4096, device='cuda') # Expert 1: 1024->4096
>>> out0, out1 = asym_dual_gmm_separated(input0, input1, weight0, weight1)
"""
# Allocate outputs if not provided
if output_expert0 is None or output_expert1 is None:
alloc_out0, alloc_out1 = _allocate_asymmetric_dual_outputs(
input_expert0, input_expert1, weight_expert0, weight_expert1
)
if output_expert0 is None:
output_expert0 = alloc_out0
if output_expert1 is None:
output_expert1 = alloc_out1
# Call optimized C++ backend kernel
backend.asym_dual_gmm(
input_expert0, input_expert1,
weight_expert0, weight_expert1,
output_expert0, output_expert1,
trans_a, trans_b
)
return output_expert0, output_expert1
def permute(input: torch.Tensor,
indices: torch.Tensor,
num_out_tokens: int,
workspace: torch.Tensor,
max_expanded_token_num: int) -> torch.Tensor:
"""
Permute input tokens according to expert assignment indices for MoE routing.
This function reorders tokens based on their assigned experts to enable
efficient grouped processing. Used in the forward pass of MoE layers.
Args:
input (torch.Tensor): Input tokens to permute
indices (torch.Tensor): Expert assignment indices for each token
num_out_tokens (int): Number of output tokens after expansion
workspace (torch.Tensor): Temporary workspace tensor for intermediate computations
max_expanded_token_num (int): Maximum number of tokens after top-k expansion
Returns:
torch.Tensor: Permuted tokens grouped by expert assignment
Note:
This is typically used with top-k expert selection where each token
can be routed to multiple experts.
"""
return backend.permute(input, indices, num_out_tokens, workspace, max_expanded_token_num)
def unpermute(input: torch.Tensor,
row_id_map: torch.Tensor,
prob: torch.Tensor,
max_tokens: int,
num_topK: int) -> torch.Tensor:
"""
Unpermute expert outputs back to original token order with probability weighting.
This function reverses the permutation applied in the forward pass and combines
outputs from multiple experts using their routing probabilities.
Args:
input (torch.Tensor): Permuted expert outputs to unpermute
row_id_map (torch.Tensor): Mapping from permuted positions to original positions
prob (torch.Tensor): Expert routing probabilities for weighted combination
max_tokens (int): Maximum number of tokens in the sequence
num_topK (int): Number of top experts selected per token
Returns:
torch.Tensor: Unpermuted tokens in original order with expert outputs combined
Note:
The output combines multiple expert predictions for each token using
the routing probabilities as weights.
"""
return backend.unpermute(input, row_id_map, prob, max_tokens, num_topK)
def unpermute_bwd(input_bwd: torch.Tensor,
input_fwd: torch.Tensor,
row_id_map: torch.Tensor,
prob: Optional[torch.Tensor]) -> torch.Tensor:
"""
Backward pass for unpermute operation with gradient flow.
This function handles the backward pass through the unpermute operation,
ensuring proper gradient flow for training MoE models.
Args:
input_bwd (torch.Tensor): Backward gradients from the next layer
input_fwd (torch.Tensor): Forward pass inputs (for gradient computation)
row_id_map (torch.Tensor): Row mapping used in forward unpermute
prob (torch.Tensor, optional): Expert probabilities. If None, uniform weights are used.
Returns:
torch.Tensor: Gradients with respect to the input of unpermute forward pass
Note:
If prob is None, uniform probabilities are assumed for gradient computation.
"""
# Handle case where probabilities are not provided
if prob is None:
prob = torch.ones([input_bwd.size(0), 1], dtype=torch.float32, device=input_bwd.device)
return backend.unpermute_bwd(input_bwd, input_fwd, row_id_map, prob)
def rope(q: torch.Tensor,
k: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
q_out: torch.Tensor,
k_out: torch.Tensor,
mrope_section_doubled: bool) -> None:
"""
Apply RoPE (Rotary Position Embedding) to query and key tensors.
Applies rotary position embeddings to query and key tensors using precomputed
cosine and sine values. Supports both standard RoPE and multi-dimensional RoPE (mRoPE).
Args:
q (torch.Tensor): Query tensor to apply RoPE to
k (torch.Tensor): Key tensor to apply RoPE to
cos (torch.Tensor): Precomputed cosine values for rotation
sin (torch.Tensor): Precomputed sine values for rotation
q_out (torch.Tensor): Output tensor for rotated queries (in-place operation supported)
k_out (torch.Tensor): Output tensor for rotated keys (in-place operation supported)
mrope_section_doubled (bool): Whether using multi-dimensional RoPE with doubled sections
Note:
This function performs in-place operations if q_out and k_out point to the same
memory as q and k respectively. The rotation is applied using the standard
RoPE formulation with complex number rotation.
"""
return backend.rope(q, k, cos, sin, q_out, k_out, mrope_section_doubled)
def rope_bwd(grad_q_out: torch.Tensor,
grad_k_out: torch.Tensor,
q: torch.Tensor,
k: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
grad_q: torch.Tensor,
grad_k: torch.Tensor,
mrope_section_doubled: bool) -> None:
"""
Backward pass for RoPE operation with gradient computation.
Computes gradients with respect to the input query and key tensors
for the RoPE operation used in transformer attention mechanisms.
Args:
grad_q_out (torch.Tensor): Gradient with respect to output queries
grad_k_out (torch.Tensor): Gradient with respect to output keys
q (torch.Tensor): Original query tensor from forward pass
k (torch.Tensor): Original key tensor from forward pass
cos (torch.Tensor): Cosine values used in forward pass
sin (torch.Tensor): Sine values used in forward pass
grad_q (torch.Tensor): Output tensor for query gradients
grad_k (torch.Tensor): Output tensor for key gradients
mrope_section_doubled (bool): Whether using multi-dimensional RoPE configuration
Note:
This function computes the analytical gradient of the RoPE operation,
which involves the inverse rotation compared to the forward pass.
"""
return backend.rope_bwd(grad_q_out, grad_k_out, q, k, cos, sin, grad_q, grad_k, mrope_section_doubled)
wall_x/fusions/ops.py
+397
View File
@@ -0,0 +1,397 @@
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)