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

Update Wall-X to 1.1.0 (#104)

Browse Source
This commit is contained in:
Starrick Liu
2026-06-15 11:40:00 +08:00
committed by GitHub
parent e23a586846
commit 72834e7de5
200 changed files with 33916 additions and 16771 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/model/core/ops/csrc/binding.cu
+69
View File
@@ -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);
}
wall_x/model/core/ops/csrc/common/activation_types.h
+11
View File
@@ -0,0 +1,11 @@
#pragma once
namespace wallx_cuda {
enum class ActivationType {
RELU,
SILU,
GELU
};
} // namespace wallx_cuda
wall_x/model/core/ops/csrc/common/activations.cuh
+30
View File
@@ -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
wall_x/model/core/ops/csrc/common/cuda_utils.h
+111
View File
@@ -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");
}
}
wall_x/model/core/ops/csrc/get_rope_index/get_rope_index.cu
+692
View File
@@ -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)
wall_x/model/core/ops/csrc/m_rope/m_rope_kernel.cu
+840
View File
@@ -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
wall_x/model/core/ops/csrc/permute_unpermute/permute.cu
+581
View File
@@ -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
wall_x/model/core/ops/csrc/rope/rope_kernel.cu
+1098
View File
File diff suppressed because it is too large Load Diff
wall_x/model/core/ops/csrc/rot_pos/rot_pos.cu
+284
View File
@@ -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
wall_x/model/core/ops/csrc/window_index/window_index.cu
+284
View File
@@ -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