Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,17 @@
|
||||
"""Operator proxy layer with runtime fallback backends."""
|
||||
|
||||
from wall_x.model.core.ops.index import get_rope_index, get_window_index
|
||||
from wall_x.model.core.ops.moe import permute, unpermute
|
||||
from wall_x.model.core.ops.norm import rmsnorm
|
||||
from wall_x.model.core.ops.rope import m_rope, rope, rot_pos_emb
|
||||
|
||||
__all__ = [
|
||||
"rmsnorm",
|
||||
"rope",
|
||||
"m_rope",
|
||||
"rot_pos_emb",
|
||||
"permute",
|
||||
"unpermute",
|
||||
"get_rope_index",
|
||||
"get_window_index",
|
||||
]
|
||||
@@ -0,0 +1,35 @@
|
||||
"""Install-time compiled CUDA kernels for wall_x ops.
|
||||
|
||||
Auto-generated by scripts/extract_cuda_kernels.py.
|
||||
Do not edit manually.
|
||||
"""
|
||||
|
||||
_module = None
|
||||
|
||||
|
||||
def load():
|
||||
"""Load the CUDA extension module built by setup.py."""
|
||||
global _module
|
||||
if _module is not None:
|
||||
return _module
|
||||
|
||||
try:
|
||||
from wall_x.model.core.ops import _cuda_ext_bin
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Wall-X CUDA operators were not built. Install requirements first, "
|
||||
"then reinstall Wall-X with: "
|
||||
"MAX_JOBS=8 pip install --no-build-isolation -e ."
|
||||
) from exc
|
||||
|
||||
_module = _cuda_ext_bin
|
||||
return _module
|
||||
|
||||
|
||||
def is_available() -> bool:
|
||||
"""Check whether the install-time CUDA extension can be imported."""
|
||||
try:
|
||||
load()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
@@ -0,0 +1,750 @@
|
||||
"""Auto-generated CUDA kernel wrappers with autograd support.
|
||||
|
||||
Generated by scripts/extract_cuda_kernels.py from internal CUDA wrappers.
|
||||
Do not edit manually.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch.autograd import Function
|
||||
|
||||
|
||||
def _m():
|
||||
"""Get the compiled CUDA extension module."""
|
||||
from wall_x.model.core.ops._cuda_ext import load
|
||||
return load()
|
||||
|
||||
|
||||
|
||||
class _RopeFunction(Function):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
q,
|
||||
k,
|
||||
cos,
|
||||
sin,
|
||||
interleave,
|
||||
fwd_kernel=None,
|
||||
bwd_kernel=None,
|
||||
):
|
||||
q_embed = q
|
||||
k_embed = k
|
||||
ctx.save_for_backward(q, k, cos, sin)
|
||||
ctx.interleave = interleave
|
||||
ctx.bwd_kernel = bwd_kernel
|
||||
fwd_kernel(q, k, q_embed, k_embed, cos, sin, interleave)
|
||||
return q_embed, k_embed
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_q_embed, grad_k_embed):
|
||||
q, k, cos, sin = ctx.saved_tensors
|
||||
interleave = ctx.interleave
|
||||
bwd_kernel = ctx.bwd_kernel
|
||||
grad_q = grad_q_embed
|
||||
grad_k = grad_k_embed
|
||||
|
||||
bwd_kernel(
|
||||
grad_q_embed,
|
||||
grad_k_embed,
|
||||
grad_q, # output
|
||||
grad_k, # output
|
||||
cos,
|
||||
sin,
|
||||
interleave,
|
||||
)
|
||||
|
||||
return grad_q, grad_k, None, None, None, None, None
|
||||
|
||||
|
||||
class _RopePackFunction(Function):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
qkv,
|
||||
q_num_heads,
|
||||
kv_num_heads,
|
||||
cos,
|
||||
sin,
|
||||
interleave,
|
||||
fwd_kernel=None,
|
||||
bwd_kernel=None,
|
||||
):
|
||||
ctx.save_for_backward(cos, sin)
|
||||
ctx.q_num_heads = q_num_heads
|
||||
ctx.kv_num_heads = kv_num_heads
|
||||
ctx.interleave = interleave
|
||||
ctx.bwd_kernel = bwd_kernel
|
||||
|
||||
fwd_kernel(qkv, cos, sin, q_num_heads, kv_num_heads, interleave)
|
||||
return qkv
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dqkv):
|
||||
if dqkv is None:
|
||||
return None, None, None, None, None, None, None, None
|
||||
|
||||
cos, sin = ctx.saved_tensors
|
||||
if dqkv.stride(-1) != 1:
|
||||
dqkv = dqkv.contiguous()
|
||||
|
||||
ctx.bwd_kernel(
|
||||
dqkv,
|
||||
cos,
|
||||
sin,
|
||||
ctx.q_num_heads,
|
||||
ctx.kv_num_heads,
|
||||
ctx.interleave,
|
||||
)
|
||||
return dqkv, None, None, None, None, None, None, None
|
||||
|
||||
|
||||
class Rope:
|
||||
def __init__(self):
|
||||
pass
|
||||
self.fwd_kernel = _m().rope
|
||||
self.inplace_kernel = _m().rope_inplace
|
||||
self.pack_kernel = _m().rope_inplace_pack
|
||||
self.pack_bwd_kernel = _m().rope_inplace_pack_bwd
|
||||
self.bwd_kernel = _m().rope_bwd
|
||||
|
||||
def pack(
|
||||
self,
|
||||
qkv,
|
||||
q_num_heads,
|
||||
kv_num_heads,
|
||||
cos,
|
||||
sin,
|
||||
interleave=False,
|
||||
inference=False,
|
||||
):
|
||||
"""packed qkvq/krope,qkv.
|
||||
|
||||
Split/view/offset calculations are done on the C++ side. Callers consume qkv through the paired
|
||||
flash_attn pack interface without splitting q/k/v in Python.
|
||||
autograd,backwarddqkvdq/dk.
|
||||
|
||||
Args:
|
||||
qkv: [seq_len, q_dim + 2*kv_dim]
|
||||
q_num_heads: number of query heads
|
||||
kv_num_heads: number of key/value heads
|
||||
cos, sin: rotary position encoding [1, seq_len, head_dim/2]
|
||||
interleave: whether to use interleaved RoPE
|
||||
inference: kernel,autograd
|
||||
"""
|
||||
if inference:
|
||||
self.pack_kernel(qkv, cos, sin, q_num_heads, kv_num_heads, interleave)
|
||||
return qkv
|
||||
|
||||
return _RopePackFunction.apply(
|
||||
qkv,
|
||||
q_num_heads,
|
||||
kv_num_heads,
|
||||
cos,
|
||||
sin,
|
||||
interleave,
|
||||
self.pack_kernel,
|
||||
self.pack_bwd_kernel,
|
||||
)
|
||||
|
||||
def pack_backward(
|
||||
self, dqkv, q_num_heads, kv_num_heads, cos, sin, interleave=False
|
||||
):
|
||||
"""packed dqkvdq/dkrope backward,dqkv."""
|
||||
self.pack_bwd_kernel(dqkv, cos, sin, q_num_heads, kv_num_heads, interleave)
|
||||
return dqkv
|
||||
|
||||
def __call__(self, q, k, cos, sin, interleave=False, inference=False):
|
||||
if inference:
|
||||
self.inplace_kernel(q, k, cos, sin, interleave)
|
||||
return q, k
|
||||
|
||||
return _RopeFunction.apply(
|
||||
q,
|
||||
k,
|
||||
cos,
|
||||
sin,
|
||||
interleave,
|
||||
self.fwd_kernel,
|
||||
self.bwd_kernel,
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
class _MRopeFunction(Function):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
q,
|
||||
k,
|
||||
cos,
|
||||
sin,
|
||||
mrope_section,
|
||||
fwd_kernel=None,
|
||||
bwd_kernel=None,
|
||||
):
|
||||
first = mrope_section[0]
|
||||
second = mrope_section[1]
|
||||
ctx.save_for_backward(q, k, cos, sin)
|
||||
ctx.mrope_section = mrope_section
|
||||
ctx.bwd_kernel = bwd_kernel
|
||||
q_embed = q
|
||||
k_embed = k
|
||||
fwd_kernel(q, k, q_embed, k_embed, cos, sin, first, second)
|
||||
return q_embed, k_embed
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_q_embed, grad_k_embed):
|
||||
if grad_q_embed.stride(-1) != 1:
|
||||
grad_q_embed = grad_q_embed.contiguous()
|
||||
if grad_k_embed.stride(-1) != 1:
|
||||
grad_k_embed = grad_k_embed.contiguous()
|
||||
q, k, cos, sin = ctx.saved_tensors
|
||||
mrope_section = ctx.mrope_section
|
||||
bwd_kernel = ctx.bwd_kernel
|
||||
grad_q = grad_q_embed
|
||||
grad_k = grad_k_embed
|
||||
first = mrope_section[0]
|
||||
second = mrope_section[1]
|
||||
|
||||
bwd_kernel(
|
||||
grad_q_embed,
|
||||
grad_k_embed,
|
||||
grad_q, # output
|
||||
grad_k, # output
|
||||
cos,
|
||||
sin,
|
||||
first,
|
||||
second,
|
||||
)
|
||||
|
||||
return grad_q, grad_k, None, None, None, None, None
|
||||
|
||||
|
||||
class _MRopePackFunction(Function):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
qkv,
|
||||
q_num_heads,
|
||||
kv_num_heads,
|
||||
cos,
|
||||
sin,
|
||||
mrope_section,
|
||||
fwd_kernel=None,
|
||||
bwd_kernel=None,
|
||||
):
|
||||
first = mrope_section[0]
|
||||
second = mrope_section[1]
|
||||
ctx.save_for_backward(cos, sin)
|
||||
ctx.q_num_heads = q_num_heads
|
||||
ctx.kv_num_heads = kv_num_heads
|
||||
ctx.mrope_section = mrope_section
|
||||
ctx.bwd_kernel = bwd_kernel
|
||||
|
||||
fwd_kernel(qkv, cos, sin, q_num_heads, kv_num_heads, first, second)
|
||||
return qkv
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dqkv):
|
||||
if dqkv is None:
|
||||
return None, None, None, None, None, None, None, None
|
||||
|
||||
cos, sin = ctx.saved_tensors
|
||||
if dqkv.stride(-1) != 1:
|
||||
dqkv = dqkv.contiguous()
|
||||
first = ctx.mrope_section[0]
|
||||
second = ctx.mrope_section[1]
|
||||
|
||||
ctx.bwd_kernel(
|
||||
dqkv,
|
||||
cos,
|
||||
sin,
|
||||
ctx.q_num_heads,
|
||||
ctx.kv_num_heads,
|
||||
first,
|
||||
second,
|
||||
)
|
||||
return dqkv, None, None, None, None, None, None, None
|
||||
|
||||
|
||||
class MRope:
|
||||
def __init__(self):
|
||||
pass
|
||||
self.fwd_kernel = _m().m_rope
|
||||
self.inplace_kernel = _m().m_rope_inplace
|
||||
self.pack_kernel = _m().m_rope_inplace_pack
|
||||
self.pack_bwd_kernel = _m().m_rope_inplace_pack_bwd
|
||||
self.bwd_kernel = _m().m_rope_bwd
|
||||
|
||||
def pack(
|
||||
self,
|
||||
qkv,
|
||||
q_num_heads,
|
||||
kv_num_heads,
|
||||
cos,
|
||||
sin,
|
||||
mrope_section,
|
||||
inference=False,
|
||||
):
|
||||
"""packed qkvq/kmrope,qkv.GQA.
|
||||
|
||||
Split/view/offset calculations are done on the C++ side. qkv must be 3D: [bz, seq_len, q_dim+2*kv_dim].
|
||||
|
||||
Args:
|
||||
qkv: [bz, seq_len, q_dim + 2*kv_dim]
|
||||
q_num_heads: number of query heads
|
||||
kv_num_heads: number of key/value heads
|
||||
cos, sin: rotary position encoding [3, bz, seq_len, head_dim/2]
|
||||
mrope_section: (first, second) M-RoPE section parameters
|
||||
inference: kernel,autograd
|
||||
"""
|
||||
if inference:
|
||||
self.pack_kernel(
|
||||
qkv,
|
||||
cos,
|
||||
sin,
|
||||
q_num_heads,
|
||||
kv_num_heads,
|
||||
mrope_section[0],
|
||||
mrope_section[1],
|
||||
)
|
||||
return qkv
|
||||
|
||||
return _MRopePackFunction.apply(
|
||||
qkv,
|
||||
q_num_heads,
|
||||
kv_num_heads,
|
||||
cos,
|
||||
sin,
|
||||
mrope_section,
|
||||
self.pack_kernel,
|
||||
self.pack_bwd_kernel,
|
||||
)
|
||||
|
||||
def pack_backward(self, dqkv, q_num_heads, kv_num_heads, cos, sin, mrope_section):
|
||||
"""packed dqkvdq/dkmrope backward,dqkv."""
|
||||
self.pack_bwd_kernel(
|
||||
dqkv,
|
||||
cos,
|
||||
sin,
|
||||
q_num_heads,
|
||||
kv_num_heads,
|
||||
mrope_section[0],
|
||||
mrope_section[1],
|
||||
)
|
||||
return dqkv
|
||||
|
||||
def __call__(self, q, k, cos, sin, mrope_section, inference=False):
|
||||
if inference:
|
||||
first = mrope_section[0]
|
||||
second = mrope_section[1]
|
||||
self.inplace_kernel(q, k, cos, sin, first, second)
|
||||
return q, k
|
||||
|
||||
return _MRopeFunction.apply(
|
||||
q,
|
||||
k,
|
||||
cos,
|
||||
sin,
|
||||
mrope_section,
|
||||
self.fwd_kernel,
|
||||
self.bwd_kernel,
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
class RotPos:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, inv_freq, grid_thw, spatial_merge_size):
|
||||
assert (
|
||||
inv_freq.dtype == torch.float32
|
||||
), f"Expected float32, got {inv_freq.dtype}"
|
||||
get_token_counts_kernel = _m().get_token_counts
|
||||
rot_pos_kernel = _m().rot_pos
|
||||
|
||||
num_grids = grid_thw.size(0)
|
||||
token_counts = torch.zeros(
|
||||
(num_grids), dtype=grid_thw.dtype, device=grid_thw.device
|
||||
)
|
||||
|
||||
get_token_counts_kernel(grid_thw, token_counts, spatial_merge_size)
|
||||
cumsum_tokens = torch.cat(
|
||||
[
|
||||
torch.zeros(1, dtype=token_counts.dtype, device=token_counts.device),
|
||||
token_counts.cumsum(dim=0),
|
||||
],
|
||||
dim=0,
|
||||
).to(grid_thw.dtype)
|
||||
|
||||
output = torch.empty(
|
||||
(cumsum_tokens[-1].item(), inv_freq.size(0) * 2),
|
||||
dtype=torch.float,
|
||||
device=inv_freq.device,
|
||||
)
|
||||
rot_pos_kernel(inv_freq, grid_thw, output, cumsum_tokens, spatial_merge_size)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
|
||||
|
||||
class GetRopeIndex:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
input_ids,
|
||||
image_grid_thw,
|
||||
video_grid_thw,
|
||||
second_per_grid_ts,
|
||||
attention_mask,
|
||||
spatial_merge_size,
|
||||
image_token_id,
|
||||
video_token_id,
|
||||
vision_start_token_id,
|
||||
tokens_per_second,
|
||||
):
|
||||
get_workspace = _m().get_rope_index_getworkspace
|
||||
get_rope_index_kernel = _m().get_rope_index
|
||||
work_space_size = get_workspace(input_ids, image_grid_thw, video_grid_thw)
|
||||
|
||||
workspace = torch.empty(
|
||||
work_space_size, dtype=torch.uint8, device=input_ids.device
|
||||
)
|
||||
batch_size = input_ids.size(0)
|
||||
seq_len = input_ids.size(1)
|
||||
position_ids = torch.empty(
|
||||
(3, batch_size, seq_len), dtype=torch.int64, device=input_ids.device
|
||||
)
|
||||
mrope_deltas = torch.empty(
|
||||
(batch_size, 1), dtype=torch.int64, device=input_ids.device
|
||||
)
|
||||
get_rope_index_kernel(
|
||||
input_ids,
|
||||
image_grid_thw,
|
||||
video_grid_thw,
|
||||
second_per_grid_ts,
|
||||
attention_mask,
|
||||
position_ids,
|
||||
mrope_deltas,
|
||||
workspace,
|
||||
spatial_merge_size,
|
||||
image_token_id,
|
||||
video_token_id,
|
||||
vision_start_token_id,
|
||||
tokens_per_second,
|
||||
)
|
||||
|
||||
return position_ids, mrope_deltas
|
||||
|
||||
|
||||
|
||||
get_window_index_kernel = _m().get_window_index
|
||||
get_totals_kernel = _m().get_totals
|
||||
|
||||
|
||||
def get_window_index_cuda(
|
||||
grid_thw, window_size, spatial_merge_size, patch_size, spatial_merge_unit=1
|
||||
):
|
||||
|
||||
if grid_thw.size(0) == 0:
|
||||
return (
|
||||
torch.empty(0, dtype=grid_thw.dtype, device=grid_thw.device),
|
||||
torch.zeros(1, dtype=grid_thw.dtype, device=grid_thw.device),
|
||||
)
|
||||
|
||||
vit_merger_window_size = window_size // spatial_merge_size // patch_size
|
||||
|
||||
grid_info_tensor = torch.empty(
|
||||
(grid_thw.size(0), 6), dtype=grid_thw.dtype, device=grid_thw.device
|
||||
)
|
||||
global_totals_tensor = torch.zeros(
|
||||
(2), dtype=grid_thw.dtype, device=grid_thw.device
|
||||
)
|
||||
get_totals_kernel(
|
||||
grid_thw,
|
||||
grid_info_tensor,
|
||||
global_totals_tensor,
|
||||
spatial_merge_size,
|
||||
vit_merger_window_size,
|
||||
)
|
||||
total_elements = global_totals_tensor[0].item()
|
||||
total_windows = global_totals_tensor[1].item()
|
||||
if total_elements == 0 or total_windows == 0:
|
||||
return (
|
||||
torch.empty(0, dtype=grid_thw.dtype, device=grid_thw.device),
|
||||
torch.zeros(1, dtype=grid_thw.dtype, device=grid_thw.device),
|
||||
)
|
||||
|
||||
window_indices = torch.empty(
|
||||
total_elements, dtype=grid_thw.dtype, device=grid_thw.device
|
||||
)
|
||||
cu_window_seqlens = torch.empty(
|
||||
(total_windows + 1), dtype=grid_thw.dtype, device=grid_thw.device
|
||||
)
|
||||
window_counts_tensor = torch.empty(
|
||||
(total_windows), dtype=grid_thw.dtype, device=grid_thw.device
|
||||
)
|
||||
|
||||
max_grid_t = grid_thw[:, 0].max().item()
|
||||
get_window_index_kernel(
|
||||
grid_thw,
|
||||
grid_info_tensor,
|
||||
window_indices,
|
||||
cu_window_seqlens,
|
||||
window_counts_tensor,
|
||||
max_grid_t,
|
||||
spatial_merge_size,
|
||||
vit_merger_window_size,
|
||||
patch_size,
|
||||
spatial_merge_unit,
|
||||
)
|
||||
|
||||
return window_indices, cu_window_seqlens
|
||||
|
||||
|
||||
|
||||
################################################################################################
|
||||
##
|
||||
## PermuteMoE topK
|
||||
##
|
||||
################################################################################################
|
||||
|
||||
|
||||
|
||||
class PermuteMoE_topK(torch.autograd.Function):
|
||||
|
||||
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]
|
||||
"""
|
||||
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!"
|
||||
)
|
||||
|
||||
# Data type check
|
||||
if indices.dtype != torch.int32:
|
||||
indices = indices.to(torch.int32)
|
||||
|
||||
# Contiguous check
|
||||
if not input_act.is_contiguous():
|
||||
input_act = input_act.contiguous()
|
||||
if not indices.is_contiguous():
|
||||
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
|
||||
|
||||
permute_kernel = _m().permute
|
||||
sorted_indices = torch.empty(
|
||||
PermuteMoE_topK.max_expanded_token_num,
|
||||
dtype=torch.int32,
|
||||
device=input_act.device,
|
||||
)
|
||||
row_id = torch.arange(
|
||||
PermuteMoE_topK.max_expanded_token_num,
|
||||
dtype=torch.int32,
|
||||
device=input_act.device,
|
||||
)
|
||||
sorted_row_id = torch.empty(
|
||||
PermuteMoE_topK.max_expanded_token_num,
|
||||
dtype=torch.int32,
|
||||
device=input_act.device,
|
||||
)
|
||||
get_storage_bytes_kernel = _m().cub_sort_pair_get_storage_bytes
|
||||
temp_storage_bytes = get_storage_bytes_kernel(
|
||||
PermuteMoE_topK.max_expanded_token_num
|
||||
)
|
||||
temp_storage = torch.empty(
|
||||
temp_storage_bytes, dtype=torch.int8, device=input_act.device
|
||||
)
|
||||
num_out = (
|
||||
num_out_tokens if (num_out_tokens > 0) else (indices.size(0) * num_topK)
|
||||
)
|
||||
permuted_output = torch.empty(
|
||||
(num_out, input_act.size(1)), dtype=input_act.dtype, device=input_act.device
|
||||
)
|
||||
row_id_map = torch.empty(
|
||||
(indices.size(0) * num_topK), dtype=torch.int32, device=input_act.device
|
||||
)
|
||||
permute_kernel(
|
||||
input_act,
|
||||
indices,
|
||||
sorted_indices,
|
||||
row_id,
|
||||
sorted_row_id,
|
||||
temp_storage,
|
||||
permuted_output,
|
||||
row_id_map,
|
||||
num_out_tokens,
|
||||
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_output, row_id_map
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, permuted_act_grad, _):
|
||||
if not permuted_act_grad.numel():
|
||||
return permuted_act_grad, None, None, None
|
||||
|
||||
unpermute_kernel = _m().unpermute
|
||||
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
|
||||
num_cols = permuted_act_grad.size(1)
|
||||
unpermuted_output = torch.empty(
|
||||
(num_tokens, num_cols),
|
||||
dtype=permuted_act_grad.dtype,
|
||||
device=permuted_act_grad.device,
|
||||
)
|
||||
unpermute_kernel(
|
||||
permuted_act_grad, row_id_map, None, unpermuted_output, num_tokens, num_topK
|
||||
)
|
||||
|
||||
return unpermuted_output, 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,
|
||||
):
|
||||
|
||||
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:
|
||||
row_id_map = row_id_map.cuda()
|
||||
if probs is not None and probs.is_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:
|
||||
row_id_map = row_id_map.to(torch.int32)
|
||||
if probs is not None and probs.dtype != torch.float32:
|
||||
probs = probs.to(torch.float32)
|
||||
|
||||
# Contiguous check
|
||||
if not input_act.is_contiguous():
|
||||
input_act = input_act.contiguous()
|
||||
if not row_id_map.is_contiguous():
|
||||
row_id_map = row_id_map.contiguous()
|
||||
if probs is not None and not probs.is_contiguous():
|
||||
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
|
||||
unpermute_kernel = _m().unpermute
|
||||
num_cols = input_act.size(1)
|
||||
unpermuted_output = torch.empty(
|
||||
(num_tokens, num_cols), dtype=input_act.dtype, device=input_act.device
|
||||
)
|
||||
unpermute_kernel(
|
||||
input_act, row_id_map, probs, unpermuted_output, num_tokens, num_topK
|
||||
)
|
||||
|
||||
ctx.save_for_backward(input_act, row_id_map, probs)
|
||||
|
||||
return unpermuted_output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, unpermuted_act_grad):
|
||||
|
||||
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
|
||||
unpermute_bwd_kernel = _m().unpermute_bwd
|
||||
if ctx.needs_input_grad[0]:
|
||||
num_cols = unpermuted_act_grad.size(1)
|
||||
act_grad = torch.empty(
|
||||
(input_act.size(0), num_cols),
|
||||
dtype=unpermuted_act_grad.dtype,
|
||||
device=unpermuted_act_grad.device,
|
||||
)
|
||||
prob_grad = torch.empty(
|
||||
(probs.size(0), probs.size(1)),
|
||||
dtype=torch.float32,
|
||||
device=unpermuted_act_grad.device,
|
||||
)
|
||||
unpermute_bwd_kernel(
|
||||
unpermuted_act_grad, input_act, row_id_map, probs, act_grad, prob_grad
|
||||
)
|
||||
|
||||
if not ctx.needs_input_grad[2]:
|
||||
prob_grad = None
|
||||
|
||||
return act_grad, None, prob_grad
|
||||
|
||||
|
||||
def permute_kernel(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_kernel(input_act, row_id_map, probs=None):
|
||||
return UnpermuteMoE_topK.apply(input_act, row_id_map, probs)
|
||||
@@ -0,0 +1,143 @@
|
||||
"""Base class for operator proxies with lazy backend resolution."""
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
import threading
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class OpsProxy:
|
||||
"""Callable proxy that lazily resolves to the best available backend.
|
||||
|
||||
Three-level fallback:
|
||||
Level 1: external_accel (fastest, requires external_accel package)
|
||||
Level 2: CUDA inline kernel (JIT compiled, requires CUDA)
|
||||
Level 3: Pure PyTorch (always available)
|
||||
|
||||
Subclasses define _external_accel_name, _pytorch_fallback, and optionally _cuda_kernel.
|
||||
|
||||
Set _external_accel_name to None to skip external_accel resolution (e.g., when the
|
||||
external_accel API is incompatible).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._resolved_fn = None
|
||||
self._backend = None
|
||||
self._resolve_lock = threading.Lock()
|
||||
|
||||
def _resolve(self):
|
||||
"""Resolve the best available backend. Called once on first use (thread-safe)."""
|
||||
if self._resolved_fn is not None:
|
||||
return
|
||||
with self._resolve_lock:
|
||||
# Double-check after acquiring lock
|
||||
if self._resolved_fn is not None:
|
||||
return
|
||||
|
||||
# Level 1: try external_accel (skip if _external_accel_name is None)
|
||||
fn = self._probe_external_accel()
|
||||
if fn is not None:
|
||||
self._backend = "external_accel"
|
||||
self._resolved_fn = fn
|
||||
return
|
||||
|
||||
# Level 2: try CUDA inline kernel (subclass override)
|
||||
cuda_fn = self._get_cuda_kernel()
|
||||
if cuda_fn is not None:
|
||||
self._backend = "cuda_inline"
|
||||
self._resolved_fn = cuda_fn
|
||||
return
|
||||
|
||||
# Level 3: PyTorch fallback
|
||||
self._backend = "pytorch"
|
||||
self._resolved_fn = self._pytorch_fallback
|
||||
logger.info(f"{self.__class__.__name__}: using PyTorch fallback")
|
||||
|
||||
def _get_cuda_kernel(self):
|
||||
"""Override in subclass to provide Level 2 CUDA kernel. Returns None if unavailable."""
|
||||
return None
|
||||
|
||||
@property
|
||||
def _external_accel_name(self):
|
||||
"""Return external_accel attribute name, or None to skip external_accel resolution.
|
||||
|
||||
Base returns None (skip external_accel). Subclasses override to return the
|
||||
attribute name when external_accel integration is desired.
|
||||
"""
|
||||
return None
|
||||
|
||||
def _pytorch_fallback(self, *args, **kwargs):
|
||||
raise NotImplementedError(f"{self.__class__.__name__} has no PyTorch fallback")
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
if self._resolved_fn is None:
|
||||
self._resolve()
|
||||
return self._resolved_fn(*args, **kwargs)
|
||||
|
||||
@property
|
||||
def backend(self) -> str:
|
||||
if self._resolved_fn is None:
|
||||
self._resolve()
|
||||
return self._backend
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Multi-backend API (for testing / benchmarking)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _probe_external_accel(self):
|
||||
"""Return the external_accel function for this op, or None if unavailable."""
|
||||
if self._external_accel_name is None:
|
||||
return None
|
||||
try:
|
||||
mod = importlib.import_module("external_accel")
|
||||
return getattr(mod, self._external_accel_name)
|
||||
except (ImportError, AttributeError):
|
||||
return None
|
||||
|
||||
def available_backends(self):
|
||||
"""Return list of available backend names for this operator."""
|
||||
backends = ["pytorch"]
|
||||
if self._probe_external_accel() is not None:
|
||||
backends.append("external_accel")
|
||||
if self._get_cuda_kernel() is not None:
|
||||
backends.append("cuda_inline")
|
||||
return backends
|
||||
|
||||
def _get_backend_fn(self, backend):
|
||||
"""Return the callable for a specific backend, or None if unavailable."""
|
||||
if backend == "pytorch":
|
||||
return self._pytorch_fallback
|
||||
elif backend == "external_accel":
|
||||
return self._probe_external_accel()
|
||||
elif backend == "cuda_inline":
|
||||
return self._get_cuda_kernel()
|
||||
else:
|
||||
raise ValueError(f"Unknown backend: {backend}")
|
||||
|
||||
def call_with_backend(self, backend, *args, **kwargs):
|
||||
"""Call this operator using a specific backend.
|
||||
|
||||
This does NOT change the default backend used by ``__call__``.
|
||||
The ``backend`` property still reflects the auto-resolved default.
|
||||
|
||||
Args:
|
||||
backend: One of "external_accel", "cuda_inline", "pytorch".
|
||||
|
||||
Returns:
|
||||
Operator output from the specified backend.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the requested backend is not available.
|
||||
"""
|
||||
fn = self._get_backend_fn(backend)
|
||||
if fn is None:
|
||||
raise RuntimeError(
|
||||
f"{self.__class__.__name__}: backend '{backend}' not available. "
|
||||
f"Available: {self.available_backends()}"
|
||||
)
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
def __repr__(self):
|
||||
backend = self._backend or "unresolved"
|
||||
return f"<{self.__class__.__name__} backend={backend}>"
|
||||
@@ -0,0 +1,69 @@
|
||||
// Auto-generated by scripts/extract_cuda_kernels.py
|
||||
// This file registers CUDA kernel bindings via pybind11.
|
||||
//
|
||||
// Kernel .cu files are compiled separately so ninja can parallelize builds.
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <torch/extension.h>
|
||||
|
||||
namespace wallx_cuda_get_rope_index {
|
||||
void GetRopeIndex(const c10::optional<at::Tensor>& input_ids, const c10::optional<at::Tensor>& image_grid_thw, const c10::optional<at::Tensor>& video_grid_thw, const c10::optional<at::Tensor>& second_per_grid_ts, const c10::optional<at::Tensor>& attention_mask, const at::Tensor& position_ids, const at::Tensor& mrope_deltas, const at::Tensor& workspace, int spatial_merge_size, int image_token_id, int video_token_id, int vision_start_token_id, float tokens_per_second);
|
||||
int64_t GetRopeIndexGetWorkSpace(const c10::optional<at::Tensor>& input_ids, const c10::optional<at::Tensor>& image_grid_thw, const c10::optional<at::Tensor>& video_grid_thw);
|
||||
} // namespace wallx_cuda_get_rope_index
|
||||
|
||||
namespace wallx_cuda_m_rope {
|
||||
void MRope(const at::Tensor& q, const at::Tensor& k, const at::Tensor& q_embed, const at::Tensor& k_embed, const at::Tensor& cos, const at::Tensor& sin, const int first, const int second );
|
||||
void MRopeInplace(const at::Tensor& q, const at::Tensor& k, const at::Tensor& cos, const at::Tensor& sin, const int first, const int second );
|
||||
void MRopeInplacePack(const at::Tensor& qkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, const int first, const int second);
|
||||
void MRopeInplacePackBackward(const at::Tensor& dqkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, const int first, const int second);
|
||||
void MRopeBackward( const at::Tensor& grad_q_embed, const at::Tensor& grad_k_embed, const at::Tensor& grad_q, const at::Tensor& grad_k, const at::Tensor& cos, const at::Tensor& sin, const int first, const int second );
|
||||
} // namespace wallx_cuda_m_rope
|
||||
|
||||
namespace wallx_cuda_permute {
|
||||
void MoePermuteTopKOp( const at::Tensor& input, const at::Tensor& indices, const at::Tensor& sorted_indices, const at::Tensor& row_id, const at::Tensor& sorted_row_id, const at::Tensor& temp_storage, const at::Tensor& permuted_output, const at::Tensor& row_id_map, int64_t num_out_tokens, int64_t max_expanded_token_num );
|
||||
void MoeRecoverTopKOp( const at::Tensor& input, const at::Tensor& row_id_map, const c10::optional<at::Tensor>& prob, const at::Tensor& unpermuted_output, int64_t num_tokens, int64_t num_topK );
|
||||
void MoeRecoverTopKBwdOp( const at::Tensor& input_bwd, const at::Tensor& input_fwd, const at::Tensor& row_id_map, const at::Tensor& prob, const at::Tensor& act_grad, const at::Tensor& prob_grad );
|
||||
size_t CubSortPairGetStorageBytes(int64_t num_items);
|
||||
} // namespace wallx_cuda_permute
|
||||
|
||||
namespace wallx_cuda_rope {
|
||||
void Rope(const at::Tensor& q, const at::Tensor& k, const at::Tensor& q_embed, const at::Tensor& k_embed, const at::Tensor& cos, const at::Tensor& sin, bool interleave );
|
||||
void RopeInplace(const at::Tensor& q, const at::Tensor& k, const at::Tensor& cos, const at::Tensor& sin, bool interleave );
|
||||
void RopeInplacePack(const at::Tensor& qkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, bool interleave);
|
||||
void RopeInplacePackBackward(const at::Tensor& dqkv, const at::Tensor& cos, const at::Tensor& sin, int64_t q_num_heads, int64_t kv_num_heads, bool interleave);
|
||||
void RopeBackward( const at::Tensor& grad_q_embed, const at::Tensor& grad_k_embed, const at::Tensor& grad_q, const at::Tensor& grad_k, const at::Tensor& cos, const at::Tensor& sin, bool interleave );
|
||||
} // namespace wallx_cuda_rope
|
||||
|
||||
namespace wallx_cuda_rot_pos {
|
||||
void GetTokenCounts( const at::Tensor& grid_thw, const at::Tensor& token_counts, int spatial_merge_size );
|
||||
void RotPosEmb( const at::Tensor& inv_freq, const at::Tensor& grid_thw, const at::Tensor& output, const at::Tensor& cumsum_tokens, int spatial_merge_size );
|
||||
} // namespace wallx_cuda_rot_pos
|
||||
|
||||
namespace wallx_cuda_window_index {
|
||||
void GetWindowIndex( const at::Tensor& grid_thw, const at::Tensor& grid_info_tensor, const at::Tensor& window_indices, const at::Tensor& cu_window_seqlens, const at::Tensor& window_counts_tensor, int max_grid_t, int spatial_merge_size, int vit_merger_window_size, int patch_size, int spatial_merge_unit);
|
||||
void GetTotals( const at::Tensor& grid_thw, const at::Tensor& grid_info_tensor, const at::Tensor& global_totals_tensor, int spatial_merge_size, int vit_merger_window_size );
|
||||
} // namespace wallx_cuda_window_index
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("rope", &wallx_cuda_rope::Rope);
|
||||
m.def("rope_inplace", &wallx_cuda_rope::RopeInplace);
|
||||
m.def("rope_inplace_pack", &wallx_cuda_rope::RopeInplacePack);
|
||||
m.def("rope_inplace_pack_bwd", &wallx_cuda_rope::RopeInplacePackBackward);
|
||||
m.def("rope_bwd", &wallx_cuda_rope::RopeBackward);
|
||||
m.def("m_rope", &wallx_cuda_m_rope::MRope);
|
||||
m.def("m_rope_inplace", &wallx_cuda_m_rope::MRopeInplace);
|
||||
m.def("m_rope_inplace_pack", &wallx_cuda_m_rope::MRopeInplacePack);
|
||||
m.def("m_rope_inplace_pack_bwd", &wallx_cuda_m_rope::MRopeInplacePackBackward);
|
||||
m.def("m_rope_bwd", &wallx_cuda_m_rope::MRopeBackward);
|
||||
m.def("get_token_counts", &wallx_cuda_rot_pos::GetTokenCounts);
|
||||
m.def("rot_pos", &wallx_cuda_rot_pos::RotPosEmb);
|
||||
m.def("get_rope_index", &wallx_cuda_get_rope_index::GetRopeIndex);
|
||||
m.def("get_rope_index_getworkspace", &wallx_cuda_get_rope_index::GetRopeIndexGetWorkSpace);
|
||||
m.def("get_window_index", &wallx_cuda_window_index::GetWindowIndex);
|
||||
m.def("get_totals", &wallx_cuda_window_index::GetTotals);
|
||||
m.def("permute", &wallx_cuda_permute::MoePermuteTopKOp);
|
||||
m.def("unpermute", &wallx_cuda_permute::MoeRecoverTopKOp);
|
||||
m.def("unpermute_bwd", &wallx_cuda_permute::MoeRecoverTopKBwdOp);
|
||||
m.def("cub_sort_pair_get_storage_bytes", &wallx_cuda_permute::CubSortPairGetStorageBytes);
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
#pragma once
|
||||
|
||||
namespace wallx_cuda {
|
||||
|
||||
enum class ActivationType {
|
||||
RELU,
|
||||
SILU,
|
||||
GELU
|
||||
};
|
||||
|
||||
} // namespace wallx_cuda
|
||||
@@ -0,0 +1,30 @@
|
||||
#pragma once
|
||||
|
||||
#include <cassert>
|
||||
#include <cuda_fp16.h>
|
||||
#include "activation_types.h"
|
||||
|
||||
namespace wallx_cuda {
|
||||
|
||||
template<typename T, ActivationType activation_type>
|
||||
__forceinline__ __device__ T applyActivation(const T &x) {
|
||||
if constexpr (activation_type == ActivationType::RELU) {
|
||||
return x > (T)0.0f ? x : (T)0.0f;
|
||||
}
|
||||
else if constexpr (activation_type == ActivationType::SILU) {
|
||||
return (T)((float)x / (1.0f + __expf((float)-x)));
|
||||
}
|
||||
else if constexpr (activation_type == ActivationType::GELU) {
|
||||
// GELU implementation from vllm (gelu_new_kernel)
|
||||
const float x_f = (float)x;
|
||||
const float x3 = x_f * x_f * x_f;
|
||||
const float t = tanhf(0.79788456f * (x_f + 0.044715f * x3));
|
||||
return (T)(0.5f * x_f * (1.0f + t));
|
||||
}
|
||||
else {
|
||||
// No activation matches
|
||||
assert(false);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace wallx_cuda
|
||||
@@ -0,0 +1,111 @@
|
||||
// Re-enable CUDA half operators (PyTorch disables them)
|
||||
#undef __CUDA_NO_HALF_OPERATORS__
|
||||
#undef __CUDA_NO_HALF_CONVERSIONS__
|
||||
#undef __CUDA_NO_HALF2_OPERATORS__
|
||||
#undef __CUDA_NO_BFLOAT16_CONVERSIONS__
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
#define CUDA_CHECK(cmd) do { \
|
||||
cudaError_t result = cmd; \
|
||||
if (result != cudaSuccess) { \
|
||||
printf("[ERROR] CUDA error %s:%d '%s': (%d) %s\n", __FILE__, __LINE__, #cmd, (int)result, cudaGetErrorString(result)); \
|
||||
exit(-1); \
|
||||
} \
|
||||
} while(0)
|
||||
|
||||
inline void syncAndCheck(const char* const file, int const line, bool force_check = false) {
|
||||
#ifdef DEBUG
|
||||
force_check = true;
|
||||
#endif
|
||||
if (force_check) {
|
||||
cudaDeviceSynchronize();
|
||||
cudaError_t result = cudaGetLastError();
|
||||
if (result) {
|
||||
throw std::runtime_error(std::string("[ST] CUDA runtime error: ") + cudaGetErrorString(result) + " "
|
||||
+ file + ":" + std::to_string(line) + " \n");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define sync_check_cuda_error() syncAndCheck(__FILE__, __LINE__, false)
|
||||
#define sync_check_cuda_error_force() syncAndCheck(__FILE__, __LINE__, true)
|
||||
|
||||
// #ifdef DEBUG
|
||||
#define ASSERT_CHECK(__cond) \
|
||||
do { \
|
||||
const bool __cond_var = (__cond); \
|
||||
if (!__cond_var) { \
|
||||
::std::string __err_msg = \
|
||||
::std::string("`") + #__cond + \
|
||||
"` check failed at " + \
|
||||
__FILE__ + ":" + \
|
||||
::std::to_string(__LINE__); \
|
||||
throw std::runtime_error(__err_msg); \
|
||||
} \
|
||||
} while (0)
|
||||
// #else
|
||||
// #define ASSERT_CHECK(__cond) do { } while (0)
|
||||
// #endif
|
||||
|
||||
// Some stuff for indexing into an 1-D array
|
||||
#define INDEX_2D(dim1, dim2, index1, index2) \
|
||||
(((int64_t)index1) * (dim2) + (index2))
|
||||
#define INDEX_3D(dim1, dim2, dim3, index1, index2, index3) \
|
||||
(((int64_t)index1) * (dim2) * (dim3) + ((int64_t)index2) * (dim3) + (index3))
|
||||
#define INDEX_4D(dim1, dim2, dim3, dim4, index1, index2, index3, index4) \
|
||||
(((int64_t)index1) * (dim2) * (dim3) * (dim4) + ((int64_t)index2) * (dim3) * (dim4) + ((int64_t)index3) * (dim4) + (index4))
|
||||
#define INDEX_5D(dim1, dim2, dim3, dim4, dim5, index1, index2, index3, index4, index5) \
|
||||
(((int64_t)index1) * (dim2) * (dim3) * (dim4) * (dim5) + ((int64_t)index2) * (dim3) * (dim4) * (dim5) + ((int64_t)index3) * (dim4) * (dim5) + (index4) * (dim5) + (index5))
|
||||
|
||||
template<typename T>
|
||||
struct VecTraits;
|
||||
|
||||
template<>
|
||||
struct VecTraits<float> {
|
||||
using Type = float4;
|
||||
static constexpr int vec_size = 4;
|
||||
__device__ static inline Type load(const float* p) { return *reinterpret_cast<const float4*>(p); }
|
||||
__device__ static inline void store(float* p, const Type& v) { *reinterpret_cast<float4*>(p) = v; }
|
||||
};
|
||||
|
||||
template<>
|
||||
struct VecTraits<__half> {
|
||||
using Type = __half2;
|
||||
static constexpr int vec_size = 2;
|
||||
__device__ static inline Type load(const __half* p) { return *reinterpret_cast<const __half2*>(p); }
|
||||
__device__ static inline void store(__half* p, const Type& v) { *reinterpret_cast<__half2*>(p) = v; }
|
||||
};
|
||||
|
||||
template<>
|
||||
struct VecTraits<__nv_bfloat16> {
|
||||
using Type = __nv_bfloat162;
|
||||
static constexpr int vec_size = 2;
|
||||
__device__ static inline Type load(const __nv_bfloat16* p) { return *reinterpret_cast<const __nv_bfloat162*>(p); }
|
||||
__device__ static inline void store(__nv_bfloat16* p, const Type& v) { *reinterpret_cast<__nv_bfloat162*>(p) = v; }
|
||||
};
|
||||
|
||||
template<typename T>
|
||||
cudaDataType_t getCudaDataType() {
|
||||
if (std::is_same<T, half>::value) {
|
||||
return CUDA_R_16F;
|
||||
}
|
||||
else if (std::is_same<T, __nv_bfloat16>::value) {
|
||||
return CUDA_R_16BF;
|
||||
}
|
||||
else if (std::is_same<T, float>::value) {
|
||||
return CUDA_R_32F;
|
||||
}
|
||||
else {
|
||||
throw std::runtime_error("Cuda data type: Unsupported type");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,692 @@
|
||||
#undef __CUDA_NO_HALF_OPERATORS__
|
||||
#undef __CUDA_NO_HALF_CONVERSIONS__
|
||||
#undef __CUDA_NO_BFLOAT16_CONVERSIONS__
|
||||
#undef __CUDA_NO_HALF2_OPERATORS__
|
||||
|
||||
#include <vector>
|
||||
#include <tuple>
|
||||
#include <cstdint>
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda.h>
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include "../common/cuda_utils.h"
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
|
||||
namespace wallx_cuda_get_rope_index {
|
||||
|
||||
// Constants and macros
|
||||
#define MAX_SEQ_LEN 8192
|
||||
#define MAX_VISION_TOKENS 64
|
||||
#define WARP_SIZE 32
|
||||
#define MAX_THREADS_PER_BLOCK 1024
|
||||
|
||||
// Vision content descriptor
|
||||
struct VisionDescriptor
|
||||
{
|
||||
int64_t start_pos; // Starting position in the sequence
|
||||
int64_t token_pos; // Position of the vision token
|
||||
int64_t patch_count; // Number of patches
|
||||
int64_t grid_t, grid_h, grid_w; // Grid dimensions (T, H, W)
|
||||
float time_interval; // Time interval (for video)
|
||||
int64_t is_video; // 0 = image, 1 = video
|
||||
int64_t position_offset; // Positional encoding offset
|
||||
};
|
||||
|
||||
// Device function: fast integer division using reciprocal multiplication
|
||||
__device__ __forceinline__ int64_t fast_div(int64_t a, int64_t b)
|
||||
{
|
||||
return __float2ll_rd(__ll2float_rn(a) * __frcp_rn(__ll2float_rn(b)));
|
||||
}
|
||||
|
||||
// Device function: map 3D patch index to (t, h, w) coordinates
|
||||
__device__ __forceinline__ void get_3d_coords(int64_t patch_idx, int64_t H, int64_t W,
|
||||
int64_t &t, int64_t &h, int64_t &w)
|
||||
{
|
||||
int64_t hw = H * W;
|
||||
t = fast_div(patch_idx, hw);
|
||||
int64_t remaining = patch_idx - t * hw;
|
||||
h = fast_div(remaining, W);
|
||||
w = remaining - h * W;
|
||||
}
|
||||
|
||||
// Stage 1: Kernel to count image and video tokens per batch
|
||||
__global__ void compute_vision_counts(
|
||||
const int64_t *input_ids, // (batch_size, seq_len)
|
||||
const int64_t *attention_mask, // (batch_size, seq_len)
|
||||
int64_t *image_counts, // (batch_size,) - output: number of images per batch
|
||||
int64_t *video_counts, // (batch_size,) - output: number of videos per batch
|
||||
const int64_t batch_size,
|
||||
const int64_t seq_len,
|
||||
const int64_t image_token_id,
|
||||
const int64_t video_token_id,
|
||||
const int64_t vision_start_token_id)
|
||||
{
|
||||
int64_t batch_idx = blockIdx.x;
|
||||
int64_t thread_idx = threadIdx.x;
|
||||
|
||||
if (batch_idx >= batch_size)
|
||||
return;
|
||||
|
||||
__shared__ int64_t shared_image_counts[MAX_THREADS_PER_BLOCK];
|
||||
__shared__ int64_t shared_video_counts[MAX_THREADS_PER_BLOCK];
|
||||
|
||||
// Initialize per-thread counters
|
||||
int64_t thread_image_count = 0;
|
||||
int64_t thread_video_count = 0;
|
||||
|
||||
// Parallel scan over sequence tokens (each thread processes multiple tokens)
|
||||
for (int64_t i = thread_idx; i < seq_len - 1; i += blockDim.x)
|
||||
{
|
||||
// Skip if masked out
|
||||
if ((attention_mask != nullptr) && attention_mask[batch_idx * seq_len + i] == 0)
|
||||
continue;
|
||||
|
||||
int64_t token_id = input_ids[batch_idx * seq_len + i];
|
||||
|
||||
// Check if current token is a vision start token followed by image/video token
|
||||
if (token_id == vision_start_token_id && i + 1 < seq_len)
|
||||
{
|
||||
int64_t next_token = input_ids[batch_idx * seq_len + i + 1];
|
||||
|
||||
if (next_token == image_token_id)
|
||||
{
|
||||
thread_image_count++;
|
||||
}
|
||||
else if (next_token == video_token_id)
|
||||
{
|
||||
thread_video_count++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Store per-thread counts in shared memory
|
||||
shared_image_counts[thread_idx] = thread_image_count;
|
||||
shared_video_counts[thread_idx] = thread_video_count;
|
||||
__syncthreads();
|
||||
|
||||
// Parallel reduction to sum counts across threads
|
||||
for (int64_t stride = blockDim.x / 2; stride > 0; stride /= 2)
|
||||
{
|
||||
if (thread_idx < stride)
|
||||
{
|
||||
shared_image_counts[thread_idx] += shared_image_counts[thread_idx + stride];
|
||||
shared_video_counts[thread_idx] += shared_video_counts[thread_idx + stride];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Write final result for this batch to global memory
|
||||
if (thread_idx == 0)
|
||||
{
|
||||
image_counts[batch_idx] = shared_image_counts[0];
|
||||
video_counts[batch_idx] = shared_video_counts[0];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// Stage 2: Preprocessing kernel - parse and analyze all vision content
|
||||
__global__ void preprocess_vision_tokens(
|
||||
const int64_t *input_ids, // (batch_size, seq_len)
|
||||
const int64_t *attention_mask, // (batch_size, seq_len)
|
||||
const int64_t *image_grid_thw, // (max_images, 3)
|
||||
const int64_t *video_grid_thw, // (max_videos, 3)
|
||||
const float *second_per_grid_ts, // (max_videos,)
|
||||
const int64_t *image_counts, // (batch_size,)
|
||||
const int64_t *video_counts, // (batch_size,)
|
||||
VisionDescriptor *vision_desc, // (batch_size, MAX_VISION_TOKENS)
|
||||
int64_t *vision_counts, // (batch_size,) - total vision tokens per batch
|
||||
int64_t *text_lengths, // (batch_size, MAX_VISION_TOKENS+1) - lengths of text segments
|
||||
int64_t *position_offsets, // (batch_size, MAX_VISION_TOKENS+1) - cumulative position offsets
|
||||
const int64_t batch_size,
|
||||
const int64_t seq_len,
|
||||
const int64_t spatial_merge_size,
|
||||
const int64_t image_token_id,
|
||||
const int64_t video_token_id,
|
||||
const int64_t vision_start_token_id,
|
||||
const float tokens_per_second)
|
||||
{
|
||||
// Each thread processes one batch sequence
|
||||
int64_t batch_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (batch_idx >= batch_size)
|
||||
return;
|
||||
|
||||
// Compute cumulative image/video indices before this batch
|
||||
int64_t image_idx = 0, video_idx = 0;
|
||||
for (int64_t i = 0; i < batch_idx; i++)
|
||||
{
|
||||
image_idx += image_counts[i];
|
||||
video_idx += video_counts[i];
|
||||
}
|
||||
|
||||
// Local state
|
||||
int64_t vision_count = 0;
|
||||
int64_t current_pos = 0;
|
||||
int64_t position_offset = 0;
|
||||
|
||||
// Sequential scan over the sequence (single-threaded per batch)
|
||||
for (int64_t i = 0; i < seq_len - 1; i++)
|
||||
{
|
||||
// Skip masked tokens
|
||||
if ((attention_mask != nullptr) && attention_mask[batch_idx * seq_len + i] == 0)
|
||||
continue;
|
||||
|
||||
int64_t token_id = input_ids[batch_idx * seq_len + i];
|
||||
|
||||
// Check for vision start token
|
||||
if (token_id == vision_start_token_id)
|
||||
{
|
||||
int64_t next_token = input_ids[batch_idx * seq_len + i + 1];
|
||||
|
||||
if (next_token == image_token_id || next_token == video_token_id)
|
||||
{
|
||||
// Record length of preceding text segment
|
||||
int64_t vision_pos = i + 1;
|
||||
text_lengths[batch_idx * (MAX_VISION_TOKENS + 1) + vision_count] = vision_pos - current_pos;
|
||||
position_offsets[batch_idx * (MAX_VISION_TOKENS + 1) + vision_count] = position_offset;
|
||||
position_offset += (vision_pos - current_pos);
|
||||
|
||||
// Fetch grid dimensions
|
||||
int64_t T, H, W;
|
||||
float time_interval = 0.0f;
|
||||
int64_t is_video = (next_token == video_token_id) ? 1 : 0;
|
||||
|
||||
if (is_video == 0)
|
||||
{
|
||||
// Image case
|
||||
T = image_grid_thw[image_idx * 3 + 0];
|
||||
H = image_grid_thw[image_idx * 3 + 1];
|
||||
W = image_grid_thw[image_idx * 3 + 2];
|
||||
image_idx++;
|
||||
}
|
||||
else
|
||||
{
|
||||
// Video case
|
||||
T = video_grid_thw[video_idx * 3 + 0];
|
||||
H = video_grid_thw[video_idx * 3 + 1];
|
||||
W = video_grid_thw[video_idx * 3 + 2];
|
||||
time_interval = second_per_grid_ts[video_idx];
|
||||
video_idx++;
|
||||
}
|
||||
|
||||
// Compute number of patches after spatial merging
|
||||
int64_t H_merged = H / spatial_merge_size;
|
||||
int64_t W_merged = W / spatial_merge_size;
|
||||
int64_t patch_count = T * H_merged * W_merged;
|
||||
|
||||
// Populate vision descriptor
|
||||
if (vision_count < MAX_VISION_TOKENS)
|
||||
{
|
||||
VisionDescriptor &desc = vision_desc[batch_idx * MAX_VISION_TOKENS + vision_count];
|
||||
desc.start_pos = current_pos;
|
||||
desc.token_pos = vision_pos;
|
||||
desc.patch_count = patch_count;
|
||||
desc.grid_t = T;
|
||||
desc.grid_h = H_merged;
|
||||
desc.grid_w = W_merged;
|
||||
desc.time_interval = time_interval;
|
||||
desc.is_video = is_video;
|
||||
desc.position_offset = position_offset;
|
||||
|
||||
// Update position offset based on content type
|
||||
if (is_video)
|
||||
{
|
||||
position_offset += max(static_cast<int64_t>((T - 1) * time_interval * tokens_per_second) + 1,
|
||||
static_cast<int64_t>(max(H_merged, W_merged)));
|
||||
}
|
||||
else
|
||||
{
|
||||
position_offset += max(H_merged, W_merged);
|
||||
}
|
||||
|
||||
current_pos = vision_pos + patch_count;
|
||||
vision_count++;
|
||||
}
|
||||
|
||||
// Skip over the vision patch tokens
|
||||
i = vision_pos + patch_count - 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Store total vision token count for this batch
|
||||
vision_counts[batch_idx] = vision_count;
|
||||
|
||||
// Handle final text segment after last vision token
|
||||
int64_t effective_len = seq_len;
|
||||
if (attention_mask != nullptr)
|
||||
{
|
||||
// Find last valid (unmasked) token
|
||||
for (int64_t i = seq_len - 1; i >= 0; i--)
|
||||
{
|
||||
if (attention_mask[batch_idx * seq_len + i] != 0)
|
||||
{
|
||||
effective_len = i + 1;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
text_lengths[batch_idx * (MAX_VISION_TOKENS + 1) + vision_count] = effective_len - current_pos;
|
||||
position_offsets[batch_idx * (MAX_VISION_TOKENS + 1) + vision_count] = position_offset;
|
||||
}
|
||||
|
||||
|
||||
// Stage 3: Main kernel - compute 3D position IDs for all tokens in parallel
|
||||
__global__ void compute_3d_positions(
|
||||
const int64_t *input_ids, // (batch_size, seq_len)
|
||||
const int64_t *attention_mask, // (batch_size, seq_len)
|
||||
const VisionDescriptor *vision_desc, // (batch_size, MAX_VISION_TOKENS)
|
||||
const int64_t *vision_counts, // (batch_size,)
|
||||
const int64_t *text_lengths, // (batch_size, MAX_VISION_TOKENS+1)
|
||||
const int64_t *position_offsets, // (batch_size, MAX_VISION_TOKENS+1)
|
||||
int64_t *position_ids, // (3, batch_size, seq_len) - output 3D position IDs
|
||||
int64_t *mrope_deltas, // (batch_size,) - RoPE length adjustment
|
||||
const int64_t batch_size,
|
||||
const int64_t seq_len,
|
||||
const float tokens_per_second)
|
||||
{
|
||||
// Grid: (batch_size), Block: (threads_per_block)
|
||||
int64_t batch_idx = blockIdx.x;
|
||||
int64_t thread_idx = threadIdx.x;
|
||||
|
||||
if (batch_idx >= batch_size)
|
||||
return;
|
||||
|
||||
__shared__ VisionDescriptor shared_visions[MAX_VISION_TOKENS];
|
||||
__shared__ int64_t shared_position_offsets[MAX_VISION_TOKENS + 1];
|
||||
__shared__ int64_t shared_max_positions[MAX_THREADS_PER_BLOCK]; // For block-wide reduction
|
||||
|
||||
int64_t shared_vision_count = vision_counts[batch_idx];
|
||||
|
||||
// Cooperative loading of vision descriptors into shared memory
|
||||
for (int64_t i = thread_idx; i < MAX_VISION_TOKENS; i += blockDim.x)
|
||||
{
|
||||
if (i < shared_vision_count)
|
||||
{
|
||||
shared_visions[i] = vision_desc[batch_idx * MAX_VISION_TOKENS + i];
|
||||
}
|
||||
}
|
||||
|
||||
for (int64_t i = thread_idx; i < MAX_VISION_TOKENS + 1; i += blockDim.x)
|
||||
{
|
||||
shared_position_offsets[i] = position_offsets[batch_idx * (MAX_VISION_TOKENS + 1) + i];
|
||||
}
|
||||
|
||||
// Initialize per-thread max position
|
||||
int64_t thread_max_position = -1;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Parallel processing: each thread handles multiple tokens
|
||||
for (int64_t token_idx = thread_idx; token_idx < seq_len; token_idx += blockDim.x)
|
||||
{
|
||||
// Check validity via attention mask
|
||||
int mask_offset = 0;
|
||||
bool is_valid_token = true;
|
||||
if (attention_mask != nullptr)
|
||||
{
|
||||
is_valid_token = attention_mask[batch_idx * seq_len + token_idx] != 0;
|
||||
for (int i = 0; i < token_idx; i++) {
|
||||
if (attention_mask[batch_idx * seq_len + i] == 0) {
|
||||
mask_offset += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!is_valid_token)
|
||||
{
|
||||
// Set masked tokens to default position (1)
|
||||
position_ids[0 * batch_size * seq_len + batch_idx * seq_len + token_idx] = 1;
|
||||
position_ids[1 * batch_size * seq_len + batch_idx * seq_len + token_idx] = 1;
|
||||
position_ids[2 * batch_size * seq_len + batch_idx * seq_len + token_idx] = 1;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Determine which segment this token belongs to
|
||||
int64_t segment_idx = -9999999;
|
||||
int64_t local_pos = token_idx;
|
||||
|
||||
// Linear search (vision count is small, so this is efficient)
|
||||
for (int64_t v = 0; v < shared_vision_count; v++)
|
||||
{
|
||||
if (token_idx < shared_visions[v].token_pos)
|
||||
{
|
||||
// Token belongs to text segment before vision v
|
||||
segment_idx = v;
|
||||
local_pos = token_idx - (v > 0 ? shared_visions[v - 1].token_pos + shared_visions[v - 1].patch_count : 0);
|
||||
break;
|
||||
}
|
||||
else if (token_idx < shared_visions[v].token_pos + shared_visions[v].patch_count)
|
||||
{
|
||||
// Token belongs to vision patch v
|
||||
segment_idx = -(v + 1); // Negative index indicates vision
|
||||
local_pos = token_idx - shared_visions[v].token_pos;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (segment_idx == -9999999)
|
||||
{
|
||||
// Token belongs to final text segment
|
||||
segment_idx = shared_vision_count;
|
||||
int64_t last_vision_end = 0;
|
||||
if (shared_vision_count > 0)
|
||||
{
|
||||
last_vision_end = shared_visions[shared_vision_count - 1].token_pos +
|
||||
shared_visions[shared_vision_count - 1].patch_count;
|
||||
}
|
||||
local_pos = token_idx - last_vision_end;
|
||||
}
|
||||
|
||||
int64_t pos_t, pos_h, pos_w;
|
||||
|
||||
if (segment_idx >= 0)
|
||||
{
|
||||
// Text token: use 1D position encoding
|
||||
int64_t offset = shared_position_offsets[segment_idx];
|
||||
pos_t = pos_h = pos_w = offset + local_pos;
|
||||
}
|
||||
else
|
||||
{
|
||||
// Vision patch: compute 3D position encoding
|
||||
int64_t vision_idx = -(segment_idx + 1);
|
||||
const VisionDescriptor &desc = shared_visions[vision_idx];
|
||||
|
||||
// Map linear patch index to (t, h, w)
|
||||
int64_t t, h, w;
|
||||
get_3d_coords(local_pos, desc.grid_h, desc.grid_w, t, h, w);
|
||||
|
||||
// Compute 3D positions with temporal scaling for video
|
||||
pos_t = static_cast<int64_t>(t * desc.time_interval * tokens_per_second) + desc.position_offset;
|
||||
pos_h = h + desc.position_offset;
|
||||
pos_w = w + desc.position_offset;
|
||||
}
|
||||
|
||||
// Write 3D position IDs
|
||||
position_ids[0 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_t - mask_offset;
|
||||
position_ids[1 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_h - mask_offset;
|
||||
position_ids[2 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_w - mask_offset;
|
||||
|
||||
// Track per-thread maximum position
|
||||
int64_t max_pos = max(pos_t - mask_offset, max(pos_h - mask_offset, pos_w - mask_offset));
|
||||
thread_max_position = max(thread_max_position, max_pos);
|
||||
}
|
||||
|
||||
// Store per-thread max into shared memory
|
||||
shared_max_positions[thread_idx] = thread_max_position;
|
||||
__syncthreads();
|
||||
|
||||
// Block-wide reduction to find global max position
|
||||
for (int64_t stride = blockDim.x / 2; stride > 0; stride /= 2)
|
||||
{
|
||||
if (thread_idx < stride)
|
||||
{
|
||||
shared_max_positions[thread_idx] = max(shared_max_positions[thread_idx],
|
||||
shared_max_positions[thread_idx + stride]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Compute mRoPE delta: (max_pos + 1) - seq_len
|
||||
if (thread_idx == 0)
|
||||
{
|
||||
int64_t global_max_position = shared_max_positions[0];
|
||||
mrope_deltas[batch_idx] = global_max_position + 1 - seq_len;
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback kernel when no vision tokens exist: assign sequential positions based on attention mask
|
||||
__global__ void compute_3d_positions_mask_text(
|
||||
const int64_t* __restrict__ attention_mask, // (batch_size, seq_len)
|
||||
int64_t* __restrict__ position_ids, // (3, batch_size, seq_len) - flattened row-major
|
||||
int64_t* __restrict__ mrope_deltas, // (batch_size,)
|
||||
const int64_t batch_size,
|
||||
const int64_t seq_len
|
||||
) {
|
||||
// One block per batch
|
||||
int batch_idx = blockIdx.x;
|
||||
if (batch_idx >= batch_size) return;
|
||||
|
||||
extern __shared__ int64_t shared_mask[];
|
||||
int64_t* mask = shared_mask;
|
||||
|
||||
// Load attention mask for current batch into shared memory
|
||||
for (int i = threadIdx.x; i < seq_len; i += blockDim.x) {
|
||||
mask[i] = attention_mask[batch_idx * seq_len + i];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int tid = threadIdx.x;
|
||||
int64_t base_offset = batch_idx * seq_len;
|
||||
int64_t total_valid = 0;
|
||||
|
||||
// Compute prefix sum (cumulative count of valid tokens)
|
||||
// Simple O(seq_len^2) per-thread loop - acceptable for moderate seq_len
|
||||
for (int idx = tid; idx < seq_len; idx += blockDim.x) {
|
||||
int64_t sum = 0;
|
||||
for (int i = 0; i <= idx; ++i) {
|
||||
sum += mask[i];
|
||||
}
|
||||
int64_t pos_val = (mask[idx] ? (sum - 1) : 1); // Use 0-based if valid, else 1
|
||||
|
||||
// Write same position to all three dimensions
|
||||
int64_t flat_idx0 = 0 * batch_size * seq_len + base_offset + idx;
|
||||
int64_t flat_idx1 = 1 * batch_size * seq_len + base_offset + idx;
|
||||
int64_t flat_idx2 = 2 * batch_size * seq_len + base_offset + idx;
|
||||
|
||||
position_ids[flat_idx0] = pos_val;
|
||||
position_ids[flat_idx1] = pos_val;
|
||||
position_ids[flat_idx2] = pos_val;
|
||||
|
||||
// Only thread 0 tracks total valid tokens
|
||||
if (tid == 0) {
|
||||
total_valid = sum; // After last idx, sum = total number of 1s
|
||||
}
|
||||
}
|
||||
|
||||
// Compute mrope_deltas: (total_valid) - seq_len
|
||||
if (tid == 0) {
|
||||
// Re-compute total_valid robustly
|
||||
total_valid = 0;
|
||||
for (int i = 0; i < seq_len; ++i) {
|
||||
total_valid += mask[i];
|
||||
}
|
||||
mrope_deltas[batch_idx] = total_valid - seq_len;
|
||||
}
|
||||
}
|
||||
|
||||
// Kernel to fill position_ids with arange(seq_len) in 3D layout when no vision or mask
|
||||
__global__ void arange_3d(int64_t seq_len,
|
||||
int64_t* out) {
|
||||
int i = blockIdx.x; // dimension index: 0, 1, 2
|
||||
int j = blockIdx.y; // batch index
|
||||
int64_t k = blockIdx.z * blockDim.x + threadIdx.x; // token index
|
||||
|
||||
if (k >= seq_len) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Flatten index: (dim, batch, token) -> linear
|
||||
int64_t idx = i * gridDim.y * seq_len + j * seq_len + k;
|
||||
out[idx] = k;
|
||||
}
|
||||
|
||||
// Kernel to zero-initialize a 1D tensor
|
||||
__global__ void fill_zeros_1d(int64_t numel, int64_t* data) {
|
||||
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < numel) {
|
||||
data[idx] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
// Compute required workspace size (in bytes)
|
||||
int64_t GetRopeIndexGetWorkSpace(const c10::optional<at::Tensor>& input_ids,
|
||||
const c10::optional<at::Tensor>& image_grid_thw,
|
||||
const c10::optional<at::Tensor>& video_grid_thw) {
|
||||
|
||||
if (!image_grid_thw.has_value() && !video_grid_thw.has_value()) {
|
||||
return 1; // minimal workspace
|
||||
} else {
|
||||
const at::Tensor& input_ids_tensor = input_ids.value();
|
||||
const int batch_size = input_ids_tensor.size(0);
|
||||
return (batch_size * MAX_VISION_TOKENS * static_cast<int64_t>(sizeof(VisionDescriptor)) +
|
||||
3 * batch_size * sizeof(int64_t) +
|
||||
2 * batch_size * (MAX_VISION_TOKENS + 1) * sizeof(int64_t));
|
||||
}
|
||||
}
|
||||
|
||||
// Main entry function: compute RoPE position IDs and mRoPE deltas
|
||||
void GetRopeIndex(const c10::optional<at::Tensor>& input_ids,
|
||||
const c10::optional<at::Tensor>& image_grid_thw,
|
||||
const c10::optional<at::Tensor>& video_grid_thw,
|
||||
const c10::optional<at::Tensor>& second_per_grid_ts,
|
||||
const c10::optional<at::Tensor>& attention_mask,
|
||||
const at::Tensor& position_ids,
|
||||
const at::Tensor& mrope_deltas,
|
||||
const at::Tensor& workspace,
|
||||
int spatial_merge_size,
|
||||
int image_token_id,
|
||||
int video_token_id,
|
||||
int vision_start_token_id,
|
||||
float tokens_per_second) {
|
||||
|
||||
ASSERT_CHECK(input_ids.has_value());
|
||||
|
||||
const at::Tensor& input_ids_tensor = input_ids.value();
|
||||
ASSERT_CHECK(input_ids_tensor.dim() == 2);
|
||||
ASSERT_CHECK(input_ids_tensor.is_contiguous());
|
||||
|
||||
const auto batch_size = input_ids_tensor.size(0);
|
||||
const auto seq_len = input_ids_tensor.size(1);
|
||||
const auto device = input_ids_tensor.device();
|
||||
|
||||
const int64_t *input_ids_ptr = static_cast<const int64_t*>(input_ids_tensor.data_ptr());
|
||||
|
||||
const int64_t *attention_mask_ptr = nullptr;
|
||||
|
||||
if (attention_mask.has_value()) {
|
||||
const at::Tensor& attention_mask_tensor = attention_mask.value();
|
||||
ASSERT_CHECK(attention_mask_tensor.is_contiguous());
|
||||
attention_mask_ptr = static_cast<const int64_t*>(attention_mask_tensor.data_ptr());
|
||||
}
|
||||
|
||||
const int64_t *image_grid_thw_ptr = nullptr;
|
||||
|
||||
if (image_grid_thw.has_value()) {
|
||||
const at::Tensor& image_grid_thw_tensor = image_grid_thw.value();
|
||||
ASSERT_CHECK(image_grid_thw_tensor.dim() == 2 && image_grid_thw_tensor.size(1) == 3);
|
||||
ASSERT_CHECK(image_grid_thw_tensor.is_contiguous());
|
||||
image_grid_thw_ptr = static_cast<const int64_t*>(image_grid_thw_tensor.data_ptr());
|
||||
}
|
||||
|
||||
const int64_t *video_grid_thw_ptr = nullptr;
|
||||
|
||||
if (video_grid_thw.has_value()) {
|
||||
const at::Tensor& video_grid_thw_tensor = video_grid_thw.value();
|
||||
ASSERT_CHECK(video_grid_thw_tensor.is_contiguous());
|
||||
ASSERT_CHECK(video_grid_thw_tensor.dim() == 2 && video_grid_thw_tensor.size(1) == 3);
|
||||
video_grid_thw_ptr = static_cast<const int64_t*>(video_grid_thw_tensor.data_ptr());
|
||||
}
|
||||
|
||||
const float *second_per_grid_ts_ptr = nullptr;
|
||||
|
||||
if (second_per_grid_ts.has_value()) {
|
||||
const at::Tensor& second_per_grid_ts_tensor = second_per_grid_ts.value();
|
||||
ASSERT_CHECK(second_per_grid_ts_tensor.is_contiguous());
|
||||
ASSERT_CHECK(second_per_grid_ts_tensor.dim() == 1);
|
||||
second_per_grid_ts_ptr = static_cast<const float*>(second_per_grid_ts_tensor.data_ptr());
|
||||
}
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
// Case: no vision tokens -> use simple positional encoding
|
||||
if (!image_grid_thw.has_value() && !video_grid_thw.has_value()) {
|
||||
if (attention_mask.has_value()) {
|
||||
// Use attention mask to compute valid positions
|
||||
int block_size = min(1024, (int)seq_len);
|
||||
dim3 grid(batch_size);
|
||||
dim3 block(block_size);
|
||||
size_t shared_mem_size = seq_len * sizeof(int64_t);
|
||||
|
||||
compute_3d_positions_mask_text<<<grid, block, shared_mem_size>>>(
|
||||
attention_mask_ptr, static_cast<int64_t*>(position_ids.data_ptr()), static_cast<int64_t*>(mrope_deltas.data_ptr()), batch_size, seq_len
|
||||
);
|
||||
return;
|
||||
} else {
|
||||
// No mask, no vision -> use simple arange
|
||||
dim3 thread_num(256);
|
||||
dim3 grid(3, batch_size, (seq_len + 255) / 256);
|
||||
arange_3d<<<grid, thread_num, 0, stream>>>(seq_len, static_cast<int64_t*>(position_ids.data_ptr()));
|
||||
|
||||
int block_size = 256;
|
||||
int grid_size = (batch_size + block_size - 1) / block_size;
|
||||
fill_zeros_1d<<<grid_size, block_size, 0, stream>>>(batch_size, static_cast<int64_t*>(mrope_deltas.data_ptr()));
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Allocate workspace buffers
|
||||
void* work_ptr = workspace.data_ptr();
|
||||
VisionDescriptor *d_vision_desc = reinterpret_cast<VisionDescriptor*>(work_ptr);
|
||||
|
||||
int64_t *d_vision_counts = reinterpret_cast<int64_t*>(reinterpret_cast<VisionDescriptor*>(work_ptr) + batch_size * MAX_VISION_TOKENS);
|
||||
int64_t *d_text_lengths = d_vision_counts + batch_size;
|
||||
int64_t *d_position_offsets = d_text_lengths + batch_size * (MAX_VISION_TOKENS + 1);
|
||||
int64_t *d_image_counts = d_position_offsets + batch_size * (MAX_VISION_TOKENS + 1);
|
||||
int64_t *d_video_counts = d_image_counts + batch_size;
|
||||
|
||||
// Stage 1: count vision tokens
|
||||
dim3 index_grid(static_cast<unsigned int>(batch_size));
|
||||
dim3 index_block(256);
|
||||
|
||||
compute_vision_counts<<<index_grid, index_block, 0, stream>>>(
|
||||
input_ids_ptr, attention_mask_ptr,
|
||||
d_image_counts, d_video_counts,
|
||||
batch_size, seq_len, image_token_id, video_token_id, vision_start_token_id);
|
||||
|
||||
// Stage 2: preprocess vision tokens
|
||||
int64_t threads_per_block = std::min(static_cast<int64_t>(batch_size), static_cast<int64_t>(256));
|
||||
int64_t num_blocks = (batch_size + threads_per_block - 1) / threads_per_block;
|
||||
|
||||
dim3 preprocess_grid(static_cast<unsigned int>(num_blocks));
|
||||
dim3 preprocess_block(static_cast<unsigned int>(threads_per_block));
|
||||
|
||||
preprocess_vision_tokens<<<preprocess_grid, preprocess_block, 0, stream>>>(
|
||||
input_ids_ptr, attention_mask_ptr, image_grid_thw_ptr, video_grid_thw_ptr,
|
||||
second_per_grid_ts_ptr, d_image_counts, d_video_counts,
|
||||
d_vision_desc, d_vision_counts, d_text_lengths, d_position_offsets,
|
||||
batch_size, seq_len, spatial_merge_size,
|
||||
image_token_id, video_token_id, vision_start_token_id, tokens_per_second);
|
||||
|
||||
// Stage 3: compute 3D positions
|
||||
threads_per_block = std::min(static_cast<int64_t>(seq_len), static_cast<int64_t>(MAX_THREADS_PER_BLOCK));
|
||||
if (threads_per_block < 32) threads_per_block = 32;
|
||||
|
||||
int64_t power_of_2 = 1;
|
||||
while (power_of_2 < threads_per_block) power_of_2 *= 2;
|
||||
if (power_of_2 > MAX_THREADS_PER_BLOCK) power_of_2 = MAX_THREADS_PER_BLOCK;
|
||||
threads_per_block = power_of_2;
|
||||
|
||||
dim3 compute_grid(static_cast<unsigned int>(batch_size));
|
||||
dim3 compute_block(static_cast<unsigned int>(threads_per_block));
|
||||
|
||||
compute_3d_positions<<<compute_grid, compute_block, 0, stream>>>(
|
||||
input_ids_ptr, attention_mask_ptr, d_vision_desc, d_vision_counts,
|
||||
d_text_lengths, d_position_offsets,
|
||||
static_cast<int64_t*>(position_ids.data_ptr()), static_cast<int64_t*>(mrope_deltas.data_ptr()),
|
||||
batch_size, seq_len, tokens_per_second);
|
||||
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
} // wallx_cuda_get_rope_index
|
||||
|
||||
// TVM FFI exports
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
@@ -0,0 +1,840 @@
|
||||
// Re-enable CUDA half operators (PyTorch disables them)
|
||||
#undef __CUDA_NO_HALF_OPERATORS__
|
||||
#undef __CUDA_NO_HALF_CONVERSIONS__
|
||||
#undef __CUDA_NO_HALF2_OPERATORS__
|
||||
#undef __CUDA_NO_BFLOAT16_CONVERSIONS__
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cassert>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cmath>
|
||||
#include <cfloat>
|
||||
#include "../common/cuda_utils.h"
|
||||
|
||||
namespace wallx_cuda_m_rope {
|
||||
|
||||
// Helper: convert vector of T to float4
|
||||
template<typename T>
|
||||
__device__ __forceinline__ float4 to_float4(const T* ptr) {
|
||||
if constexpr (std::is_same_v<T, half>) {
|
||||
half2 h2_0 = reinterpret_cast<const half2*>(ptr)[0];
|
||||
half2 h2_1 = reinterpret_cast<const half2*>(ptr)[1];
|
||||
float2 f2_0 = __half22float2(h2_0);
|
||||
float2 f2_1 = __half22float2(h2_1);
|
||||
return make_float4(f2_0.x, f2_0.y, f2_1.x, f2_1.y);
|
||||
} else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
|
||||
__nv_bfloat162 b2_0 = reinterpret_cast<const __nv_bfloat162*>(ptr)[0];
|
||||
__nv_bfloat162 b2_1 = reinterpret_cast<const __nv_bfloat162*>(ptr)[1];
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
float2 f2_0 = __bfloat1622float2(b2_0);
|
||||
float2 f2_1 = __bfloat1622float2(b2_1);
|
||||
return make_float4(f2_0.x, f2_0.y, f2_1.x, f2_1.y);
|
||||
#else
|
||||
assert(false && "Unsupported on arch < 800");
|
||||
#endif
|
||||
} else if constexpr (std::is_same_v<T, float>) {
|
||||
return reinterpret_cast<const float4*>(ptr)[0];
|
||||
} else {
|
||||
static_assert(sizeof(T) == 0, "Unsupported type");
|
||||
}
|
||||
}
|
||||
|
||||
// Helper: convert float4 back to vector of T
|
||||
template<typename T>
|
||||
__device__ __forceinline__ void from_float4(T* ptr, const float4& f4) {
|
||||
if constexpr (std::is_same_v<T, half>) {
|
||||
half2 h2_0 = __float22half2_rn(make_float2(f4.x, f4.y));
|
||||
half2 h2_1 = __float22half2_rn(make_float2(f4.z, f4.w));
|
||||
reinterpret_cast<half2*>(ptr)[0] = h2_0;
|
||||
reinterpret_cast<half2*>(ptr)[1] = h2_1;
|
||||
} else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
|
||||
#if __CUDA_ARCH__ >= 800
|
||||
__nv_bfloat162 b2_0 = __float22bfloat162_rn(make_float2(f4.x, f4.y));
|
||||
__nv_bfloat162 b2_1 = __float22bfloat162_rn(make_float2(f4.z, f4.w));
|
||||
reinterpret_cast<__nv_bfloat162*>(ptr)[0] = b2_0;
|
||||
reinterpret_cast<__nv_bfloat162*>(ptr)[1] = b2_1;
|
||||
#else
|
||||
assert(false && "Unsupported on arch < 800");
|
||||
#endif
|
||||
} else if constexpr (std::is_same_v<T, float>) {
|
||||
reinterpret_cast<float4*>(ptr)[0] = f4;
|
||||
} else {
|
||||
static_assert(sizeof(T) == 0, "Unsupported type");
|
||||
}
|
||||
}
|
||||
|
||||
template<class T>
|
||||
__global__ void MRopeKernel(const float* cos,
|
||||
const float* sin,
|
||||
const T* q,
|
||||
const int q_h,
|
||||
const T* k,
|
||||
const int k_h,
|
||||
T* q_embed,
|
||||
T* k_embed,
|
||||
const int first,
|
||||
const int second,
|
||||
const int d,
|
||||
const int qb_stride,
|
||||
const int qs_stride,
|
||||
const int qh_stride,
|
||||
const int kb_stride,
|
||||
const int ks_stride,
|
||||
const int kh_stride,
|
||||
const int qeb_stride,
|
||||
const int qes_stride,
|
||||
const int qeh_stride,
|
||||
const int keb_stride,
|
||||
const int kes_stride,
|
||||
const int keh_stride) {
|
||||
extern __shared__ char cos_sin[];
|
||||
const int half_dim = d / 2;
|
||||
float* cos_smem = reinterpret_cast<float*>(cos_sin);
|
||||
float* sin_smem = cos_smem + half_dim;
|
||||
int b = blockIdx.x;
|
||||
int s = blockIdx.y;
|
||||
|
||||
int64_t offset = gridDim.x * gridDim.y * half_dim;
|
||||
int64_t cos_sin_b_stride = gridDim.y * half_dim;
|
||||
int64_t cos_sin_s_stride = half_dim;
|
||||
#define SIN_GMEM(a, b, c, d) sin[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)]
|
||||
#define COS_GMEM(a, b, c, d) cos[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)]
|
||||
|
||||
for (int i = threadIdx.x; i < half_dim; i += blockDim.x) {
|
||||
if (i < first) {
|
||||
cos_smem[i] = COS_GMEM(0, b, s, i);
|
||||
sin_smem[i] = SIN_GMEM(0, b, s, i);
|
||||
} else if (i < (second + first)) {
|
||||
cos_smem[i] = COS_GMEM(1, b, s, i);
|
||||
sin_smem[i] = SIN_GMEM(1, b, s, i);
|
||||
} else {
|
||||
cos_smem[i] = COS_GMEM(2, b, s, i);
|
||||
sin_smem[i] = SIN_GMEM(2, b, s, i);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
#define Q_GMEM(a, b, c, d) q[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + d]
|
||||
#define K_GMEM(a, b, c, d) k[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + d]
|
||||
#define Q_EMBED_GMEM(a, b, c, d) q_embed[(a) * qeb_stride + (b) * qes_stride + (c) * qeh_stride + d]
|
||||
#define K_EMBED_GMEM(a, b, c, d) k_embed[(a) * keb_stride + (b) * kes_stride + (c) * keh_stride + d]
|
||||
|
||||
for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) {
|
||||
int h_idx = j / half_dim;
|
||||
int base_d = j % half_dim;
|
||||
if (base_d + 3 >= half_dim) continue;
|
||||
|
||||
float4 q_x0 = to_float4<T>(&Q_GMEM(b, s, h_idx, base_d));
|
||||
float4 q_x1 = to_float4<T>(&Q_GMEM(b, s, h_idx, base_d + half_dim));
|
||||
|
||||
float4 cos_vec = to_float4<float>(&cos_smem[base_d]);
|
||||
float4 sin_vec = to_float4<float>(&sin_smem[base_d]);
|
||||
|
||||
float4 q0_rot = make_float4(
|
||||
q_x0.x * cos_vec.x - q_x1.x * sin_vec.x,
|
||||
q_x0.y * cos_vec.y - q_x1.y * sin_vec.y,
|
||||
q_x0.z * cos_vec.z - q_x1.z * sin_vec.z,
|
||||
q_x0.w * cos_vec.w - q_x1.w * sin_vec.w
|
||||
);
|
||||
|
||||
float4 q1_rot = make_float4(
|
||||
q_x1.x * cos_vec.x + q_x0.x * sin_vec.x,
|
||||
q_x1.y * cos_vec.y + q_x0.y * sin_vec.y,
|
||||
q_x1.z * cos_vec.z + q_x0.z * sin_vec.z,
|
||||
q_x1.w * cos_vec.w + q_x0.w * sin_vec.w
|
||||
);
|
||||
|
||||
from_float4<T>(&Q_EMBED_GMEM(b, s, h_idx, base_d), q0_rot);
|
||||
from_float4<T>(&Q_EMBED_GMEM(b, s, h_idx, base_d + half_dim), q1_rot);
|
||||
}
|
||||
|
||||
for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) {
|
||||
int h_idx = j / half_dim;
|
||||
int base_d = j % half_dim;
|
||||
if (base_d + 3 >= half_dim) continue;
|
||||
|
||||
float4 k_x0 = to_float4<T>(&K_GMEM(b, s, h_idx, base_d));
|
||||
float4 k_x1 = to_float4<T>(&K_GMEM(b, s, h_idx, base_d + half_dim));
|
||||
|
||||
float4 cos_vec = to_float4<float>(&cos_smem[base_d]);
|
||||
float4 sin_vec = to_float4<float>(&sin_smem[base_d]);
|
||||
|
||||
float4 k0_rot = make_float4(
|
||||
k_x0.x * cos_vec.x - k_x1.x * sin_vec.x,
|
||||
k_x0.y * cos_vec.y - k_x1.y * sin_vec.y,
|
||||
k_x0.z * cos_vec.z - k_x1.z * sin_vec.z,
|
||||
k_x0.w * cos_vec.w - k_x1.w * sin_vec.w
|
||||
);
|
||||
|
||||
float4 k1_rot = make_float4(
|
||||
k_x1.x * cos_vec.x + k_x0.x * sin_vec.x,
|
||||
k_x1.y * cos_vec.y + k_x0.y * sin_vec.y,
|
||||
k_x1.z * cos_vec.z + k_x0.z * sin_vec.z,
|
||||
k_x1.w * cos_vec.w + k_x0.w * sin_vec.w
|
||||
);
|
||||
|
||||
from_float4<T>(&K_EMBED_GMEM(b, s, h_idx, base_d), k0_rot);
|
||||
from_float4<T>(&K_EMBED_GMEM(b, s, h_idx, base_d + half_dim), k1_rot);
|
||||
}
|
||||
}
|
||||
|
||||
template<class T>
|
||||
__global__ void MRopeInplaceKernel(float* cos, float* sin,
|
||||
T* q, const int q_h, T* k, const int k_h,
|
||||
const int first, const int second,
|
||||
const int d, const int qb_stride, const int qs_stride, const int qh_stride,
|
||||
const int kb_stride, const int ks_stride, const int kh_stride) {
|
||||
extern __shared__ char cos_sin[];
|
||||
const int half_dim = d / 2;
|
||||
float* cos_smem = reinterpret_cast<float*>(cos_sin);
|
||||
float* sin_smem = cos_smem + half_dim;
|
||||
int b = blockIdx.x;
|
||||
int s = blockIdx.y;
|
||||
|
||||
int64_t offset = gridDim.x * gridDim.y * half_dim;
|
||||
int64_t cos_sin_b_stride = gridDim.y * half_dim;
|
||||
int64_t cos_sin_s_stride = half_dim;
|
||||
#define SIN_GMEM(a, b, c, d) sin[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)]
|
||||
#define COS_GMEM(a, b, c, d) cos[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)]
|
||||
|
||||
for (int i = threadIdx.x; i < half_dim; i += blockDim.x) {
|
||||
if (i < first) {
|
||||
cos_smem[i] = COS_GMEM(0, b, s, i);
|
||||
sin_smem[i] = SIN_GMEM(0, b, s, i);
|
||||
} else if (i < (second + first)) {
|
||||
cos_smem[i] = COS_GMEM(1, b, s, i);
|
||||
sin_smem[i] = SIN_GMEM(1, b, s, i);
|
||||
} else {
|
||||
cos_smem[i] = COS_GMEM(2, b, s, i);
|
||||
sin_smem[i] = SIN_GMEM(2, b, s, i);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
#define Q_GMEM(a, b, c, d) q[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + d]
|
||||
#define K_GMEM(a, b, c, d) k[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + d]
|
||||
|
||||
for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) {
|
||||
int h_idx = j / half_dim;
|
||||
int base_d = j % half_dim;
|
||||
if (base_d + 3 >= half_dim) continue;
|
||||
|
||||
float4 q_x0 = to_float4<T>(&Q_GMEM(b, s, h_idx, base_d));
|
||||
float4 q_x1 = to_float4<T>(&Q_GMEM(b, s, h_idx, base_d + half_dim));
|
||||
|
||||
float4 cos_vec = to_float4<float>(&cos_smem[base_d]);
|
||||
float4 sin_vec = to_float4<float>(&sin_smem[base_d]);
|
||||
|
||||
float4 q0_rot = make_float4(
|
||||
q_x0.x * cos_vec.x - q_x1.x * sin_vec.x,
|
||||
q_x0.y * cos_vec.y - q_x1.y * sin_vec.y,
|
||||
q_x0.z * cos_vec.z - q_x1.z * sin_vec.z,
|
||||
q_x0.w * cos_vec.w - q_x1.w * sin_vec.w
|
||||
);
|
||||
|
||||
float4 q1_rot = make_float4(
|
||||
q_x1.x * cos_vec.x + q_x0.x * sin_vec.x,
|
||||
q_x1.y * cos_vec.y + q_x0.y * sin_vec.y,
|
||||
q_x1.z * cos_vec.z + q_x0.z * sin_vec.z,
|
||||
q_x1.w * cos_vec.w + q_x0.w * sin_vec.w
|
||||
);
|
||||
|
||||
from_float4<T>(&Q_GMEM(b, s, h_idx, base_d), q0_rot);
|
||||
from_float4<T>(&Q_GMEM(b, s, h_idx, base_d + half_dim), q1_rot);
|
||||
}
|
||||
|
||||
for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) {
|
||||
int h_idx = j / half_dim;
|
||||
int base_d = j % half_dim;
|
||||
if (base_d + 3 >= half_dim) continue;
|
||||
|
||||
float4 k_x0 = to_float4<T>(&K_GMEM(b, s, h_idx, base_d));
|
||||
float4 k_x1 = to_float4<T>(&K_GMEM(b, s, h_idx, base_d + half_dim));
|
||||
|
||||
float4 cos_vec = to_float4<float>(&cos_smem[base_d]);
|
||||
float4 sin_vec = to_float4<float>(&sin_smem[base_d]);
|
||||
|
||||
float4 k0_rot = make_float4(
|
||||
k_x0.x * cos_vec.x - k_x1.x * sin_vec.x,
|
||||
k_x0.y * cos_vec.y - k_x1.y * sin_vec.y,
|
||||
k_x0.z * cos_vec.z - k_x1.z * sin_vec.z,
|
||||
k_x0.w * cos_vec.w - k_x1.w * sin_vec.w
|
||||
);
|
||||
|
||||
float4 k1_rot = make_float4(
|
||||
k_x1.x * cos_vec.x + k_x0.x * sin_vec.x,
|
||||
k_x1.y * cos_vec.y + k_x0.y * sin_vec.y,
|
||||
k_x1.z * cos_vec.z + k_x0.z * sin_vec.z,
|
||||
k_x1.w * cos_vec.w + k_x0.w * sin_vec.w
|
||||
);
|
||||
|
||||
from_float4<T>(&K_GMEM(b, s, h_idx, base_d), k0_rot);
|
||||
from_float4<T>(&K_GMEM(b, s, h_idx, base_d + half_dim), k1_rot);
|
||||
}
|
||||
}
|
||||
|
||||
template<class T>
|
||||
__global__ void MRopeKernelBackward(
|
||||
const float* cos,
|
||||
const float* sin,
|
||||
const T* grad_q_embed,
|
||||
const T* grad_k_embed,
|
||||
T* grad_q,
|
||||
T* grad_k,
|
||||
const int first, const int second,
|
||||
const int q_h,
|
||||
const int k_h,
|
||||
const int d,
|
||||
const int qb_stride, const int qs_stride, const int qh_stride,
|
||||
const int kb_stride, const int ks_stride, const int kh_stride) {
|
||||
|
||||
extern __shared__ char cos_sin[];
|
||||
const int half_dim = d / 2;
|
||||
float* cos_smem = reinterpret_cast<float*>(cos_sin);
|
||||
float* sin_smem = cos_smem + half_dim;
|
||||
|
||||
int b = blockIdx.x;
|
||||
int s = blockIdx.y;
|
||||
|
||||
int64_t offset = gridDim.x * gridDim.y * half_dim;
|
||||
int64_t cos_sin_b_stride = gridDim.y * half_dim;
|
||||
int64_t cos_sin_s_stride = half_dim;
|
||||
|
||||
#define SIN_GMEM(a, b, c, d) sin[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)]
|
||||
#define COS_GMEM(a, b, c, d) cos[(a) * offset + (b) * cos_sin_b_stride + (c) * cos_sin_s_stride + (d)]
|
||||
|
||||
for (int i = threadIdx.x; i < half_dim; i += blockDim.x) {
|
||||
if (i < first) {
|
||||
cos_smem[i] = COS_GMEM(0, b, s, i);
|
||||
sin_smem[i] = SIN_GMEM(0, b, s, i);
|
||||
} else if (i < (second + first)) {
|
||||
cos_smem[i] = COS_GMEM(1, b, s, i);
|
||||
sin_smem[i] = SIN_GMEM(1, b, s, i);
|
||||
} else {
|
||||
cos_smem[i] = COS_GMEM(2, b, s, i);
|
||||
sin_smem[i] = SIN_GMEM(2, b, s, i);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
#define GQ_EMBED(a, b, c, d) grad_q_embed[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + (d)]
|
||||
#define GK_EMBED(a, b, c, d) grad_k_embed[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + (d)]
|
||||
#define GQ(a, b, c, d) grad_q[(a) * qb_stride + (b) * qs_stride + (c) * qh_stride + (d)]
|
||||
#define GK(a, b, c, d) grad_k[(a) * kb_stride + (b) * ks_stride + (c) * kh_stride + (d)]
|
||||
|
||||
// Process grad_q
|
||||
for (int j = threadIdx.x * 4; j < q_h * half_dim; j += blockDim.x * 4) {
|
||||
int h_idx = j / half_dim;
|
||||
int base_d = j % half_dim;
|
||||
if (base_d + 3 >= half_dim) continue;
|
||||
|
||||
// float path
|
||||
float4 gq0_rot = to_float4<T>(&GQ_EMBED(b, s, h_idx, base_d));
|
||||
float4 gq1_rot = to_float4<T>(&GQ_EMBED(b, s, h_idx, base_d + half_dim));
|
||||
|
||||
float4 cos_vec = to_float4<float>(&cos_smem[base_d]);
|
||||
float4 sin_vec = to_float4<float>(&sin_smem[base_d]);
|
||||
|
||||
float4 gq0 = make_float4(
|
||||
gq0_rot.x * cos_vec.x + gq1_rot.x * sin_vec.x,
|
||||
gq0_rot.y * cos_vec.y + gq1_rot.y * sin_vec.y,
|
||||
gq0_rot.z * cos_vec.z + gq1_rot.z * sin_vec.z,
|
||||
gq0_rot.w * cos_vec.w + gq1_rot.w * sin_vec.w
|
||||
);
|
||||
float4 gq1 = make_float4(
|
||||
gq1_rot.x * cos_vec.x - gq0_rot.x * sin_vec.x,
|
||||
gq1_rot.y * cos_vec.y - gq0_rot.y * sin_vec.y,
|
||||
gq1_rot.z * cos_vec.z - gq0_rot.z * sin_vec.z,
|
||||
gq1_rot.w * cos_vec.w - gq0_rot.w * sin_vec.w
|
||||
);
|
||||
|
||||
from_float4<T>(&GQ(b, s, h_idx, base_d), gq0);
|
||||
from_float4<T>(&GQ(b, s, h_idx, base_d + half_dim), gq1);
|
||||
}
|
||||
|
||||
// Process grad_k
|
||||
for (int j = threadIdx.x * 4; j < k_h * half_dim; j += blockDim.x * 4) {
|
||||
int h_idx = j / half_dim;
|
||||
int base_d = j % half_dim;
|
||||
if (base_d + 3 >= half_dim) continue;
|
||||
|
||||
// float path
|
||||
float4 gk0_rot = to_float4<T>(&GK_EMBED(b, s, h_idx, base_d));
|
||||
float4 gk1_rot = to_float4<T>(&GK_EMBED(b, s, h_idx, base_d + half_dim));
|
||||
|
||||
float4 cos_vec = to_float4<float>(&cos_smem[base_d]);
|
||||
float4 sin_vec = to_float4<float>(&sin_smem[base_d]);
|
||||
|
||||
float4 gk0 = make_float4(
|
||||
gk0_rot.x * cos_vec.x + gk1_rot.x * sin_vec.x,
|
||||
gk0_rot.y * cos_vec.y + gk1_rot.y * sin_vec.y,
|
||||
gk0_rot.z * cos_vec.z + gk1_rot.z * sin_vec.z,
|
||||
gk0_rot.w * cos_vec.w + gk1_rot.w * sin_vec.w
|
||||
);
|
||||
float4 gk1 = make_float4(
|
||||
gk1_rot.x * cos_vec.x - gk0_rot.x * sin_vec.x,
|
||||
gk1_rot.y * cos_vec.y - gk0_rot.y * sin_vec.y,
|
||||
gk1_rot.z * cos_vec.z - gk0_rot.z * sin_vec.z,
|
||||
gk1_rot.w * cos_vec.w - gk0_rot.w * sin_vec.w
|
||||
);
|
||||
|
||||
from_float4<T>(&GK(b, s, h_idx, base_d), gk0);
|
||||
from_float4<T>(&GK(b, s, h_idx, base_d + half_dim), gk1);
|
||||
}
|
||||
}
|
||||
|
||||
void MRope(const at::Tensor& q, // [b, s, h, d]
|
||||
const at::Tensor& k, // [b, s, h_k, d]
|
||||
const at::Tensor& q_embed, // [b, s, h, d]
|
||||
const at::Tensor& k_embed, // [b, s, h_k, d]
|
||||
const at::Tensor& cos, // [3, b, s, d / 2]
|
||||
const at::Tensor& sin, // [3, b, s, d / 2]
|
||||
const int first,
|
||||
const int second
|
||||
) {
|
||||
int Nthreads = 256;
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
ASSERT_CHECK(q.scalar_type() == k.scalar_type());
|
||||
ASSERT_CHECK(q.scalar_type() == q_embed.scalar_type());
|
||||
ASSERT_CHECK(k.scalar_type() == k_embed.scalar_type());
|
||||
ASSERT_CHECK(q_embed.size(0) == q.size(0));
|
||||
ASSERT_CHECK(q_embed.size(1) == q.size(1));
|
||||
ASSERT_CHECK(q_embed.size(2) == q.size(2));
|
||||
ASSERT_CHECK(q_embed.size(3) == q.size(3));
|
||||
ASSERT_CHECK(k_embed.size(0) == k.size(0));
|
||||
ASSERT_CHECK(k_embed.size(1) == k.size(1));
|
||||
ASSERT_CHECK(k_embed.size(2) == k.size(2));
|
||||
ASSERT_CHECK(k_embed.size(3) == k.size(3));
|
||||
ASSERT_CHECK(q.size(3) % 8 == 0);
|
||||
ASSERT_CHECK(cos.size(0) == 3);
|
||||
ASSERT_CHECK(cos.size(3) == q.size(3) / 2);
|
||||
ASSERT_CHECK(cos.scalar_type() == at::kFloat);
|
||||
ASSERT_CHECK(sin.scalar_type() == at::kFloat);
|
||||
ASSERT_CHECK(sin.size(0) == 3);
|
||||
ASSERT_CHECK(q.stride(3) == 1);
|
||||
ASSERT_CHECK(k.stride(3) == 1);
|
||||
ASSERT_CHECK(q_embed.stride(3) == 1);
|
||||
ASSERT_CHECK(k_embed.stride(3) == 1);
|
||||
ASSERT_CHECK(cos.stride(3) == 1);
|
||||
ASSERT_CHECK(sin.stride(3) == 1);
|
||||
ASSERT_CHECK(cos.is_contiguous());
|
||||
ASSERT_CHECK(sin.is_contiguous());
|
||||
|
||||
int64_t batch = q.size(0);
|
||||
int64_t q_head_num = q.size(2);
|
||||
int64_t k_head_num = k.size(2);
|
||||
int64_t seq_len = q.size(1);
|
||||
int64_t dim = q.size(3);
|
||||
int64_t qb_stride = q.stride(0);
|
||||
int64_t kb_stride = k.stride(0);
|
||||
int64_t qs_stride = q.stride(1);
|
||||
int64_t ks_stride = k.stride(1);
|
||||
int64_t qh_stride = q.stride(2);
|
||||
int64_t kh_stride = k.stride(2);
|
||||
int64_t qeb_stride = q_embed.stride(0);
|
||||
int64_t keb_stride = k_embed.stride(0);
|
||||
int64_t qes_stride = q_embed.stride(1);
|
||||
int64_t kes_stride = k_embed.stride(1);
|
||||
int64_t qeh_stride = q_embed.stride(2);
|
||||
int64_t keh_stride = k_embed.stride(2);
|
||||
|
||||
dim3 grid(batch, seq_len);
|
||||
|
||||
if (q.scalar_type() == at::kHalf) {
|
||||
const half* q_data = static_cast<const half*>(q.data_ptr());
|
||||
const half* k_data = static_cast<const half*>(k.data_ptr());
|
||||
half* q_embed_data = static_cast<half*>(q_embed.data_ptr());
|
||||
half* k_embed_data = static_cast<half*>(k_embed.data_ptr());
|
||||
const float* cos_data = static_cast<const float*>(cos.data_ptr());
|
||||
const float* sin_data = static_cast<const float*>(sin.data_ptr());
|
||||
MRopeKernel<half><<<grid, Nthreads, dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
q_data, q_head_num,
|
||||
k_data, k_head_num,
|
||||
q_embed_data, k_embed_data,
|
||||
first, second, dim,
|
||||
qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride,
|
||||
qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride);
|
||||
} else if (q.scalar_type() == at::kBFloat16) {
|
||||
const __nv_bfloat16* q_data = static_cast<const __nv_bfloat16*>(q.data_ptr());
|
||||
const __nv_bfloat16* k_data = static_cast<const __nv_bfloat16*>(k.data_ptr());
|
||||
__nv_bfloat16* q_embed_data = static_cast<__nv_bfloat16*>(q_embed.data_ptr());
|
||||
__nv_bfloat16* k_embed_data = static_cast<__nv_bfloat16*>(k_embed.data_ptr());
|
||||
const float* cos_data = static_cast<const float*>(cos.data_ptr());
|
||||
const float* sin_data = static_cast<const float*>(sin.data_ptr());
|
||||
MRopeKernel<__nv_bfloat16><<<grid, Nthreads, dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
q_data, q_head_num,
|
||||
k_data, k_head_num,
|
||||
q_embed_data, k_embed_data,
|
||||
first, second, dim,
|
||||
qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride,
|
||||
qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride);
|
||||
} else if (q.scalar_type() == at::kFloat) {
|
||||
const float* q_data = static_cast<const float*>(q.data_ptr());
|
||||
const float* k_data = static_cast<const float*>(k.data_ptr());
|
||||
float* q_embed_data = static_cast<float*>(q_embed.data_ptr());
|
||||
float* k_embed_data = static_cast<float*>(k_embed.data_ptr());
|
||||
const float* cos_data = static_cast<const float*>(cos.data_ptr());
|
||||
const float* sin_data = static_cast<const float*>(sin.data_ptr());
|
||||
MRopeKernel<float><<<grid, Nthreads, dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
q_data, q_head_num,
|
||||
k_data, k_head_num,
|
||||
q_embed_data, k_embed_data,
|
||||
first, second, dim,
|
||||
qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride,
|
||||
qeb_stride, qes_stride, qeh_stride, keb_stride, kes_stride, keh_stride
|
||||
);
|
||||
} else {
|
||||
throw std::runtime_error("Unsupported data type for m_rope");
|
||||
}
|
||||
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
void MRopeInplace(const at::Tensor& q, // [b, s, h, d]
|
||||
const at::Tensor& k, // [b, s, h_k, d]
|
||||
const at::Tensor& cos, // [3, b, s, d / 2]
|
||||
const at::Tensor& sin, // [3, b, s, d / 2]
|
||||
const int first,
|
||||
const int second
|
||||
) {
|
||||
int Nthreads = 256;
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
ASSERT_CHECK(q.scalar_type() == k.scalar_type());
|
||||
ASSERT_CHECK(q.size(3) % 8 == 0);
|
||||
ASSERT_CHECK(cos.size(0) == 3);
|
||||
ASSERT_CHECK(cos.size(3) == q.size(3) / 2);
|
||||
ASSERT_CHECK(cos.scalar_type() == at::kFloat);
|
||||
ASSERT_CHECK(sin.scalar_type() == at::kFloat);
|
||||
ASSERT_CHECK(sin.size(0) == 3);
|
||||
ASSERT_CHECK(q.stride(3) == 1);
|
||||
ASSERT_CHECK(k.stride(3) == 1);
|
||||
ASSERT_CHECK(cos.stride(3) == 1);
|
||||
ASSERT_CHECK(sin.stride(3) == 1);
|
||||
ASSERT_CHECK(cos.is_contiguous());
|
||||
ASSERT_CHECK(sin.is_contiguous());
|
||||
|
||||
int64_t batch = q.size(0);
|
||||
int64_t q_head_num = q.size(2);
|
||||
int64_t k_head_num = k.size(2);
|
||||
int64_t seq_len = q.size(1);
|
||||
int64_t dim = q.size(3);
|
||||
int64_t qb_stride = q.stride(0);
|
||||
int64_t kb_stride = k.stride(0);
|
||||
int64_t qs_stride = q.stride(1);
|
||||
int64_t ks_stride = k.stride(1);
|
||||
int64_t qh_stride = q.stride(2);
|
||||
int64_t kh_stride = k.stride(2);
|
||||
|
||||
dim3 grid(batch, seq_len);
|
||||
|
||||
if (q.scalar_type() == at::kHalf) {
|
||||
half* q_data = static_cast<half*>(q.data_ptr());
|
||||
half* k_data = static_cast<half*>(k.data_ptr());
|
||||
float* cos_data = static_cast<float*>(cos.data_ptr());
|
||||
float* sin_data = static_cast<float*>(sin.data_ptr());
|
||||
MRopeInplaceKernel<half><<<grid, Nthreads, dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
q_data, q_head_num,
|
||||
k_data, k_head_num,
|
||||
first, second, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride);
|
||||
} else if (q.scalar_type() == at::kBFloat16) {
|
||||
__nv_bfloat16* q_data = static_cast<__nv_bfloat16*>(q.data_ptr());
|
||||
__nv_bfloat16* k_data = static_cast<__nv_bfloat16*>(k.data_ptr());
|
||||
float* cos_data = static_cast<float*>(cos.data_ptr());
|
||||
float* sin_data = static_cast<float*>(sin.data_ptr());
|
||||
MRopeInplaceKernel<__nv_bfloat16><<<grid, Nthreads, dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
q_data, q_head_num,
|
||||
k_data, k_head_num,
|
||||
first, second, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride);
|
||||
} else if (q.scalar_type() == at::kFloat) {
|
||||
float* q_data = static_cast<float*>(q.data_ptr());
|
||||
float* k_data = static_cast<float*>(k.data_ptr());
|
||||
float* cos_data = static_cast<float*>(cos.data_ptr());
|
||||
float* sin_data = static_cast<float*>(sin.data_ptr());
|
||||
MRopeInplaceKernel<float><<<grid, Nthreads, dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
q_data, q_head_num,
|
||||
k_data, k_head_num,
|
||||
first, second, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride
|
||||
);
|
||||
} else {
|
||||
throw std::runtime_error("Unsupported data type for m_rope");
|
||||
}
|
||||
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
void MRopeBackward(
|
||||
const at::Tensor& grad_q_embed,
|
||||
const at::Tensor& grad_k_embed,
|
||||
const at::Tensor& grad_q, // output: [b, s, q_h, d]
|
||||
const at::Tensor& grad_k, // output: [b, s, k_h, d]
|
||||
const at::Tensor& cos, // [3, b, s, d/2]
|
||||
const at::Tensor& sin, // [3, b, s, d/2]
|
||||
const int first,
|
||||
const int second
|
||||
) {
|
||||
int Nthreads = 256;
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
ASSERT_CHECK(grad_q_embed.scalar_type() == grad_k_embed.scalar_type());
|
||||
ASSERT_CHECK(grad_q.scalar_type() == grad_k.scalar_type());
|
||||
ASSERT_CHECK(grad_q_embed.size(3) == grad_q.size(3));
|
||||
ASSERT_CHECK(grad_k_embed.size(3) == grad_k.size(3));
|
||||
ASSERT_CHECK(grad_q_embed.size(3) % 8 == 0);
|
||||
ASSERT_CHECK(grad_q_embed.stride(3) == 1);
|
||||
ASSERT_CHECK(grad_k_embed.stride(3) == 1);
|
||||
ASSERT_CHECK(grad_q.stride(0) == grad_q_embed.stride(0));
|
||||
ASSERT_CHECK(grad_q.stride(1) == grad_q_embed.stride(1));
|
||||
ASSERT_CHECK(grad_q.stride(2) == grad_q_embed.stride(2));
|
||||
ASSERT_CHECK(grad_q.stride(3) == grad_q_embed.stride(3));
|
||||
ASSERT_CHECK(grad_k.stride(0) == grad_k_embed.stride(0));
|
||||
ASSERT_CHECK(grad_k.stride(1) == grad_k_embed.stride(1));
|
||||
ASSERT_CHECK(grad_k.stride(2) == grad_k_embed.stride(2));
|
||||
ASSERT_CHECK(grad_k.stride(3) == grad_k_embed.stride(3));
|
||||
|
||||
int64_t batch = grad_q_embed.size(0);
|
||||
int64_t seq_len = grad_q_embed.size(1);
|
||||
int64_t q_head_num = grad_q.size(2);
|
||||
int64_t k_head_num = grad_k.size(2);
|
||||
int64_t dim = grad_q_embed.size(3);
|
||||
int64_t qb_stride = grad_q_embed.stride(0);
|
||||
int64_t kb_stride = grad_k_embed.stride(0);
|
||||
int64_t qs_stride = grad_q_embed.stride(1);
|
||||
int64_t ks_stride = grad_k_embed.stride(1);
|
||||
int64_t qh_stride = grad_q_embed.stride(2);
|
||||
int64_t kh_stride = grad_k_embed.stride(2);
|
||||
|
||||
dim3 grid(batch, seq_len);
|
||||
|
||||
if (grad_q_embed.scalar_type() == at::kHalf) {
|
||||
const half* gq_embed = static_cast<const half*>(grad_q_embed.data_ptr());
|
||||
const half* gk_embed = static_cast<const half*>(grad_k_embed.data_ptr());
|
||||
half* gq = static_cast<half*>(grad_q.data_ptr());
|
||||
half* gk = static_cast<half*>(grad_k.data_ptr());
|
||||
const float* cos_data = static_cast<const float*>(cos.data_ptr());
|
||||
const float* sin_data = static_cast<const float*>(sin.data_ptr());
|
||||
MRopeKernelBackward<half><<<grid, Nthreads, dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
gq_embed, gk_embed,
|
||||
gq, gk,
|
||||
first, second,
|
||||
q_head_num, k_head_num, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride
|
||||
);
|
||||
} else if (grad_q_embed.scalar_type() == at::kBFloat16) {
|
||||
const __nv_bfloat16* gq_embed = static_cast<const __nv_bfloat16*>(grad_q_embed.data_ptr());
|
||||
const __nv_bfloat16* gk_embed = static_cast<const __nv_bfloat16*>(grad_k_embed.data_ptr());
|
||||
__nv_bfloat16* gq = static_cast<__nv_bfloat16*>(grad_q.data_ptr());
|
||||
__nv_bfloat16* gk = static_cast<__nv_bfloat16*>(grad_k.data_ptr());
|
||||
const float* cos_data = static_cast<const float*>(cos.data_ptr());
|
||||
const float* sin_data = static_cast<const float*>(sin.data_ptr());
|
||||
MRopeKernelBackward<__nv_bfloat16><<<grid, Nthreads, dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
gq_embed, gk_embed,
|
||||
gq, gk,
|
||||
first, second,
|
||||
q_head_num, k_head_num, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride
|
||||
);
|
||||
} else if (grad_q_embed.scalar_type() == at::kFloat) {
|
||||
const float* gq_embed = static_cast<const float*>(grad_q_embed.data_ptr());
|
||||
const float* gk_embed = static_cast<const float*>(grad_k_embed.data_ptr());
|
||||
float* gq = static_cast<float*>(grad_q.data_ptr());
|
||||
float* gk = static_cast<float*>(grad_k.data_ptr());
|
||||
const float* cos_data = static_cast<const float*>(cos.data_ptr());
|
||||
const float* sin_data = static_cast<const float*>(sin.data_ptr());
|
||||
MRopeKernelBackward<float><<<grid, Nthreads, dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
gq_embed, gk_embed,
|
||||
gq, gk,
|
||||
first, second,
|
||||
q_head_num, k_head_num, dim, qb_stride, qs_stride, qh_stride, kb_stride, ks_stride, kh_stride
|
||||
);
|
||||
} else {
|
||||
throw std::runtime_error("Unsupported data type for m_rope backward");
|
||||
}
|
||||
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
// Pack interface: accepts qkv [bz, seq_len, q_dim + 2*kv_dim], applies mrope
|
||||
// inplace on q and k without Python-side split/view overhead.
|
||||
// Supports GQA (q_num_heads != kv_num_heads). CUDA kernels are not modified.
|
||||
void MRopeInplacePack(const at::Tensor& qkv, // [bz, seq_len, q_dim + 2*kv_dim]
|
||||
const at::Tensor& cos, // [3, bz, seq_len, head_dim/2]
|
||||
const at::Tensor& sin, // [3, bz, seq_len, head_dim/2]
|
||||
int64_t q_num_heads,
|
||||
int64_t kv_num_heads,
|
||||
const int first,
|
||||
const int second) {
|
||||
int Nthreads = 256;
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
ASSERT_CHECK(qkv.dim() == 3);
|
||||
ASSERT_CHECK(qkv.stride(2) == 1);
|
||||
ASSERT_CHECK(cos.scalar_type() == at::kFloat);
|
||||
ASSERT_CHECK(sin.scalar_type() == at::kFloat);
|
||||
ASSERT_CHECK(cos.size(0) == 3);
|
||||
ASSERT_CHECK(sin.size(0) == 3);
|
||||
ASSERT_CHECK(cos.is_contiguous());
|
||||
ASSERT_CHECK(sin.is_contiguous());
|
||||
|
||||
int64_t bz = qkv.size(0);
|
||||
int64_t seq_len = qkv.size(1);
|
||||
int64_t head_dim = cos.size(3) * 2; // cos: [3, bz, seq, half_dim]
|
||||
int64_t q_dim = q_num_heads * head_dim;
|
||||
int64_t kv_dim = kv_num_heads * head_dim;
|
||||
ASSERT_CHECK(qkv.size(2) == q_dim + 2 * kv_dim);
|
||||
ASSERT_CHECK(head_dim % 8 == 0);
|
||||
|
||||
// q at offset 0, k at offset q_dim; logical [bz, seq, *_num_heads, head_dim]
|
||||
// with strides [qkv.stride(0), qkv.stride(1), head_dim, 1]
|
||||
int64_t qkv_b_stride = qkv.stride(0);
|
||||
int64_t qkv_s_stride = qkv.stride(1);
|
||||
|
||||
dim3 grid(bz, seq_len);
|
||||
|
||||
if (qkv.scalar_type() == at::kHalf) {
|
||||
half* q_data = static_cast<half*>(qkv.data_ptr());
|
||||
half* k_data = q_data + q_dim;
|
||||
float* cos_data = static_cast<float*>(cos.data_ptr());
|
||||
float* sin_data = static_cast<float*>(sin.data_ptr());
|
||||
MRopeInplaceKernel<half><<<grid, Nthreads, head_dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
q_data, q_num_heads, k_data, kv_num_heads,
|
||||
first, second, head_dim,
|
||||
qkv_b_stride, qkv_s_stride, head_dim,
|
||||
qkv_b_stride, qkv_s_stride, head_dim);
|
||||
} else if (qkv.scalar_type() == at::kBFloat16) {
|
||||
__nv_bfloat16* q_data = static_cast<__nv_bfloat16*>(qkv.data_ptr());
|
||||
__nv_bfloat16* k_data = q_data + q_dim;
|
||||
float* cos_data = static_cast<float*>(cos.data_ptr());
|
||||
float* sin_data = static_cast<float*>(sin.data_ptr());
|
||||
MRopeInplaceKernel<__nv_bfloat16><<<grid, Nthreads, head_dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
q_data, q_num_heads, k_data, kv_num_heads,
|
||||
first, second, head_dim,
|
||||
qkv_b_stride, qkv_s_stride, head_dim,
|
||||
qkv_b_stride, qkv_s_stride, head_dim);
|
||||
} else if (qkv.scalar_type() == at::kFloat) {
|
||||
float* q_data = static_cast<float*>(qkv.data_ptr());
|
||||
float* k_data = q_data + q_dim;
|
||||
float* cos_data = static_cast<float*>(cos.data_ptr());
|
||||
float* sin_data = static_cast<float*>(sin.data_ptr());
|
||||
MRopeInplaceKernel<float><<<grid, Nthreads, head_dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
q_data, q_num_heads, k_data, kv_num_heads,
|
||||
first, second, head_dim,
|
||||
qkv_b_stride, qkv_s_stride, head_dim,
|
||||
qkv_b_stride, qkv_s_stride, head_dim);
|
||||
} else {
|
||||
throw std::runtime_error("Unsupported data type for m_rope_inplace_pack");
|
||||
}
|
||||
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
// Pack backward interface: accepts dqkv [bz, seq_len, q_dim + 2*kv_dim] and
|
||||
// applies the inverse mrope transform inplace on the dq and dk slices. The dv
|
||||
// slice is left unchanged.
|
||||
void MRopeInplacePackBackward(const at::Tensor& dqkv, // [bz, seq_len, q_dim + 2*kv_dim]
|
||||
const at::Tensor& cos, // [3, bz, seq_len, head_dim/2]
|
||||
const at::Tensor& sin, // [3, bz, seq_len, head_dim/2]
|
||||
int64_t q_num_heads,
|
||||
int64_t kv_num_heads,
|
||||
const int first,
|
||||
const int second) {
|
||||
int Nthreads = 256;
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
ASSERT_CHECK(dqkv.dim() == 3);
|
||||
ASSERT_CHECK(dqkv.stride(2) == 1);
|
||||
ASSERT_CHECK(cos.scalar_type() == at::kFloat);
|
||||
ASSERT_CHECK(sin.scalar_type() == at::kFloat);
|
||||
ASSERT_CHECK(cos.size(0) == 3);
|
||||
ASSERT_CHECK(sin.size(0) == 3);
|
||||
ASSERT_CHECK(cos.is_contiguous());
|
||||
ASSERT_CHECK(sin.is_contiguous());
|
||||
|
||||
int64_t bz = dqkv.size(0);
|
||||
int64_t seq_len = dqkv.size(1);
|
||||
int64_t head_dim = cos.size(3) * 2; // cos: [3, bz, seq, half_dim]
|
||||
int64_t q_dim = q_num_heads * head_dim;
|
||||
int64_t kv_dim = kv_num_heads * head_dim;
|
||||
ASSERT_CHECK(dqkv.size(2) == q_dim + 2 * kv_dim);
|
||||
ASSERT_CHECK(head_dim % 8 == 0);
|
||||
|
||||
int64_t dqkv_b_stride = dqkv.stride(0);
|
||||
int64_t dqkv_s_stride = dqkv.stride(1);
|
||||
|
||||
dim3 grid(bz, seq_len);
|
||||
|
||||
if (dqkv.scalar_type() == at::kHalf) {
|
||||
half* dq_data = static_cast<half*>(dqkv.data_ptr());
|
||||
half* dk_data = dq_data + q_dim;
|
||||
const float* cos_data = static_cast<const float*>(cos.data_ptr());
|
||||
const float* sin_data = static_cast<const float*>(sin.data_ptr());
|
||||
MRopeKernelBackward<half><<<grid, Nthreads, head_dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
dq_data, dk_data,
|
||||
dq_data, dk_data,
|
||||
first, second,
|
||||
q_num_heads, kv_num_heads, head_dim,
|
||||
dqkv_b_stride, dqkv_s_stride, head_dim,
|
||||
dqkv_b_stride, dqkv_s_stride, head_dim);
|
||||
} else if (dqkv.scalar_type() == at::kBFloat16) {
|
||||
__nv_bfloat16* dq_data = static_cast<__nv_bfloat16*>(dqkv.data_ptr());
|
||||
__nv_bfloat16* dk_data = dq_data + q_dim;
|
||||
const float* cos_data = static_cast<const float*>(cos.data_ptr());
|
||||
const float* sin_data = static_cast<const float*>(sin.data_ptr());
|
||||
MRopeKernelBackward<__nv_bfloat16><<<grid, Nthreads, head_dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
dq_data, dk_data,
|
||||
dq_data, dk_data,
|
||||
first, second,
|
||||
q_num_heads, kv_num_heads, head_dim,
|
||||
dqkv_b_stride, dqkv_s_stride, head_dim,
|
||||
dqkv_b_stride, dqkv_s_stride, head_dim);
|
||||
} else if (dqkv.scalar_type() == at::kFloat) {
|
||||
float* dq_data = static_cast<float*>(dqkv.data_ptr());
|
||||
float* dk_data = dq_data + q_dim;
|
||||
const float* cos_data = static_cast<const float*>(cos.data_ptr());
|
||||
const float* sin_data = static_cast<const float*>(sin.data_ptr());
|
||||
MRopeKernelBackward<float><<<grid, Nthreads, head_dim * sizeof(float), stream>>>(
|
||||
cos_data, sin_data,
|
||||
dq_data, dk_data,
|
||||
dq_data, dk_data,
|
||||
first, second,
|
||||
q_num_heads, kv_num_heads, head_dim,
|
||||
dqkv_b_stride, dqkv_s_stride, head_dim,
|
||||
dqkv_b_stride, dqkv_s_stride, head_dim);
|
||||
} else {
|
||||
throw std::runtime_error("Unsupported data type for m_rope_inplace_pack_bwd");
|
||||
}
|
||||
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
|
||||
} // wallx_cuda_m_rope
|
||||
@@ -0,0 +1,581 @@
|
||||
// Re-enable CUDA half operators (PyTorch disables them)
|
||||
#undef __CUDA_NO_HALF_OPERATORS__
|
||||
#undef __CUDA_NO_HALF_CONVERSIONS__
|
||||
#undef __CUDA_NO_HALF2_OPERATORS__
|
||||
#undef __CUDA_NO_BFLOAT16_CONVERSIONS__
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cub/cub.cuh>
|
||||
#include "../common/cuda_utils.h"
|
||||
#include "cuda_runtime.h"
|
||||
#include "device_launch_parameters.h"
|
||||
|
||||
// Global FP8 guard: use CUDA_VERSION (defined in both host and device passes),
|
||||
// not __CUDA_ARCH__ (device-pass only - leaves host dispatchers unable to see
|
||||
// FP8 typedefs/branches). FP8 storage types arrived in CUDA 11.8.
|
||||
#if defined(CUDA_VERSION) && (CUDA_VERSION >= 11800)
|
||||
#include <cuda_fp8.h>
|
||||
#define XCOMPUTE_HAS_FP8 1
|
||||
#else
|
||||
#define XCOMPUTE_HAS_FP8 0
|
||||
#endif
|
||||
|
||||
namespace wallx_cuda_permute {
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// CUTLASS-free helpers: replace cutlass::Array, NumericArrayConverter,
|
||||
// arch::global_load, and type aliases with standard CUDA equivalents.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Vectorized array: replaces cutlass::Array<T, N>
|
||||
// __align__(16) ensures safe reinterpret_cast<float4*> in vec_load/store.
|
||||
template <typename T, int N>
|
||||
struct __align__(16) VecArray {
|
||||
T data[N];
|
||||
__device__ __forceinline__ T& at(int i) { return data[i]; }
|
||||
__device__ __forceinline__ const T& at(int i) const { return data[i]; }
|
||||
__device__ __forceinline__ T* raw() { return data; }
|
||||
__device__ __forceinline__ const T* raw() const { return data; }
|
||||
__device__ __forceinline__ void clear() {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < N; i++) data[i] = T(0);
|
||||
}
|
||||
};
|
||||
|
||||
// Scalar multiply: VecArray * scalar
|
||||
template <typename T, int N>
|
||||
__device__ __forceinline__ VecArray<T, N> operator*(const VecArray<T, N>& a, T s) {
|
||||
VecArray<T, N> r;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < N; i++) r.data[i] = a.data[i] * s;
|
||||
return r;
|
||||
}
|
||||
|
||||
// Element-wise convert: replaces cutlass::NumericArrayConverter
|
||||
template <typename To, typename From, int N>
|
||||
struct ArrayConverter {
|
||||
__device__ __forceinline__ VecArray<To, N> operator()(const VecArray<From, N>& src) const {
|
||||
VecArray<To, N> dst;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < N; i++) dst.data[i] = To(src.data[i]);
|
||||
return dst;
|
||||
}
|
||||
};
|
||||
|
||||
// Vectorized global load: replaces cutlass::arch::global_load
|
||||
template <typename Fragment>
|
||||
__device__ __forceinline__ void vec_load(Fragment& frag, const void* ptr) {
|
||||
static_assert(sizeof(Fragment) == sizeof(float4) || sizeof(Fragment) == 16,
|
||||
"vec_load expects 16-byte fragment");
|
||||
*reinterpret_cast<float4*>(&frag) = *reinterpret_cast<const float4*>(ptr);
|
||||
}
|
||||
|
||||
// Type aliases: replace cutlass types with standard CUDA types
|
||||
using half_t = __half;
|
||||
using bfloat16_t = __nv_bfloat16;
|
||||
#if XCOMPUTE_HAS_FP8
|
||||
using float_e5m2_t = __nv_fp8_e5m2;
|
||||
using float_e4m3_t = __nv_fp8_e4m3; // fn variant; CUDA's __nv_fp8_e4m3 IS the "fn" form
|
||||
#endif
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Kernels (unchanged logic, only CUTLASS types/calls replaced)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
static __global__ void moe_permute_topK_row_map(
|
||||
const int *sorted_row_id,
|
||||
int *row_id_map,
|
||||
const int num_rows,
|
||||
const int num_topK,
|
||||
const int num_out_tokens)
|
||||
{
|
||||
const int bid = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
const int idx = bid * blockDim.x + tid;
|
||||
|
||||
if (idx >= num_rows * num_topK)
|
||||
return;
|
||||
|
||||
int source_row = sorted_row_id[idx];
|
||||
int source_token_id = source_row / num_topK;
|
||||
int source_topK_id = source_row % num_topK;
|
||||
|
||||
if (idx >= num_out_tokens)
|
||||
{
|
||||
row_id_map[source_topK_id * num_rows + source_token_id] = -1;
|
||||
}
|
||||
else
|
||||
{
|
||||
row_id_map[source_topK_id * num_rows + source_token_id] = idx;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename TCompute, int kElementsPerAccess, bool hasProb>
|
||||
__global__ void moe_recover_topK_kernel(const T *input,
|
||||
T *unpermuted_output,
|
||||
const int *row_id_map,
|
||||
const float *prob,
|
||||
const int num_rows,
|
||||
const int num_topK,
|
||||
const int num_cols)
|
||||
{
|
||||
extern __shared__ int8_t s_mem[];
|
||||
TCompute *s_prob = reinterpret_cast<TCompute *>(s_mem);
|
||||
|
||||
using FragLS = VecArray<T, kElementsPerAccess>;
|
||||
using FragC = VecArray<TCompute, kElementsPerAccess>;
|
||||
|
||||
ArrayConverter<TCompute, T, kElementsPerAccess> src_converter;
|
||||
ArrayConverter<T, TCompute, kElementsPerAccess> dst_converter;
|
||||
|
||||
const int source_token = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
|
||||
if (hasProb)
|
||||
{
|
||||
for (int i = tid; i < num_topK; i += blockDim.x * blockDim.y)
|
||||
{
|
||||
s_prob[i] = TCompute(prob[source_token * num_topK + i]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess)
|
||||
{
|
||||
FragLS frag_load_store;
|
||||
FragC frag_elem;
|
||||
FragC frag_sum;
|
||||
|
||||
int source_row = row_id_map[source_token];
|
||||
|
||||
if (source_row != -1)
|
||||
{
|
||||
const T *source_row_ptr = input + source_row * num_cols;
|
||||
|
||||
vec_load(frag_load_store, source_row_ptr + i);
|
||||
frag_sum = src_converter(frag_load_store);
|
||||
|
||||
if (hasProb)
|
||||
{
|
||||
frag_sum = frag_sum * s_prob[0];
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
frag_sum.clear();
|
||||
}
|
||||
|
||||
for (int k = 1; k < num_topK; k++)
|
||||
{
|
||||
source_row = row_id_map[k * num_rows + source_token];
|
||||
|
||||
if (source_row == -1)
|
||||
continue;
|
||||
|
||||
const T *source_row_ptr = input + source_row * num_cols;
|
||||
|
||||
vec_load(frag_load_store, source_row_ptr + i);
|
||||
frag_elem = src_converter(frag_load_store);
|
||||
|
||||
if (hasProb)
|
||||
{
|
||||
frag_elem = frag_elem * s_prob[k];
|
||||
}
|
||||
|
||||
for (int e = 0; e < kElementsPerAccess; e++)
|
||||
{
|
||||
frag_sum.at(e) = frag_sum.at(e) + frag_elem.at(e);
|
||||
}
|
||||
}
|
||||
|
||||
T *dest_row_ptr = unpermuted_output + source_token * num_cols;
|
||||
frag_load_store = dst_converter(frag_sum);
|
||||
*(float4 *)(dest_row_ptr + i) = *(float4 *)(frag_load_store.raw());
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T,
|
||||
typename TCompute,
|
||||
int kElementsPerAccess,
|
||||
int topKTile,
|
||||
bool hasProb>
|
||||
__global__ void moe_permute_topK_kernel(const T *input_bwd,
|
||||
const T *input_fwd,
|
||||
T *act_grad,
|
||||
const float *prob,
|
||||
float *prob_grad,
|
||||
const int *row_id_map,
|
||||
const int num_rows,
|
||||
const int num_topK,
|
||||
const int num_cols)
|
||||
{
|
||||
extern __shared__ int8_t s_mem[];
|
||||
TCompute *s_prob = reinterpret_cast<TCompute *>(s_mem);
|
||||
|
||||
using FragLS = VecArray<T, kElementsPerAccess>;
|
||||
using FragC = VecArray<TCompute, kElementsPerAccess>;
|
||||
|
||||
ArrayConverter<TCompute, T, kElementsPerAccess> src_converter;
|
||||
ArrayConverter<T, TCompute, kElementsPerAccess> dst_converter;
|
||||
|
||||
const int source_token = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
|
||||
if (hasProb)
|
||||
{
|
||||
for (int i = tid; i < num_topK; i += blockDim.x)
|
||||
{
|
||||
s_prob[i] = TCompute(prob[source_token * num_topK + i]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
float accum[topKTile] = {0.0f};
|
||||
FragLS frag_load_store;
|
||||
|
||||
const T *source_row_ptr = input_bwd + source_token * num_cols;
|
||||
for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess)
|
||||
{
|
||||
vec_load(frag_load_store, source_row_ptr + i);
|
||||
FragC frag_src = src_converter(frag_load_store);
|
||||
|
||||
int index = source_token;
|
||||
|
||||
for (int k = 0; k < topKTile; k++)
|
||||
{
|
||||
if (k == num_topK)
|
||||
break;
|
||||
|
||||
int dest_row = row_id_map[index];
|
||||
index += num_rows;
|
||||
|
||||
if (dest_row == -1)
|
||||
continue;
|
||||
|
||||
if (hasProb)
|
||||
{
|
||||
frag_load_store = dst_converter(frag_src * s_prob[k]);
|
||||
}
|
||||
else
|
||||
{
|
||||
frag_load_store = dst_converter(frag_src);
|
||||
}
|
||||
|
||||
T *dest_row_ptr = act_grad + dest_row * num_cols;
|
||||
*(float4 *)(dest_row_ptr + i) = *(float4 *)(frag_load_store.raw());
|
||||
|
||||
if (hasProb)
|
||||
{
|
||||
const T *input_fwd_ptr = input_fwd + dest_row * num_cols;
|
||||
vec_load(frag_load_store, input_fwd_ptr + i);
|
||||
FragC frag_input_fwd = src_converter(frag_load_store);
|
||||
|
||||
for (int e = 0; e < kElementsPerAccess; e++)
|
||||
{
|
||||
accum[k] += float(frag_src.at(e) * frag_input_fwd.at(e));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (hasProb)
|
||||
{
|
||||
for (int k = 0; k < topKTile; k++)
|
||||
{
|
||||
if (k == num_topK)
|
||||
break;
|
||||
|
||||
for (int mask = 16; mask > 0; mask /= 2)
|
||||
{
|
||||
accum[k] = accum[k] + __shfl_xor_sync(0xffffffff, accum[k], mask, 32);
|
||||
}
|
||||
}
|
||||
|
||||
if (tid == 0)
|
||||
{
|
||||
for (int k = 0; k < topKTile; k++)
|
||||
{
|
||||
if (k == num_topK)
|
||||
break;
|
||||
prob_grad[source_token * num_topK + k] = accum[k];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Launcher (unchanged)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
template <typename T, typename TCompute, bool FWD, int kElementsPerAccess>
|
||||
void moe_permute_topK_kernel_launcher(
|
||||
const T *input,
|
||||
T *output,
|
||||
const int *sorted_row_id,
|
||||
int *row_id_map,
|
||||
const float *prob,
|
||||
const int num_rows,
|
||||
const int num_topK,
|
||||
const int num_cols,
|
||||
const int num_out_tokens,
|
||||
cudaStream_t stream,
|
||||
float *prob_grad = nullptr,
|
||||
const T *input_fwd = nullptr)
|
||||
{
|
||||
if (FWD)
|
||||
{
|
||||
if (prob_grad == nullptr)
|
||||
{
|
||||
int threads = 64;
|
||||
int blocks = (num_rows * num_topK + threads - 1) / threads;
|
||||
moe_permute_topK_row_map<<<blocks, threads, 0, stream>>>(
|
||||
sorted_row_id, row_id_map, num_rows, num_topK, num_out_tokens);
|
||||
|
||||
blocks = num_rows;
|
||||
threads = std::min(num_cols / kElementsPerAccess, 1024);
|
||||
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 128, false><<<blocks, threads, 0, stream>>>(
|
||||
input, nullptr, output, nullptr, nullptr, row_id_map,
|
||||
num_rows, num_topK, num_cols);
|
||||
}
|
||||
else
|
||||
{
|
||||
int blocks = num_rows;
|
||||
int threads = 32;
|
||||
size_t smem_bytes = num_topK * sizeof(TCompute);
|
||||
|
||||
if (num_topK == 1)
|
||||
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 1, false><<<blocks, threads, 0, stream>>>(
|
||||
input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols);
|
||||
else if (num_topK <= 8)
|
||||
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 8, true><<<blocks, threads, smem_bytes, stream>>>(
|
||||
input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols);
|
||||
else if (num_topK <= 16)
|
||||
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 16, true><<<blocks, threads, smem_bytes, stream>>>(
|
||||
input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols);
|
||||
else if (num_topK <= 32)
|
||||
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 32, true><<<blocks, threads, smem_bytes, stream>>>(
|
||||
input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols);
|
||||
else if (num_topK <= 64)
|
||||
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 64, true><<<blocks, threads, smem_bytes, stream>>>(
|
||||
input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols);
|
||||
else if (num_topK <= 128)
|
||||
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 128, true><<<blocks, threads, smem_bytes, stream>>>(
|
||||
input, input_fwd, output, prob, prob_grad, row_id_map, num_rows, num_topK, num_cols);
|
||||
else
|
||||
throw std::runtime_error("num_topK cannot exceed 128.");
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
int blocks = num_rows;
|
||||
int threads = std::min(num_cols / kElementsPerAccess, 1024);
|
||||
size_t smem_bytes = num_topK * sizeof(TCompute);
|
||||
|
||||
if (num_topK == 1)
|
||||
moe_recover_topK_kernel<T, TCompute, kElementsPerAccess, false><<<blocks, threads, smem_bytes, stream>>>(
|
||||
input, output, row_id_map, prob, num_rows, num_topK, num_cols);
|
||||
else if (prob == nullptr)
|
||||
moe_recover_topK_kernel<T, TCompute, kElementsPerAccess, false><<<blocks, threads, smem_bytes, stream>>>(
|
||||
input, output, row_id_map, prob, num_rows, num_topK, num_cols);
|
||||
else
|
||||
moe_recover_topK_kernel<T, TCompute, kElementsPerAccess, true><<<blocks, threads, smem_bytes, stream>>>(
|
||||
input, output, row_id_map, prob, num_rows, num_topK, num_cols);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Host-side ops
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
void MoePermuteTopKOp(
|
||||
const at::Tensor& input,
|
||||
const at::Tensor& indices,
|
||||
const at::Tensor& sorted_indices,
|
||||
const at::Tensor& row_id,
|
||||
const at::Tensor& sorted_row_id,
|
||||
const at::Tensor& temp_storage,
|
||||
const at::Tensor& permuted_output,
|
||||
const at::Tensor& row_id_map,
|
||||
int64_t num_out_tokens,
|
||||
int64_t max_expanded_token_num
|
||||
) {
|
||||
ASSERT_CHECK(input.size(0) == indices.size(0));
|
||||
const int num_tokens = input.size(0);
|
||||
const int num_cols = input.size(1);
|
||||
const int num_topK = indices.size(1);
|
||||
|
||||
int *indices_ptr = static_cast<int*>(indices.data_ptr());
|
||||
int *sorted_indices_ptr = static_cast<int*>(sorted_indices.data_ptr());
|
||||
int *row_id_ptr = static_cast<int*>(row_id.data_ptr());
|
||||
int *sorted_row_id_ptr = static_cast<int*>(sorted_row_id.data_ptr());
|
||||
|
||||
void *d_temp_storage = static_cast<void*>(temp_storage.data_ptr());
|
||||
size_t temp_storage_bytes = std::numeric_limits<size_t>::max();
|
||||
|
||||
cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes,
|
||||
indices_ptr, sorted_indices_ptr,
|
||||
row_id_ptr, sorted_row_id_ptr, num_tokens * num_topK);
|
||||
|
||||
num_out_tokens = (num_out_tokens > 0) ? num_out_tokens : num_tokens * num_topK;
|
||||
int *row_id_map_ptr = static_cast<int*>(row_id_map.data_ptr());
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
if (input.scalar_type() == at::kFloat) {
|
||||
float *input_ptr = static_cast<float*>(input.data_ptr());
|
||||
float *out_ptr = static_cast<float*>(permuted_output.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<float, float, true, 4>(
|
||||
input_ptr, out_ptr, sorted_row_id_ptr, row_id_map_ptr,
|
||||
nullptr, num_tokens, num_topK, num_cols, num_out_tokens, stream);
|
||||
} else if (input.scalar_type() == at::kHalf) {
|
||||
half_t *input_ptr = static_cast<half_t*>(input.data_ptr());
|
||||
half_t *out_ptr = static_cast<half_t*>(permuted_output.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<half_t, half_t, true, 8>(
|
||||
input_ptr, out_ptr, sorted_row_id_ptr, row_id_map_ptr,
|
||||
nullptr, num_tokens, num_topK, num_cols, num_out_tokens, stream);
|
||||
} else if (input.scalar_type() == at::kBFloat16) {
|
||||
bfloat16_t *input_ptr = static_cast<bfloat16_t*>(input.data_ptr());
|
||||
bfloat16_t *out_ptr = static_cast<bfloat16_t*>(permuted_output.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<bfloat16_t, bfloat16_t, true, 8>(
|
||||
input_ptr, out_ptr, sorted_row_id_ptr, row_id_map_ptr,
|
||||
nullptr, num_tokens, num_topK, num_cols, num_out_tokens, stream);
|
||||
#if XCOMPUTE_HAS_FP8
|
||||
} else if (input.scalar_type() == at::kFloat8_e5m2) {
|
||||
float_e5m2_t *input_ptr = static_cast<float_e5m2_t*>(input.data_ptr());
|
||||
float_e5m2_t *out_ptr = static_cast<float_e5m2_t*>(permuted_output.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<float_e5m2_t, half_t, true, 16>(
|
||||
input_ptr, out_ptr, sorted_row_id_ptr, row_id_map_ptr,
|
||||
nullptr, num_tokens, num_topK, num_cols, num_out_tokens, stream);
|
||||
} else if (input.scalar_type() == at::kFloat8_e4m3fn) {
|
||||
float_e4m3_t *input_ptr = static_cast<float_e4m3_t*>(input.data_ptr());
|
||||
float_e4m3_t *out_ptr = static_cast<float_e4m3_t*>(permuted_output.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<float_e4m3_t, half_t, true, 16>(
|
||||
input_ptr, out_ptr, sorted_row_id_ptr, row_id_map_ptr,
|
||||
nullptr, num_tokens, num_topK, num_cols, num_out_tokens, stream);
|
||||
#endif
|
||||
} else {
|
||||
throw std::runtime_error("Unsupported data type for MoePermuteTopKOp");
|
||||
}
|
||||
}
|
||||
|
||||
void MoeRecoverTopKOp(
|
||||
const at::Tensor& input,
|
||||
const at::Tensor& row_id_map,
|
||||
const c10::optional<at::Tensor>& prob,
|
||||
const at::Tensor& unpermuted_output,
|
||||
int64_t num_tokens,
|
||||
int64_t num_topK
|
||||
) {
|
||||
const int num_cols = input.size(1);
|
||||
int *row_id_map_ptr = static_cast<int*>(row_id_map.data_ptr());
|
||||
float *prob_ptr = (prob.has_value()) ? static_cast<float*>(prob.value().data_ptr()) : nullptr;
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
if (input.scalar_type() == at::kFloat) {
|
||||
float *in_ptr = static_cast<float*>(input.data_ptr());
|
||||
float *out_ptr = static_cast<float*>(unpermuted_output.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<float, float, false, 4>(
|
||||
in_ptr, out_ptr, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream);
|
||||
} else if (input.scalar_type() == at::kHalf) {
|
||||
half_t *in_ptr = static_cast<half_t*>(input.data_ptr());
|
||||
half_t *out_ptr = static_cast<half_t*>(unpermuted_output.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<half_t, half_t, false, 8>(
|
||||
in_ptr, out_ptr, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream);
|
||||
} else if (input.scalar_type() == at::kBFloat16) {
|
||||
bfloat16_t *in_ptr = static_cast<bfloat16_t*>(input.data_ptr());
|
||||
bfloat16_t *out_ptr = static_cast<bfloat16_t*>(unpermuted_output.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<bfloat16_t, bfloat16_t, false, 8>(
|
||||
in_ptr, out_ptr, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream);
|
||||
#if XCOMPUTE_HAS_FP8
|
||||
} else if (input.scalar_type() == at::kFloat8_e5m2) {
|
||||
float_e5m2_t *in_ptr = static_cast<float_e5m2_t*>(input.data_ptr());
|
||||
float_e5m2_t *out_ptr = static_cast<float_e5m2_t*>(unpermuted_output.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<float_e5m2_t, half_t, false, 16>(
|
||||
in_ptr, out_ptr, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream);
|
||||
} else if (input.scalar_type() == at::kFloat8_e4m3fn) {
|
||||
float_e4m3_t *in_ptr = static_cast<float_e4m3_t*>(input.data_ptr());
|
||||
float_e4m3_t *out_ptr = static_cast<float_e4m3_t*>(unpermuted_output.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<float_e4m3_t, half_t, false, 16>(
|
||||
in_ptr, out_ptr, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream);
|
||||
#endif
|
||||
} else {
|
||||
throw std::runtime_error("Unsupported data type for MoeRecoverTopKOp");
|
||||
}
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
void MoeRecoverTopKBwdOp(
|
||||
const at::Tensor& input_bwd,
|
||||
const at::Tensor& input_fwd,
|
||||
const at::Tensor& row_id_map,
|
||||
const at::Tensor& prob,
|
||||
const at::Tensor& act_grad,
|
||||
const at::Tensor& prob_grad
|
||||
) {
|
||||
const int num_tokens = prob.size(0);
|
||||
const int num_topK = prob.size(1);
|
||||
const int num_cols = input_bwd.size(1);
|
||||
int *row_id_map_ptr = static_cast<int*>(row_id_map.data_ptr());
|
||||
float *prob_ptr = static_cast<float*>(prob.data_ptr());
|
||||
float *prob_grad_ptr = static_cast<float*>(prob_grad.data_ptr());
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
if (input_bwd.scalar_type() == at::kFloat) {
|
||||
float *bwd = static_cast<float*>(input_bwd.data_ptr());
|
||||
float *fwd = static_cast<float*>(input_fwd.data_ptr());
|
||||
float *grad = static_cast<float*>(act_grad.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<float, float, true, 4>(
|
||||
bwd, grad, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream, prob_grad_ptr, fwd);
|
||||
} else if (input_bwd.scalar_type() == at::kHalf) {
|
||||
half_t *bwd = static_cast<half_t*>(input_bwd.data_ptr());
|
||||
half_t *fwd = static_cast<half_t*>(input_fwd.data_ptr());
|
||||
half_t *grad = static_cast<half_t*>(act_grad.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<half_t, half_t, true, 8>(
|
||||
bwd, grad, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream, prob_grad_ptr, fwd);
|
||||
} else if (input_bwd.scalar_type() == at::kBFloat16) {
|
||||
bfloat16_t *bwd = static_cast<bfloat16_t*>(input_bwd.data_ptr());
|
||||
bfloat16_t *fwd = static_cast<bfloat16_t*>(input_fwd.data_ptr());
|
||||
bfloat16_t *grad = static_cast<bfloat16_t*>(act_grad.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<bfloat16_t, bfloat16_t, true, 8>(
|
||||
bwd, grad, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream, prob_grad_ptr, fwd);
|
||||
#if XCOMPUTE_HAS_FP8
|
||||
} else if (input_bwd.scalar_type() == at::kFloat8_e5m2) {
|
||||
float_e5m2_t *bwd = static_cast<float_e5m2_t*>(input_bwd.data_ptr());
|
||||
float_e5m2_t *fwd = static_cast<float_e5m2_t*>(input_fwd.data_ptr());
|
||||
float_e5m2_t *grad = static_cast<float_e5m2_t*>(act_grad.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<float_e5m2_t, half_t, true, 16>(
|
||||
bwd, grad, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream, prob_grad_ptr, fwd);
|
||||
} else if (input_bwd.scalar_type() == at::kFloat8_e4m3fn) {
|
||||
float_e4m3_t *bwd = static_cast<float_e4m3_t*>(input_bwd.data_ptr());
|
||||
float_e4m3_t *fwd = static_cast<float_e4m3_t*>(input_fwd.data_ptr());
|
||||
float_e4m3_t *grad = static_cast<float_e4m3_t*>(act_grad.data_ptr());
|
||||
moe_permute_topK_kernel_launcher<float_e4m3_t, half_t, true, 16>(
|
||||
bwd, grad, nullptr, row_id_map_ptr, prob_ptr, num_tokens, num_topK, num_cols, 0, stream, prob_grad_ptr, fwd);
|
||||
#endif
|
||||
} else {
|
||||
throw std::runtime_error("Unsupported data type for MoeRecoverTopKBwdOp");
|
||||
}
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
size_t CubSortPairGetStorageBytes(int64_t num_items){
|
||||
size_t temp_storage_bytes = 0;
|
||||
int *temp_ptr = nullptr;
|
||||
cub::DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes,
|
||||
temp_ptr, temp_ptr,
|
||||
temp_ptr, temp_ptr, num_items);
|
||||
return temp_storage_bytes;
|
||||
}
|
||||
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
|
||||
} // wallx_cuda_permute
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,284 @@
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cub/cub.cuh>
|
||||
#include <cassert>
|
||||
#include "../common/cuda_utils.h"
|
||||
|
||||
namespace wallx_cuda_rot_pos {
|
||||
|
||||
__global__ void fused_rot_pos_emb_kernel_int32(
|
||||
const float *__restrict__ inv_freq, // [dim/2] - precomputed inverse frequencies
|
||||
const int32_t *__restrict__ grid_thw, // [num_grids, 3] - (t, h, w) for each grid
|
||||
float *__restrict__ output, // [total_tokens, dim] - output rotary embeddings
|
||||
const int32_t *__restrict__ cumsum_tokens, // [num_grids+1] - cumulative sum of tokens per grid
|
||||
const int dim_half, // dim/2 (size of inv_freq)
|
||||
const int spatial_merge_size, // spatial merge size
|
||||
const int num_grids // number of grids
|
||||
)
|
||||
{
|
||||
const int32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int32_t total_tokens = cumsum_tokens[num_grids];
|
||||
|
||||
extern __shared__ float freq_smem[];
|
||||
for (int i = threadIdx.x; i < dim_half; i += blockDim.x) {
|
||||
freq_smem[i] = inv_freq[i];
|
||||
}
|
||||
if (tid >= total_tokens * dim_half)
|
||||
return;
|
||||
|
||||
const int32_t token_idx = tid / dim_half;
|
||||
const int freq_idx = tid % dim_half;
|
||||
|
||||
// Find which grid this token belongs to
|
||||
int grid_idx = 0;
|
||||
int32_t local_token_idx = token_idx;
|
||||
for (int g = 0; g < num_grids; g++)
|
||||
{
|
||||
if (token_idx < cumsum_tokens[g + 1])
|
||||
{
|
||||
grid_idx = g;
|
||||
local_token_idx = token_idx - cumsum_tokens[g];
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Get grid dimensions
|
||||
const int32_t h = grid_thw[grid_idx * 3 + 1];
|
||||
const int32_t w = grid_thw[grid_idx * 3 + 2];
|
||||
|
||||
// Calculate spatial dimensions after merging
|
||||
const int32_t h_merged = h / spatial_merge_size;
|
||||
const int32_t w_merged = w / spatial_merge_size;
|
||||
const int32_t spatial_tokens = h_merged * w_merged * spatial_merge_size * spatial_merge_size;
|
||||
|
||||
// Get spatial index
|
||||
const int32_t spatial_idx = local_token_idx % spatial_tokens;
|
||||
|
||||
// Decompose spatial index to get merged block and position within block
|
||||
const int32_t tokens_per_block = spatial_merge_size * spatial_merge_size;
|
||||
const int32_t block_idx = spatial_idx / tokens_per_block;
|
||||
const int32_t within_block_idx = spatial_idx % tokens_per_block;
|
||||
|
||||
// Get block coordinates in merged grid
|
||||
const int32_t block_h = block_idx / w_merged;
|
||||
const int32_t block_w = block_idx % w_merged;
|
||||
|
||||
// Get position within block
|
||||
const int32_t within_h = within_block_idx / spatial_merge_size;
|
||||
const int32_t within_w = within_block_idx % spatial_merge_size;
|
||||
|
||||
// Calculate actual h and w positions
|
||||
const int32_t h_pos = block_h * spatial_merge_size + within_h;
|
||||
const int32_t w_pos = block_w * spatial_merge_size + within_w;
|
||||
|
||||
// Compute rotary embedding
|
||||
float freq_val = freq_smem[freq_idx];
|
||||
|
||||
// Output has shape [total_tokens, dim] where dim = 2 * dim_half
|
||||
int32_t out_idx = token_idx * dim_half * 2 + freq_idx;
|
||||
output[out_idx] = h_pos * freq_val; // h_pos frequencies
|
||||
output[out_idx + dim_half] = w_pos * freq_val; // w_pos frequencies
|
||||
}
|
||||
|
||||
__global__ void fused_rot_pos_emb_kernel_int64(
|
||||
const float *__restrict__ inv_freq, // [dim/2] - precomputed inverse frequencies
|
||||
const int64_t *__restrict__ grid_thw, // [num_grids, 3] - (t, h, w) for each grid
|
||||
float *__restrict__ output, // [total_tokens, dim] - output rotary embeddings
|
||||
const int64_t *__restrict__ cumsum_tokens, // [num_grids+1] - cumulative sum of tokens per grid
|
||||
const int dim_half, // dim/2 (size of inv_freq)
|
||||
const int spatial_merge_size, // spatial merge size
|
||||
const int num_grids // number of grids
|
||||
)
|
||||
{
|
||||
const int64_t tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int64_t total_tokens = cumsum_tokens[num_grids];
|
||||
|
||||
extern __shared__ float freq_smem[];
|
||||
for (int i = threadIdx.x; i < dim_half; i += blockDim.x) {
|
||||
freq_smem[i] = inv_freq[i];
|
||||
}
|
||||
|
||||
if (tid >= total_tokens * dim_half)
|
||||
return;
|
||||
|
||||
const int64_t token_idx = tid / dim_half;
|
||||
const int freq_idx = tid % dim_half;
|
||||
|
||||
// Find which grid this token belongs to
|
||||
int grid_idx = 0;
|
||||
int64_t local_token_idx = token_idx;
|
||||
for (int g = 0; g < num_grids; g++)
|
||||
{
|
||||
if (token_idx < cumsum_tokens[g + 1])
|
||||
{
|
||||
grid_idx = g;
|
||||
local_token_idx = token_idx - cumsum_tokens[g];
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Get grid dimensions
|
||||
const int64_t h = grid_thw[grid_idx * 3 + 1];
|
||||
const int64_t w = grid_thw[grid_idx * 3 + 2];
|
||||
|
||||
// Calculate spatial dimensions after merging
|
||||
const int64_t h_merged = h / spatial_merge_size;
|
||||
const int64_t w_merged = w / spatial_merge_size;
|
||||
const int64_t spatial_tokens = h_merged * w_merged * spatial_merge_size * spatial_merge_size;
|
||||
|
||||
// Get spatial index
|
||||
const int64_t spatial_idx = local_token_idx % spatial_tokens;
|
||||
|
||||
// Decompose spatial index to get merged block and position within block
|
||||
const int64_t tokens_per_block = spatial_merge_size * spatial_merge_size;
|
||||
const int64_t block_idx = spatial_idx / tokens_per_block;
|
||||
const int64_t within_block_idx = spatial_idx % tokens_per_block;
|
||||
|
||||
// Get block coordinates in merged grid
|
||||
const int64_t block_h = block_idx / w_merged;
|
||||
const int64_t block_w = block_idx % w_merged;
|
||||
|
||||
// Get position within block
|
||||
const int64_t within_h = within_block_idx / spatial_merge_size;
|
||||
const int64_t within_w = within_block_idx % spatial_merge_size;
|
||||
|
||||
// Calculate actual h and w positions
|
||||
const int64_t h_pos = block_h * spatial_merge_size + within_h;
|
||||
const int64_t w_pos = block_w * spatial_merge_size + within_w;
|
||||
|
||||
// Compute rotary embedding
|
||||
float freq_val = freq_smem[freq_idx];
|
||||
|
||||
// Output has shape [total_tokens, dim] where dim = 2 * dim_half
|
||||
int64_t out_idx = token_idx * dim_half * 2 + freq_idx;
|
||||
output[out_idx] = h_pos * freq_val; // h_pos frequencies
|
||||
output[out_idx + dim_half] = w_pos * freq_val; // w_pos frequencies
|
||||
}
|
||||
|
||||
__global__ void compute_token_counts_kernel_int32(
|
||||
const int32_t *__restrict__ grid_thw,
|
||||
int32_t *__restrict__ token_counts,
|
||||
const int spatial_merge_size,
|
||||
const int num_grids)
|
||||
{
|
||||
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= num_grids)
|
||||
return;
|
||||
|
||||
int32_t t = grid_thw[idx * 3 + 0];
|
||||
int32_t h = grid_thw[idx * 3 + 1];
|
||||
int32_t w = grid_thw[idx * 3 + 2];
|
||||
int32_t h_merged = h / spatial_merge_size;
|
||||
int32_t w_merged = w / spatial_merge_size;
|
||||
token_counts[idx] = t * h_merged * w_merged * spatial_merge_size * spatial_merge_size;
|
||||
}
|
||||
|
||||
__global__ void compute_token_counts_kernel_int64(
|
||||
const int64_t *__restrict__ grid_thw,
|
||||
int64_t *__restrict__ token_counts,
|
||||
const int spatial_merge_size,
|
||||
const int num_grids)
|
||||
{
|
||||
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= num_grids)
|
||||
return;
|
||||
|
||||
int64_t t = grid_thw[idx * 3 + 0];
|
||||
int64_t h = grid_thw[idx * 3 + 1];
|
||||
int64_t w = grid_thw[idx * 3 + 2];
|
||||
int64_t h_merged = h / spatial_merge_size;
|
||||
int64_t w_merged = w / spatial_merge_size;
|
||||
token_counts[idx] = t * h_merged * w_merged * spatial_merge_size * spatial_merge_size;
|
||||
}
|
||||
|
||||
void GetTokenCounts(
|
||||
const at::Tensor& grid_thw, // [num_grids, 3]
|
||||
const at::Tensor& token_counts, // [num_grids]
|
||||
int spatial_merge_size
|
||||
) {
|
||||
ASSERT_CHECK(grid_thw.dim() == 2);
|
||||
ASSERT_CHECK(grid_thw.size(1) == 3);
|
||||
ASSERT_CHECK(grid_thw.size(0) == token_counts.size(0));
|
||||
|
||||
ASSERT_CHECK(spatial_merge_size > 0);
|
||||
|
||||
const int num_grids = grid_thw.size(0);
|
||||
|
||||
const int threads = 256;
|
||||
const int blocks = (num_grids + threads - 1) / threads;
|
||||
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
if (grid_thw.scalar_type() == at::kInt) {
|
||||
compute_token_counts_kernel_int32<<<blocks, threads, 0, stream>>>(
|
||||
static_cast<int32_t*>(grid_thw.data_ptr()),
|
||||
static_cast<int32_t*>(token_counts.data_ptr()),
|
||||
spatial_merge_size,
|
||||
num_grids);
|
||||
} else if (grid_thw.scalar_type() == at::kLong) {
|
||||
compute_token_counts_kernel_int64<<<blocks, threads, 0, stream>>>(
|
||||
static_cast<int64_t*>(grid_thw.data_ptr()),
|
||||
static_cast<int64_t*>(token_counts.data_ptr()),
|
||||
spatial_merge_size,
|
||||
num_grids);
|
||||
} else {
|
||||
throw std::runtime_error("Unsupported data type for GetTokenCounts");
|
||||
}
|
||||
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
void RotPosEmb(
|
||||
const at::Tensor& inv_freq, // [dim/2]
|
||||
const at::Tensor& grid_thw, // [num_grids, 3]
|
||||
const at::Tensor& output,
|
||||
const at::Tensor& cumsum_tokens,
|
||||
int spatial_merge_size
|
||||
) {
|
||||
ASSERT_CHECK(output.size(0) > 0);
|
||||
ASSERT_CHECK(inv_freq.dim() == 1);
|
||||
ASSERT_CHECK(inv_freq.scalar_type() == at::kFloat);
|
||||
|
||||
const int dim_half = inv_freq.size(0);
|
||||
const int num_grids = grid_thw.size(0);
|
||||
const int total_tokens = output.size(0);
|
||||
|
||||
const int64_t num_elements = total_tokens * dim_half;
|
||||
|
||||
const int threads_per_block = 256;
|
||||
const int num_blocks = static_cast<int>((num_elements + threads_per_block - 1) / threads_per_block);
|
||||
cudaStream_t stream = static_cast<cudaStream_t>(
|
||||
at::cuda::getCurrentCUDAStream().stream());
|
||||
|
||||
if (grid_thw.scalar_type() == at::kInt) {
|
||||
fused_rot_pos_emb_kernel_int32<<<num_blocks, threads_per_block, dim_half * sizeof(float), stream>>>(
|
||||
static_cast<float*>(inv_freq.data_ptr()),
|
||||
static_cast<int32_t*>(grid_thw.data_ptr()),
|
||||
static_cast<float*>(output.data_ptr()),
|
||||
static_cast<int32_t*>(cumsum_tokens.data_ptr()),
|
||||
dim_half,
|
||||
spatial_merge_size,
|
||||
num_grids);
|
||||
} else if (grid_thw.scalar_type() == at::kLong) {
|
||||
fused_rot_pos_emb_kernel_int64<<<num_blocks, threads_per_block, dim_half * sizeof(float), stream>>>(
|
||||
static_cast<float*>(inv_freq.data_ptr()),
|
||||
static_cast<int64_t*>(grid_thw.data_ptr()),
|
||||
static_cast<float*>(output.data_ptr()),
|
||||
static_cast<int64_t*>(cumsum_tokens.data_ptr()),
|
||||
dim_half,
|
||||
spatial_merge_size,
|
||||
num_grids);
|
||||
} else {
|
||||
throw std::runtime_error("Unsupported data type for RotPosEmb");
|
||||
}
|
||||
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
|
||||
} // wallx_cuda_rot_pos
|
||||
@@ -0,0 +1,284 @@
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cassert>
|
||||
#include "../common/cuda_utils.h"
|
||||
|
||||
namespace wallx_cuda_window_index {
|
||||
|
||||
__global__ void compute_metadata(
|
||||
const int *grid_thw, // [num_grids, 3]
|
||||
int *grid_info, // [num_grids, 6]: [grid_elements, grid_windows, llm_h, llm_w, num_windows_h, num_windows_w]
|
||||
int *global_totals, // [total_elements, total_windows]
|
||||
int num_grids,
|
||||
int spatial_merge_size,
|
||||
int vit_merger_window_size)
|
||||
{
|
||||
int grid_idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (grid_idx >= num_grids)
|
||||
return;
|
||||
|
||||
int grid_t = grid_thw[grid_idx * 3 + 0];
|
||||
int grid_h = grid_thw[grid_idx * 3 + 1];
|
||||
int grid_w = grid_thw[grid_idx * 3 + 2];
|
||||
|
||||
int llm_h = grid_h / spatial_merge_size;
|
||||
int llm_w = grid_w / spatial_merge_size;
|
||||
|
||||
int pad_h = (vit_merger_window_size - llm_h % vit_merger_window_size) % vit_merger_window_size;
|
||||
int pad_w = (vit_merger_window_size - llm_w % vit_merger_window_size) % vit_merger_window_size;
|
||||
|
||||
int num_windows_h = (llm_h + pad_h) / vit_merger_window_size;
|
||||
int num_windows_w = (llm_w + pad_w) / vit_merger_window_size;
|
||||
|
||||
int grid_elements = grid_t * llm_h * llm_w;
|
||||
int grid_windows = grid_t * num_windows_h * num_windows_w;
|
||||
|
||||
grid_info[grid_idx * 6 + 0] = grid_elements;
|
||||
grid_info[grid_idx * 6 + 1] = grid_windows;
|
||||
grid_info[grid_idx * 6 + 2] = llm_h;
|
||||
grid_info[grid_idx * 6 + 3] = llm_w;
|
||||
grid_info[grid_idx * 6 + 4] = num_windows_h;
|
||||
grid_info[grid_idx * 6 + 5] = num_windows_w;
|
||||
|
||||
atomicAdd(&global_totals[0], grid_elements);
|
||||
atomicAdd(&global_totals[1], grid_windows);
|
||||
}
|
||||
|
||||
__global__ void compute_window_counts(
|
||||
const int *grid_thw,
|
||||
const int *grid_info,
|
||||
int *window_counts,
|
||||
int vit_merger_window_size,
|
||||
int spatial_merge_unit,
|
||||
int num_grids)
|
||||
{
|
||||
int grid_idx = blockIdx.y;
|
||||
int t_idx = blockIdx.x;
|
||||
|
||||
if (grid_idx >= num_grids)
|
||||
return;
|
||||
|
||||
int grid_t = grid_thw[grid_idx * 3 + 0];
|
||||
if (t_idx >= grid_t)
|
||||
return;
|
||||
|
||||
int llm_h = grid_info[grid_idx * 6 + 2];
|
||||
int llm_w = grid_info[grid_idx * 6 + 3];
|
||||
int num_windows_h = grid_info[grid_idx * 6 + 4];
|
||||
int num_windows_w = grid_info[grid_idx * 6 + 5];
|
||||
|
||||
int window_base = 0;
|
||||
for (int g = 0; g < grid_idx; g++)
|
||||
{
|
||||
window_base += grid_info[g * 6 + 1];
|
||||
}
|
||||
|
||||
int t_window_base = window_base + t_idx * num_windows_h * num_windows_w;
|
||||
|
||||
int thread_id = threadIdx.x;
|
||||
int warp_id = thread_id / 32;
|
||||
int lane_id = thread_id % 32;
|
||||
|
||||
int windows_per_t = num_windows_h * num_windows_w;
|
||||
int warps_per_block = blockDim.x / 32;
|
||||
|
||||
if (lane_id == 0)
|
||||
{
|
||||
for (int window_idx = warp_id; window_idx < windows_per_t; window_idx += warps_per_block)
|
||||
{
|
||||
int win_h = window_idx / num_windows_w;
|
||||
int win_w = window_idx % num_windows_w;
|
||||
|
||||
int start_h = win_h * vit_merger_window_size;
|
||||
int start_w = win_w * vit_merger_window_size;
|
||||
|
||||
int valid_h = min(vit_merger_window_size, llm_h - start_h);
|
||||
int valid_w = min(vit_merger_window_size, llm_w - start_w);
|
||||
|
||||
int valid_count = (valid_h > 0 && valid_w > 0) ? valid_h * valid_w : 0;
|
||||
|
||||
window_counts[t_window_base + window_idx] = valid_count;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void compute_cu_window_seqlens(
|
||||
const int *window_counts,
|
||||
int *cu_window_seqlens, // [total_windows + 1]
|
||||
int total_windows,
|
||||
int spatial_merge_unit
|
||||
)
|
||||
{
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
if (tid == 0)
|
||||
{
|
||||
cu_window_seqlens[0] = 0;
|
||||
}
|
||||
|
||||
if (tid < total_windows)
|
||||
{
|
||||
cu_window_seqlens[tid + 1] = window_counts[tid] * spatial_merge_unit;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (tid == 0)
|
||||
{
|
||||
for (int i = 1; i <= total_windows; i++)
|
||||
{
|
||||
cu_window_seqlens[i] += cu_window_seqlens[i - 1];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void generate_window_indices(
|
||||
const int *grid_thw,
|
||||
const int *grid_info,
|
||||
const int *cu_window_seqlens,
|
||||
int *window_indices,
|
||||
int vit_merger_window_size,
|
||||
int spatial_merge_unit,
|
||||
int num_grids)
|
||||
{
|
||||
int grid_idx = blockIdx.y;
|
||||
int t_idx = blockIdx.x;
|
||||
|
||||
if (grid_idx >= num_grids)
|
||||
return;
|
||||
|
||||
int grid_t = grid_thw[grid_idx * 3 + 0];
|
||||
if (t_idx >= grid_t)
|
||||
return;
|
||||
|
||||
int llm_h = grid_info[grid_idx * 6 + 2];
|
||||
int llm_w = grid_info[grid_idx * 6 + 3];
|
||||
int num_windows_h = grid_info[grid_idx * 6 + 4];
|
||||
int num_windows_w = grid_info[grid_idx * 6 + 5];
|
||||
|
||||
int element_base = 0;
|
||||
for (int g = 0; g < grid_idx; g++)
|
||||
{
|
||||
element_base += grid_info[g * 6 + 0];
|
||||
}
|
||||
int t_element_base = element_base + t_idx * llm_h * llm_w;
|
||||
|
||||
int window_base = 0;
|
||||
for (int g = 0; g < grid_idx; g++)
|
||||
{
|
||||
window_base += grid_info[g * 6 + 1];
|
||||
}
|
||||
int t_window_base = window_base + t_idx * num_windows_h * num_windows_w;
|
||||
|
||||
int thread_id = threadIdx.x;
|
||||
int warp_id = thread_id / 32;
|
||||
int lane_id = thread_id % 32;
|
||||
|
||||
int windows_per_t = num_windows_h * num_windows_w;
|
||||
int warps_per_block = blockDim.x / 32;
|
||||
|
||||
for (int window_idx = warp_id; window_idx < windows_per_t; window_idx += warps_per_block)
|
||||
{
|
||||
int win_h = window_idx / num_windows_w;
|
||||
int win_w = window_idx % num_windows_w;
|
||||
|
||||
int global_window_idx = t_window_base + window_idx;
|
||||
int output_offset = cu_window_seqlens[global_window_idx] / spatial_merge_unit;
|
||||
|
||||
int start_h = win_h * vit_merger_window_size;
|
||||
int start_w = win_w * vit_merger_window_size;
|
||||
|
||||
int valid_h = min(vit_merger_window_size, llm_h - start_h);
|
||||
int valid_w = min(vit_merger_window_size, llm_w - start_w);
|
||||
|
||||
for (int elem_idx = lane_id; elem_idx < valid_h * valid_w; elem_idx += 32)
|
||||
{
|
||||
int local_h = elem_idx / valid_w;
|
||||
int local_w = elem_idx % valid_w;
|
||||
|
||||
int abs_h = start_h + local_h;
|
||||
int abs_w = start_w + local_w;
|
||||
|
||||
int value = t_element_base + abs_h * llm_w + abs_w;
|
||||
|
||||
int base_offset = output_offset + elem_idx;
|
||||
window_indices[base_offset] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void GetTotals(
|
||||
const at::Tensor& grid_thw,
|
||||
const at::Tensor& grid_info_tensor,
|
||||
const at::Tensor& global_totals_tensor,
|
||||
int spatial_merge_size,
|
||||
int vit_merger_window_size
|
||||
) {
|
||||
ASSERT_CHECK(grid_thw.dim() == 2 && grid_thw.size(1) == 3);
|
||||
ASSERT_CHECK(grid_thw.scalar_type() == at::kInt);
|
||||
|
||||
int num_grids = grid_thw.size(0);
|
||||
const int *d_grid_thw = static_cast<int32_t*>(grid_thw.data_ptr());
|
||||
|
||||
int *d_grid_info = static_cast<int32_t*>(grid_info_tensor.data_ptr());
|
||||
int *d_global_totals = static_cast<int32_t*>(global_totals_tensor.data_ptr());
|
||||
|
||||
int threads1 = 256;
|
||||
int blocks1 = (num_grids + threads1 - 1) / threads1;
|
||||
|
||||
compute_metadata<<<blocks1, threads1>>>(
|
||||
d_grid_thw, d_grid_info, d_global_totals,
|
||||
num_grids, spatial_merge_size, vit_merger_window_size);
|
||||
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
void GetWindowIndex(
|
||||
const at::Tensor& grid_thw,
|
||||
const at::Tensor& grid_info_tensor,
|
||||
const at::Tensor& window_indices,
|
||||
const at::Tensor& cu_window_seqlens,
|
||||
const at::Tensor& window_counts_tensor,
|
||||
int max_grid_t,
|
||||
int spatial_merge_size,
|
||||
int vit_merger_window_size,
|
||||
int patch_size,
|
||||
int spatial_merge_unit) {
|
||||
|
||||
int num_grids = grid_thw.size(0);
|
||||
|
||||
dim3 blocks2(max_grid_t, num_grids);
|
||||
dim3 threads2(256);
|
||||
|
||||
const int *d_grid_thw = static_cast<int32_t*>(grid_thw.data_ptr());
|
||||
|
||||
int *d_grid_info = static_cast<int32_t*>(grid_info_tensor.data_ptr());
|
||||
|
||||
int *d_window_indices = static_cast<int32_t*>(window_indices.data_ptr());
|
||||
int *d_cu_window_seqlens = static_cast<int32_t*>(cu_window_seqlens.data_ptr());
|
||||
|
||||
int *d_window_counts = static_cast<int32_t*>(window_counts_tensor.data_ptr());
|
||||
|
||||
compute_window_counts<<<blocks2, threads2>>>(
|
||||
d_grid_thw, d_grid_info, d_window_counts,
|
||||
vit_merger_window_size, spatial_merge_unit, num_grids);
|
||||
|
||||
int total_windows = window_counts_tensor.size(0);
|
||||
|
||||
int threads4 = 256;
|
||||
int blocks4 = (total_windows + threads4 - 1) / threads4;
|
||||
compute_cu_window_seqlens<<<blocks4, threads4>>>(
|
||||
d_window_counts, d_cu_window_seqlens, total_windows, spatial_merge_unit);
|
||||
|
||||
generate_window_indices<<<blocks2, threads2>>>(
|
||||
d_grid_thw, d_grid_info, d_cu_window_seqlens, d_window_indices,
|
||||
vit_merger_window_size, spatial_merge_unit, num_grids);
|
||||
|
||||
sync_check_cuda_error();
|
||||
}
|
||||
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
// Removed TVM FFI export (see pybind11 registration)
|
||||
|
||||
} // wallx_cuda_window_index
|
||||
@@ -0,0 +1,340 @@
|
||||
"""Index computation operators for attention."""
|
||||
|
||||
import logging
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from wall_x.model.core.ops.base import OpsProxy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Sentinel value for padding in window index computation.
|
||||
# Valid indices are always >= 0 (they are positional offsets within flattened
|
||||
# grid tensors), so any negative value is safe. -100 is chosen by convention
|
||||
# (same as HuggingFace's ignore_index for cross-entropy loss).
|
||||
_PAD_VALUE = -100
|
||||
|
||||
|
||||
def _compute_vision_position_ids(
|
||||
input_tokens,
|
||||
token_set,
|
||||
st,
|
||||
image_grid_thw,
|
||||
video_grid_thw,
|
||||
second_per_grid_ts,
|
||||
image_index,
|
||||
video_index,
|
||||
remain_images,
|
||||
remain_videos,
|
||||
image_token_id,
|
||||
video_token_id,
|
||||
spatial_merge_size,
|
||||
tokens_per_second,
|
||||
device,
|
||||
llm_pos_ids_list,
|
||||
):
|
||||
"""Compute position IDs for a single vision token (image or video).
|
||||
|
||||
Returns updated (st, image_index, video_index, remain_images, remain_videos).
|
||||
"""
|
||||
if image_token_id in token_set and remain_images > 0:
|
||||
ed_image = input_tokens.index(image_token_id, st)
|
||||
else:
|
||||
ed_image = len(input_tokens) + 1
|
||||
if video_token_id in token_set and remain_videos > 0:
|
||||
ed_video = input_tokens.index(video_token_id, st)
|
||||
else:
|
||||
ed_video = len(input_tokens) + 1
|
||||
|
||||
if ed_image < ed_video:
|
||||
if image_index >= len(image_grid_thw):
|
||||
raise IndexError(
|
||||
f"image_index {image_index} out of range (have {len(image_grid_thw)} image grids)"
|
||||
)
|
||||
t, h, w = (
|
||||
image_grid_thw[image_index][0],
|
||||
image_grid_thw[image_index][1],
|
||||
image_grid_thw[image_index][2],
|
||||
)
|
||||
# Images are single-frame: no temporal progression, so second_per_grid_t = 0.
|
||||
# This makes all image tokens share t_index = 0 (only spatial positions vary).
|
||||
second_per_grid_t = 0
|
||||
image_index += 1
|
||||
remain_images -= 1
|
||||
ed = ed_image
|
||||
else:
|
||||
if video_index >= len(video_grid_thw):
|
||||
raise IndexError(
|
||||
f"video_index {video_index} out of range (have {len(video_grid_thw)} video grids)"
|
||||
)
|
||||
t, h, w = (
|
||||
video_grid_thw[video_index][0],
|
||||
video_grid_thw[video_index][1],
|
||||
video_grid_thw[video_index][2],
|
||||
)
|
||||
if second_per_grid_ts is not None:
|
||||
second_per_grid_t = second_per_grid_ts[video_index]
|
||||
else:
|
||||
second_per_grid_t = 1.0
|
||||
video_index += 1
|
||||
remain_videos -= 1
|
||||
ed = ed_video
|
||||
|
||||
llm_grid_t, llm_grid_h, llm_grid_w = (
|
||||
int(t),
|
||||
int(h) // spatial_merge_size,
|
||||
int(w) // spatial_merge_size,
|
||||
)
|
||||
text_len = ed - st
|
||||
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
|
||||
llm_pos_ids_list.append(
|
||||
torch.arange(text_len, device=device).view(1, -1).expand(3, -1) + st_idx
|
||||
)
|
||||
range_tensor = torch.arange(llm_grid_t, device=device).view(-1, 1)
|
||||
expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)
|
||||
time_tensor = expanded_range * second_per_grid_t * tokens_per_second
|
||||
t_index = time_tensor.long().flatten()
|
||||
h_index = (
|
||||
torch.arange(llm_grid_h, device=device)
|
||||
.view(1, -1, 1)
|
||||
.expand(llm_grid_t, -1, llm_grid_w)
|
||||
.flatten()
|
||||
)
|
||||
w_index = (
|
||||
torch.arange(llm_grid_w, device=device)
|
||||
.view(1, 1, -1)
|
||||
.expand(llm_grid_t, llm_grid_h, -1)
|
||||
.flatten()
|
||||
)
|
||||
llm_pos_ids_list.append(
|
||||
torch.stack([t_index, h_index, w_index]) + text_len + st_idx
|
||||
)
|
||||
st = ed + llm_grid_t * llm_grid_h * llm_grid_w
|
||||
return st, image_index, video_index, remain_images, remain_videos
|
||||
|
||||
|
||||
def _compute_text_only_positions(input_ids, attention_mask):
|
||||
"""Compute position IDs for text-only inputs (no vision tokens)."""
|
||||
if attention_mask is not None:
|
||||
position_ids = attention_mask.long().cumsum(-1) - 1
|
||||
position_ids.masked_fill_(attention_mask == 0, 1)
|
||||
position_ids = (
|
||||
position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
|
||||
)
|
||||
max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[
|
||||
0
|
||||
]
|
||||
mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
|
||||
else:
|
||||
position_ids = (
|
||||
torch.arange(input_ids.shape[1], device=input_ids.device)
|
||||
.view(1, 1, -1)
|
||||
.expand(3, input_ids.shape[0], -1)
|
||||
)
|
||||
mrope_position_deltas = torch.zeros(
|
||||
[input_ids.shape[0], 1],
|
||||
device=input_ids.device,
|
||||
dtype=input_ids.dtype,
|
||||
)
|
||||
return position_ids, mrope_position_deltas
|
||||
|
||||
|
||||
class GetRopeIndexOp(OpsProxy):
|
||||
"""Compute 3D RoPE position indices for multimodal inputs.
|
||||
|
||||
Signature: get_rope_index(input_ids, image_grid_thw, video_grid_thw,
|
||||
second_per_grid_ts, attention_mask, spatial_merge_size,
|
||||
image_token_id, video_token_id, vision_start_token_id,
|
||||
tokens_per_second) -> (position_ids, mrope_position_deltas)
|
||||
"""
|
||||
|
||||
@property
|
||||
def _external_accel_name(self):
|
||||
return "get_rope_index"
|
||||
|
||||
def _get_cuda_kernel(self):
|
||||
try:
|
||||
from wall_x.model.core.ops._cuda_wrappers import GetRopeIndex
|
||||
|
||||
return GetRopeIndex()
|
||||
except ImportError:
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.warning("GetRopeIndexOp: CUDA kernel load failed: %s", e)
|
||||
return None
|
||||
|
||||
def _pytorch_fallback(
|
||||
self,
|
||||
input_ids: torch.LongTensor,
|
||||
image_grid_thw: Optional[torch.LongTensor],
|
||||
video_grid_thw: Optional[torch.LongTensor],
|
||||
second_per_grid_ts: Optional[torch.Tensor],
|
||||
attention_mask: Optional[torch.Tensor],
|
||||
spatial_merge_size: int,
|
||||
image_token_id: int,
|
||||
video_token_id: int,
|
||||
vision_start_token_id: int,
|
||||
tokens_per_second: int,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
if input_ids is None or (image_grid_thw is None and video_grid_thw is None):
|
||||
return _compute_text_only_positions(input_ids, attention_mask)
|
||||
|
||||
mrope_position_deltas = []
|
||||
device = input_ids.device
|
||||
total_input_ids = input_ids
|
||||
if attention_mask is None:
|
||||
attention_mask = torch.ones_like(total_input_ids)
|
||||
position_ids = torch.ones(
|
||||
3,
|
||||
input_ids.shape[0],
|
||||
input_ids.shape[1],
|
||||
dtype=input_ids.dtype,
|
||||
device=input_ids.device,
|
||||
)
|
||||
image_index, video_index = 0, 0
|
||||
attention_mask = attention_mask.to(total_input_ids.device)
|
||||
for i, input_ids_i in enumerate(total_input_ids):
|
||||
input_ids_i = input_ids_i[attention_mask[i] == 1]
|
||||
vision_start_indices = torch.argwhere(
|
||||
input_ids_i == vision_start_token_id
|
||||
).squeeze(1)
|
||||
# Boundary check: ensure vision_start + 1 doesn't exceed sequence length
|
||||
valid_mask = (vision_start_indices + 1) < len(input_ids_i)
|
||||
vision_start_indices = vision_start_indices[valid_mask]
|
||||
vision_tokens = input_ids_i[vision_start_indices + 1]
|
||||
image_nums = int((vision_tokens == image_token_id).sum())
|
||||
video_nums = int((vision_tokens == video_token_id).sum())
|
||||
input_tokens = input_ids_i.tolist()
|
||||
token_set = set(input_tokens)
|
||||
llm_pos_ids_list: list = []
|
||||
st = 0
|
||||
remain_images, remain_videos = image_nums, video_nums
|
||||
for _ in range(image_nums + video_nums):
|
||||
st, image_index, video_index, remain_images, remain_videos = (
|
||||
_compute_vision_position_ids(
|
||||
input_tokens,
|
||||
token_set,
|
||||
st,
|
||||
image_grid_thw,
|
||||
video_grid_thw,
|
||||
second_per_grid_ts,
|
||||
image_index,
|
||||
video_index,
|
||||
remain_images,
|
||||
remain_videos,
|
||||
image_token_id,
|
||||
video_token_id,
|
||||
spatial_merge_size,
|
||||
tokens_per_second,
|
||||
device,
|
||||
llm_pos_ids_list,
|
||||
)
|
||||
)
|
||||
if st < len(input_tokens):
|
||||
st_idx = (
|
||||
llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
|
||||
)
|
||||
text_len = len(input_tokens) - st
|
||||
llm_pos_ids_list.append(
|
||||
torch.arange(text_len, device=device).view(1, -1).expand(3, -1)
|
||||
+ st_idx
|
||||
)
|
||||
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
|
||||
position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(
|
||||
position_ids.device
|
||||
)
|
||||
mrope_position_deltas.append(
|
||||
llm_positions.max() + 1 - len(total_input_ids[i])
|
||||
)
|
||||
mrope_position_deltas = torch.tensor(
|
||||
mrope_position_deltas, device=total_input_ids.device
|
||||
).unsqueeze(1)
|
||||
return position_ids, mrope_position_deltas
|
||||
|
||||
|
||||
class GetWindowIndexOp(OpsProxy):
|
||||
"""Compute window attention indices for ViT.
|
||||
|
||||
Signature: get_window_index(grid_thw, window_size, spatial_merge_size,
|
||||
patch_size, spatial_merge_unit=1) -> (window_index, cu_window_seqlens)
|
||||
"""
|
||||
|
||||
@property
|
||||
def _external_accel_name(self):
|
||||
return "get_window_index"
|
||||
|
||||
def _get_cuda_kernel(self):
|
||||
try:
|
||||
from wall_x.model.core.ops._cuda_wrappers import get_window_index_cuda
|
||||
|
||||
return get_window_index_cuda
|
||||
except ImportError:
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.warning("GetWindowIndexOp: CUDA kernel load failed: %s", e)
|
||||
return None
|
||||
|
||||
def _pytorch_fallback(
|
||||
self,
|
||||
grid_thw,
|
||||
window_size,
|
||||
spatial_merge_size,
|
||||
patch_size,
|
||||
spatial_merge_unit=1,
|
||||
):
|
||||
device = grid_thw.device
|
||||
vit_merger_window_size = window_size // spatial_merge_size // patch_size
|
||||
window_index_list = []
|
||||
cu_window_seqlens = [0]
|
||||
window_index_id = 0
|
||||
for grid_t, grid_h, grid_w in grid_thw:
|
||||
llm_grid_h = grid_h // spatial_merge_size
|
||||
llm_grid_w = grid_w // spatial_merge_size
|
||||
index = torch.arange(
|
||||
grid_t * llm_grid_h * llm_grid_w, device=device
|
||||
).reshape(grid_t, llm_grid_h, llm_grid_w)
|
||||
pad_h = (
|
||||
vit_merger_window_size - llm_grid_h % vit_merger_window_size
|
||||
) % vit_merger_window_size
|
||||
pad_w = (
|
||||
vit_merger_window_size - llm_grid_w % vit_merger_window_size
|
||||
) % vit_merger_window_size
|
||||
num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size
|
||||
num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size
|
||||
index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", _PAD_VALUE)
|
||||
index_padded = (
|
||||
index_padded.reshape(
|
||||
grid_t,
|
||||
num_windows_h,
|
||||
vit_merger_window_size,
|
||||
num_windows_w,
|
||||
vit_merger_window_size,
|
||||
)
|
||||
.permute(0, 1, 3, 2, 4)
|
||||
.reshape(
|
||||
grid_t,
|
||||
num_windows_h * num_windows_w,
|
||||
vit_merger_window_size,
|
||||
vit_merger_window_size,
|
||||
)
|
||||
)
|
||||
seqlens = (index_padded != _PAD_VALUE).sum([2, 3]).reshape(-1)
|
||||
index_padded = index_padded.reshape(-1)
|
||||
index_new = index_padded[index_padded != _PAD_VALUE]
|
||||
window_index_list.append(index_new + window_index_id)
|
||||
cu_seqlens_tmp = (
|
||||
seqlens.cumsum(0) * spatial_merge_unit + cu_window_seqlens[-1]
|
||||
)
|
||||
cu_window_seqlens.extend(cu_seqlens_tmp.tolist())
|
||||
window_index_id += int(grid_t) * int(llm_grid_h) * int(llm_grid_w)
|
||||
window_index = torch.cat(window_index_list, dim=0)
|
||||
cu_window_seqlens = torch.tensor(
|
||||
cu_window_seqlens, dtype=grid_thw.dtype, device=grid_thw.device
|
||||
)
|
||||
return window_index, cu_window_seqlens
|
||||
|
||||
|
||||
get_rope_index = GetRopeIndexOp()
|
||||
get_window_index = GetWindowIndexOp()
|
||||
@@ -0,0 +1,97 @@
|
||||
"""MoE routing operators (permute/unpermute).
|
||||
|
||||
These operators reorder tokens by expert assignment for efficient
|
||||
Mixture-of-Experts processing.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
|
||||
from wall_x.model.core.ops.base import OpsProxy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PermuteOp(OpsProxy):
|
||||
"""Reorder tokens by expert assignment for MoE processing.
|
||||
|
||||
Signature: permute(tokens, indices, num_out_tokens=None, max_token_num=0) -> (permuted_tokens, sorted_indices)
|
||||
"""
|
||||
|
||||
@property
|
||||
def _external_accel_name(self):
|
||||
return "permute"
|
||||
|
||||
def _get_cuda_kernel(self):
|
||||
try:
|
||||
from wall_x.model.core.ops._cuda_wrappers import permute_kernel
|
||||
|
||||
return permute_kernel
|
||||
except ImportError:
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.warning("PermuteOp: CUDA kernel load failed: %s", e)
|
||||
return None
|
||||
|
||||
def _pytorch_fallback(self, tokens, indices, num_out_tokens=None, max_token_num=0):
|
||||
"""PyTorch fallback for permute.
|
||||
|
||||
Args:
|
||||
num_out_tokens: If set, truncate output to this many tokens (matches
|
||||
external_accel behavior of discarding padding expert assignments).
|
||||
max_token_num: Unused, kept for external_accel API compatibility.
|
||||
"""
|
||||
del max_token_num # unused, external_accel API compat
|
||||
if indices.dim() == 1:
|
||||
indices = indices.view(-1, 1)
|
||||
expand_factor = indices.size(1)
|
||||
flatten_indices = indices.view(-1)
|
||||
# Keep int64 throughout to avoid precision loss on large token counts
|
||||
sorted_indices = torch.argsort(flatten_indices, stable=True)
|
||||
permuted_tokens = tokens.index_select(0, sorted_indices // expand_factor)
|
||||
if num_out_tokens is not None:
|
||||
permuted_tokens = permuted_tokens[:num_out_tokens]
|
||||
sorted_indices = sorted_indices[:num_out_tokens]
|
||||
return permuted_tokens, sorted_indices
|
||||
|
||||
|
||||
class UnpermuteOp(OpsProxy):
|
||||
"""Restore tokens to original order after MoE processing.
|
||||
|
||||
Signature: unpermute(permuted_tokens, sorted_indices, probs=None) -> restored_tokens
|
||||
"""
|
||||
|
||||
@property
|
||||
def _external_accel_name(self):
|
||||
return "unpermute"
|
||||
|
||||
def _get_cuda_kernel(self):
|
||||
try:
|
||||
from wall_x.model.core.ops._cuda_wrappers import unpermute_kernel
|
||||
|
||||
return unpermute_kernel
|
||||
except ImportError:
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.warning("UnpermuteOp: CUDA kernel load failed: %s", e)
|
||||
return None
|
||||
|
||||
def _pytorch_fallback(self, permuted_tokens, sorted_indices, probs=None):
|
||||
if probs is not None:
|
||||
merge_factor = probs.size(1)
|
||||
else:
|
||||
merge_factor = 1
|
||||
unpermuted_tokens = torch.zeros_like(permuted_tokens)
|
||||
unpermuted_tokens.index_copy_(0, sorted_indices.long(), permuted_tokens)
|
||||
unpermuted_tokens = unpermuted_tokens.reshape(
|
||||
-1, merge_factor, permuted_tokens.size(-1)
|
||||
)
|
||||
if probs is not None:
|
||||
unpermuted_tokens = unpermuted_tokens * probs.unsqueeze(-1)
|
||||
unpermuted_tokens = unpermuted_tokens.sum(dim=1)
|
||||
return unpermuted_tokens
|
||||
|
||||
|
||||
permute = PermuteOp()
|
||||
unpermute = UnpermuteOp()
|
||||
@@ -0,0 +1,42 @@
|
||||
"""RMS normalization operator."""
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
|
||||
from wall_x.model.core.ops.base import OpsProxy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RMSNormOp(OpsProxy):
|
||||
"""RMS normalization: rmsnorm(x, weight, eps).
|
||||
|
||||
Accepts ``rmsnorm(x, weight, eps)``.
|
||||
"""
|
||||
|
||||
@property
|
||||
def _external_accel_name(self):
|
||||
return "rmsnorm"
|
||||
|
||||
def _get_cuda_kernel(self):
|
||||
try:
|
||||
from wall_x.model.core.ops._cuda_wrappers import Rmsnorm
|
||||
|
||||
return Rmsnorm()
|
||||
except ImportError:
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.warning("RMSNormOp: CUDA kernel load failed: %s", e)
|
||||
return None
|
||||
|
||||
def _pytorch_fallback(self, hidden_states, weight, eps=1e-6):
|
||||
"""Pure PyTorch RMSNorm. eps default matches external_accel (1e-6, not PyTorch's 1e-5)."""
|
||||
input_dtype = hidden_states.dtype
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
||||
hidden_states = hidden_states * torch.rsqrt(variance + eps)
|
||||
return (weight * hidden_states).to(input_dtype)
|
||||
|
||||
|
||||
rmsnorm = RMSNormOp()
|
||||
@@ -0,0 +1,234 @@
|
||||
"""Rotary position embedding operators."""
|
||||
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
|
||||
from wall_x.model.core.ops.base import OpsProxy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _rotate_half(x):
|
||||
"""Rotate half: split last dim in two halves, negate-swap, concat."""
|
||||
x1 = x[..., : x.shape[-1] // 2]
|
||||
x2 = x[..., x.shape[-1] // 2 :]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def _rotate_interleave(x):
|
||||
"""Interleaved rotation: pairs of (even, odd) elements."""
|
||||
x_even = x[..., ::2]
|
||||
x_odd = x[..., 1::2]
|
||||
return torch.stack((-x_odd, x_even), dim=-1).flatten(-2)
|
||||
|
||||
|
||||
class RoPEOp(OpsProxy):
|
||||
"""Standard rotary position embedding.
|
||||
|
||||
Signature: rope(q, k, cos, sin, interleave=False) -> (q_embed, k_embed)
|
||||
"""
|
||||
|
||||
@property
|
||||
def _external_accel_name(self):
|
||||
return "rope"
|
||||
|
||||
def _get_cuda_kernel(self):
|
||||
try:
|
||||
from wall_x.model.core.ops._cuda_wrappers import Rope
|
||||
|
||||
return Rope()
|
||||
except ImportError:
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.warning("RoPEOp: CUDA kernel load failed: %s", e)
|
||||
return None
|
||||
|
||||
def _pytorch_fallback(self, q, k, cos, sin, interleave=False, **kwargs):
|
||||
cos = cos.float()
|
||||
sin = sin.float()
|
||||
rotary_dim = cos.size(-1) * 2
|
||||
head_dim = q.size(-1)
|
||||
if rotary_dim > head_dim:
|
||||
raise ValueError(
|
||||
f"rotary_dim ({rotary_dim}) > head_dim ({head_dim}): "
|
||||
f"cos last dim ({cos.size(-1)}) is too large for q"
|
||||
)
|
||||
partial = rotary_dim < head_dim
|
||||
if partial:
|
||||
q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
|
||||
k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
|
||||
else:
|
||||
q_rot, k_rot = q, k
|
||||
cos = cos.unsqueeze(-2) # (..., 1, half_dim)
|
||||
sin = sin.unsqueeze(-2)
|
||||
if interleave:
|
||||
cos = cos.repeat_interleave(2, dim=-1)
|
||||
sin = sin.repeat_interleave(2, dim=-1)
|
||||
q_embed = q_rot.float() * cos + _rotate_interleave(q_rot.float()) * sin
|
||||
k_embed = k_rot.float() * cos + _rotate_interleave(k_rot.float()) * sin
|
||||
else:
|
||||
cos = torch.cat((cos, cos), dim=-1)
|
||||
sin = torch.cat((sin, sin), dim=-1)
|
||||
q_embed = q_rot.float() * cos + _rotate_half(q_rot.float()) * sin
|
||||
k_embed = k_rot.float() * cos + _rotate_half(k_rot.float()) * sin
|
||||
if partial:
|
||||
q_embed = torch.cat([q_embed.to(q.dtype), q_pass], dim=-1)
|
||||
k_embed = torch.cat([k_embed.to(k.dtype), k_pass], dim=-1)
|
||||
else:
|
||||
q_embed = q_embed.to(q.dtype)
|
||||
k_embed = k_embed.to(k.dtype)
|
||||
return q_embed, k_embed
|
||||
|
||||
def pack(self, *args, **kwargs):
|
||||
"""Delegate to resolved backend's pack method if available."""
|
||||
if self._resolved_fn is None:
|
||||
self._resolve()
|
||||
if hasattr(self._resolved_fn, "pack"):
|
||||
return self._resolved_fn.pack(*args, **kwargs)
|
||||
return args
|
||||
|
||||
|
||||
class MRoPEOp(OpsProxy):
|
||||
"""Multi-head rotary position embedding (used by Qwen2.5-VL models).
|
||||
|
||||
Signature: m_rope(query_states, key_states, cos, sin, mrope_section, interleaved=False) -> (q_embed, k_embed)
|
||||
|
||||
cos/sin shape: (3, B, S, D//2). mrope_section like [16, 24, 24] specifies
|
||||
how many head dims each of T/H/W occupies. After split by mrope_section * 2,
|
||||
uses i % 3 to select from corresponding temporal/height/width rows.
|
||||
"""
|
||||
|
||||
@property
|
||||
def _external_accel_name(self):
|
||||
return "m_rope"
|
||||
|
||||
def _get_cuda_kernel(self):
|
||||
try:
|
||||
from wall_x.model.core.ops._cuda_wrappers import MRope
|
||||
|
||||
return MRope()
|
||||
except ImportError:
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.warning("MRoPEOp: CUDA kernel load failed: %s", e)
|
||||
return None
|
||||
|
||||
def _pytorch_fallback(
|
||||
self,
|
||||
query_states,
|
||||
key_states,
|
||||
cos,
|
||||
sin,
|
||||
mrope_section: List[int],
|
||||
interleaved=False,
|
||||
**kwargs,
|
||||
):
|
||||
cos = cos.float()
|
||||
sin = sin.float()
|
||||
# Double cos/sin along last dim (matching external_accel kernel's internal behavior)
|
||||
cos = torch.cat((cos, cos), dim=-1)
|
||||
sin = torch.cat((sin, sin), dim=-1)
|
||||
# List concatenation: [16,24,24] -> [16,24,24,16,24,24] (NOT element-wise multiply)
|
||||
mrope_section_doubled = mrope_section + mrope_section
|
||||
cos_split = torch.cat(
|
||||
[m[i % 3] for i, m in enumerate(cos.split(mrope_section_doubled, dim=-1))],
|
||||
dim=-1,
|
||||
).unsqueeze(
|
||||
2
|
||||
) # (B, S, 1, D)
|
||||
sin_split = torch.cat(
|
||||
[m[i % 3] for i, m in enumerate(sin.split(mrope_section_doubled, dim=-1))],
|
||||
dim=-1,
|
||||
).unsqueeze(2)
|
||||
q_embed = (query_states.float() * cos_split) + (
|
||||
_rotate_half(query_states.float()) * sin_split
|
||||
)
|
||||
k_embed = (key_states.float() * cos_split) + (
|
||||
_rotate_half(key_states.float()) * sin_split
|
||||
)
|
||||
return q_embed.to(query_states.dtype), k_embed.to(key_states.dtype)
|
||||
|
||||
def pack(self, *args, **kwargs):
|
||||
"""Delegate to resolved backend's pack method if available."""
|
||||
if self._resolved_fn is None:
|
||||
self._resolve()
|
||||
if hasattr(self._resolved_fn, "pack"):
|
||||
return self._resolved_fn.pack(*args, **kwargs)
|
||||
return args
|
||||
|
||||
|
||||
class RotPosEmbOp(OpsProxy):
|
||||
"""Rotary position embedding computation for ViT (used by Qwen2.5-VL vision encoder).
|
||||
|
||||
Signature: rot_pos_emb(inv_freq, grid_thw, spatial_merge_size) -> rotary_pos_emb
|
||||
"""
|
||||
|
||||
@property
|
||||
def _external_accel_name(self):
|
||||
return "rot_pos_emb"
|
||||
|
||||
def _get_cuda_kernel(self):
|
||||
try:
|
||||
from wall_x.model.core.ops._cuda_wrappers import RotPos
|
||||
|
||||
return RotPos()
|
||||
except ImportError:
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.warning("RotPosEmbOp: CUDA kernel load failed: %s", e)
|
||||
return None
|
||||
|
||||
def _pytorch_fallback(self, inv_freq, grid_thw, spatial_merge_size):
|
||||
if inv_freq.dtype != torch.float32:
|
||||
inv_freq = inv_freq.to(torch.float32)
|
||||
pos_ids = []
|
||||
for t, h, w in grid_thw:
|
||||
t, h, w = int(t), int(h), int(w)
|
||||
if h % spatial_merge_size != 0 or w % spatial_merge_size != 0:
|
||||
raise ValueError(
|
||||
f"grid h={h}, w={w} must be divisible by spatial_merge_size={spatial_merge_size}"
|
||||
)
|
||||
hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
|
||||
hpos_ids = (
|
||||
hpos_ids.reshape(
|
||||
h // spatial_merge_size,
|
||||
spatial_merge_size,
|
||||
w // spatial_merge_size,
|
||||
spatial_merge_size,
|
||||
)
|
||||
.permute(0, 2, 1, 3)
|
||||
.flatten()
|
||||
)
|
||||
wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
|
||||
wpos_ids = (
|
||||
wpos_ids.reshape(
|
||||
h // spatial_merge_size,
|
||||
spatial_merge_size,
|
||||
w // spatial_merge_size,
|
||||
spatial_merge_size,
|
||||
)
|
||||
.permute(0, 2, 1, 3)
|
||||
.flatten()
|
||||
)
|
||||
pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
|
||||
pos_ids = torch.cat(pos_ids, dim=0).to(inv_freq.device)
|
||||
max_grid_size = grid_thw[:, 1:].max()
|
||||
seq = torch.arange(max_grid_size, device=inv_freq.device, dtype=inv_freq.dtype)
|
||||
rotary_pos_emb_full = torch.outer(seq, inv_freq)
|
||||
rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
|
||||
return rotary_pos_emb.to(torch.float)
|
||||
|
||||
def pack(self, *args, **kwargs):
|
||||
"""Delegate to resolved backend's pack method if available."""
|
||||
if self._resolved_fn is None:
|
||||
self._resolve()
|
||||
if hasattr(self._resolved_fn, "pack"):
|
||||
return self._resolved_fn.pack(*args, **kwargs)
|
||||
return args
|
||||
|
||||
|
||||
rope = RoPEOp()
|
||||
m_rope = MRoPEOp()
|
||||
rot_pos_emb = RotPosEmbOp()
|
||||
Reference in New Issue
Block a user