feat: Major optimization and robustness improvements (#31)

This release introduces significant performance optimizations, memory efficiency
improvements, and enhanced system robustness:

🚀 Performance Optimizations:
- Add three new fused CUDA kernels (rope_index, rot_pos_emb, get_window_index)
  for accelerated multimodal preprocessing
- Implement FSDP2 support for distributed training with improved memory efficiency
- Add Torch.compile integration for additional performance gains
- Optimize memory usage: reduce peak allocation from 48GB to 24GB on 8-GPU setup

🔧 System Robustness:
- Fix missing token position inputs in prediction pipeline
- Add type-robust negation operations in RoPE CUDA kernels (half/bfloat16 support)
- Fix dataset root parameter initialization in LeRobot data loader
- Enhanced error handling and input validation across fusion operators

📚 Documentation & Usability:
- Add comprehensive memory usage benchmarks and hardware recommendations
- Update citation format with proper arXiv reference
- Improve training configuration documentation with quick start guide
- Add detailed API documentation for new fusion operators

🛠️ Technical Details:
- Version bump to 1.0.1
- New CUDA kernels: rope_index.cu, rot_pos.cu, window_index.cu
- FSDP2 state dict loading with distribute_tensor support
- Enhanced multimodal RoPE with 3D position encoding
- Window attention optimization for Vision Transformers

Breaking Changes: None - all changes are backward compatible

csrc/README.md

@@ -1,6 +1,6 @@
 # Fusion Operators (CSRC)
 
-High-performance CUDA kernels for accelerating model training.
+High-performance CUDA kernels for accelerating model training, with specialized support for multimodal and MoE architectures.
 
 ## Operators
 
@@ -16,7 +16,12 @@ High-performance CUDA kernels for accelerating model training.
 ### Multimodal RoPE
 - `rope`: Rotary Position Embedding forward pass
 - `rope_bwd`: RoPE backward pass
-- Support for multimodal inputs with configurable sections
+- `rope_index`: Generates position indices for multimodal RoPE
+- `rot_pos_emb`: Fused rotary position embedding computation
+
+### Vision Transformer Optimization
+- `get_window_index`: Window attention index generation
+
 
 ## Acknowledgments
 

csrc/ops.cu

@@ -1,6 +1,9 @@
 #include "dual_asym_grouped_gemm.h"
 #include "permute.h"
 #include "rope.h"
+#include "rope_index.h"
+#include "rot_pos.h"
+#include "window_index.h"
 
 #include <torch/extension.h>
 
@@ -12,4 +15,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("unpermute_bwd", &moe_recover_topK_bwd_op, "Token un-permutation backward kernel");
   m.def("rope", &launch_multimodal_rope_forward, "Multimodal RoPE forward kernel");
   m.def("rope_bwd", &launch_multimodal_rope_backward, "Multimodal RoPE backward kernel");
+  m.def("rope_index", &get_rope_index, "Get RoPE index kernel");
+  m.def("rot_pos_emb", &fused_rot_pos_emb_cuda, "Fused Rotary Position Embedding kernel");
+  m.def("get_window_index", &get_window_index_cuda, "Get window index kernel");
 }

csrc/rope.cu

@@ -13,8 +13,6 @@
 #include <c10/cuda/CUDAStream.h>
 #include <torch/extension.h>
 
-
-
 // Type traits for CUDA types
 template <typename T>
 struct CudaTypeTraits
@@ -81,7 +79,8 @@ __global__ void multimodal_rope_forward_kernel(
 
     // Process tokens in batches across warps
 
-    for (int seq_base = warp_id + seq_paral_idx * warps_per_block; seq_base < seq_len; seq_base += seq_paral_size * warps_per_block) {
+    for (int seq_base = warp_id + seq_paral_idx * warps_per_block; seq_base < seq_len; seq_base += seq_paral_size * warps_per_block)
+    {
         int seq_idx = seq_base;
         if (seq_idx >= seq_len)
             break;
@@ -119,16 +118,16 @@ __global__ void multimodal_rope_forward_kernel(
 
                 // Load cos/sin values (coalesced access)
                 int cos_sin_idx = section_idx * batch_size * seq_len * head_dim +
-                                    batch_idx * seq_len * head_dim +
-                                    seq_idx * head_dim + cos_sin_d;
+                                  batch_idx * seq_len * head_dim +
+                                  seq_idx * head_dim + cos_sin_d;
 
                 T cos_val = cos[cos_sin_idx];
                 T sin_val = sin[cos_sin_idx];
 
                 // Calculate tensor indices
                 int tensor_idx = batch_idx * total_heads * seq_len * head_dim +
-                                    actual_head_idx * seq_len * head_dim +
-                                    seq_idx * head_dim + dim_idx;
+                                 actual_head_idx * seq_len * head_dim +
+                                 seq_idx * head_dim + dim_idx;
 
                 // Get input value
                 T input_val = is_q_head ? q[tensor_idx] : k[tensor_idx];
@@ -143,7 +142,22 @@ __global__ void multimodal_rope_forward_kernel(
                 T rotate_val = is_q_head ? q[rotate_tensor_idx] : k[rotate_tensor_idx];
                 if (dim_idx < half_dim)
                 {
-                    rotate_val = -rotate_val; // First half: negate second half
+                    if constexpr (std::is_same_v<T, half>)
+                    {
+                        rotate_val = __hneg(rotate_val);
+                    }
+                    else if constexpr (std::is_same_v<T, __nv_bfloat16>)
+                    {
+#if __CUDA_ARCH__ >= 800                                 // BFloat16 support requires Ampere or newer
+                        rotate_val = __hneg(rotate_val); // __hneg works for bfloat16 in newer CUDA
+#else
+                        rotate_val = __float2bfloat16(-__bfloat162float(rotate_val));
+#endif
+                    }
+                    else
+                    {
+                        rotate_val = -rotate_val;
+                    }
                 }
 
                 // Apply RoPE: output = input * cos + rotate_half(input) * sin
@@ -234,13 +248,13 @@ __global__ void multimodal_rope_backward_kernel(
 
                 // Global cos/sin index
                 int cos_sin_idx = section_idx * batch_size * seq_len * head_dim +
-                                    batch_idx * seq_len * head_dim +
-                                    seq_idx * head_dim + cos_sin_d;
+                                  batch_idx * seq_len * head_dim +
+                                  seq_idx * head_dim + cos_sin_d;
 
                 // Tensor index for current position
                 int tensor_idx = batch_idx * total_heads * seq_len * head_dim +
-                                    actual_head_idx * seq_len * head_dim +
-                                    seq_idx * head_dim + dim_idx;
+                                 actual_head_idx * seq_len * head_dim +
+                                 seq_idx * head_dim + dim_idx;
 
                 // Load values
                 T cos_val = cos[cos_sin_idx];
@@ -274,8 +288,8 @@ __global__ void multimodal_rope_backward_kernel(
                 }
 
                 int paired_cos_sin_idx = paired_section_idx * batch_size * seq_len * head_dim +
-                                            batch_idx * seq_len * head_dim +
-                                            seq_idx * head_dim + paired_cos_sin_d;
+                                         batch_idx * seq_len * head_dim +
+                                         seq_idx * head_dim + paired_cos_sin_d;
                 T paired_sin_val = sin[paired_cos_sin_idx];
 
                 // === Compute input gradients (the only thing we need!) ===
@@ -441,7 +455,6 @@ void launch_multimodal_rope_forward(
     torch::Tensor q_out, torch::Tensor k_out,
     std::vector<int> mrope_section_doubled)
 {
-
     cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
 
     int batch_size = q.size(0);
@@ -463,10 +476,14 @@ void launch_multimodal_rope_forward(
     {
         data_type = 2;
     }
-    int *d_mrope_section_doubled;
-    cudaMalloc(&d_mrope_section_doubled, 3 * sizeof(int));
-    cudaMemcpyAsync(d_mrope_section_doubled, mrope_section_doubled.data(), 3 * sizeof(int),
-                    cudaMemcpyHostToDevice, stream);
+
+    auto mrope_tensor = torch::from_blob(
+        mrope_section_doubled.data(),
+        {3},
+        torch::TensorOptions().dtype(torch::kInt32)
+    ).to(q.device(), /*non_blocking=*/true);
+
+    int *d_mrope_section_doubled = static_cast<int*>(mrope_tensor.data_ptr());
 
     switch (data_type)
     {
@@ -497,7 +514,6 @@ void launch_multimodal_rope_backward(
     torch::Tensor grad_q, torch::Tensor grad_k,
     std::vector<int> mrope_section_doubled)
 {
-
     cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
 
     int batch_size = q.size(0);
@@ -519,10 +535,14 @@ void launch_multimodal_rope_backward(
     {
         data_type = 2;
     }
-    int *d_mrope_section_doubled;
-    cudaMalloc(&d_mrope_section_doubled, 3 * sizeof(int));
-    cudaMemcpyAsync(d_mrope_section_doubled, mrope_section_doubled.data(), 3 * sizeof(int),
-                    cudaMemcpyHostToDevice, stream);
+
+    auto mrope_tensor = torch::from_blob(
+        mrope_section_doubled.data(),
+        {3},
+        torch::TensorOptions().dtype(torch::kInt32)
+    ).to(q.device(), /*non_blocking=*/true);
+
+    int *d_mrope_section_doubled = static_cast<int*>(mrope_tensor.data_ptr());
 
     switch (data_type)
     {

csrc/rope_index.cu

@@ -0,0 +1,569 @@
+#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 <optional>
+#include <cstdint> // for int64_t
+
+#include <cuda_runtime.h>
+#include <cuda.h>
+#include <cub/cub.cuh>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/util/BFloat16.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+
+#define MAX_SEQ_LEN 8192
+#define MAX_VISION_TOKENS 64
+#define WARP_SIZE 32
+#define MAX_THREADS_PER_BLOCK 1024
+
+struct VisionDescriptor
+{
+    int64_t start_pos;
+    int64_t token_pos;
+    int64_t patch_count;
+    int64_t grid_t, grid_h, grid_w;
+    float time_interval;
+    int64_t is_video;
+    int64_t position_offset;
+};
+
+__device__ __forceinline__ int64_t fast_div(int64_t a, int64_t b)
+{
+    return __float2ll_rd(__ll2float_rn(a) * __frcp_rn(__ll2float_rn(b)));
+}
+
+__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;
+}
+
+__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,)
+    int64_t *video_counts,                 // (batch_size,)
+    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];
+
+    int64_t thread_image_count = 0;
+    int64_t thread_video_count = 0;
+
+    for (int64_t i = thread_idx; i < seq_len - 1; i += blockDim.x)
+    {
+        if ((attention_mask != nullptr) && attention_mask[batch_idx * seq_len + i] == 0)
+            continue;
+
+        int64_t token_id = input_ids[batch_idx * seq_len + i];
+
+        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++;
+            }
+        }
+    }
+
+    shared_image_counts[thread_idx] = thread_image_count;
+    shared_video_counts[thread_idx] = thread_video_count;
+    __syncthreads();
+
+    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();
+    }
+
+    if (thread_idx == 0)
+    {
+        image_counts[batch_idx] = shared_image_counts[0];
+        video_counts[batch_idx] = shared_video_counts[0];
+    }
+}
+
+
+__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,)
+    int64_t *text_lengths,               // (batch_size, MAX_VISION_TOKENS+1)
+    int64_t *position_offsets,           // (batch_size, MAX_VISION_TOKENS+1)
+    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)
+{
+    int64_t batch_idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (batch_idx >= batch_size)
+        return;
+
+    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];
+    }
+
+    int64_t vision_count = 0;
+    int64_t current_pos = 0;
+    int64_t position_offset = 0;
+
+    for (int64_t i = 0; i < seq_len - 1; i++)
+    {
+        if ((attention_mask != nullptr) && attention_mask[batch_idx * seq_len + i] == 0)
+            continue;
+
+        int64_t token_id = input_ids[batch_idx * seq_len + i];
+
+        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)
+            {
+                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);
+
+                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)
+                {
+                    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
+                {
+                    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++;
+                }
+
+                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;
+
+                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;
+
+                    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++;
+                }
+
+                i = vision_pos + patch_count - 1;
+            }
+        }
+    }
+
+    vision_counts[batch_idx] = vision_count;
+
+    int64_t effective_len = seq_len;
+    if (attention_mask != nullptr)
+    {
+        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;
+}
+
+
+__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)
+    int64_t *mrope_deltas,                   // (batch_size,)
+    const int64_t batch_size,
+    const int64_t seq_len,
+    const float tokens_per_second)
+{
+    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];
+
+    int64_t shared_vision_count = vision_counts[batch_idx];
+
+    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];
+    }
+
+    int64_t thread_max_position = -1;
+
+    __syncthreads();
+
+    for (int64_t token_idx = thread_idx; token_idx < seq_len; token_idx += blockDim.x)
+    {
+        bool is_valid_token = true;
+        if (attention_mask != nullptr)
+        {
+            is_valid_token = attention_mask[batch_idx * seq_len + token_idx] != 0;
+        }
+
+        if (!is_valid_token)
+        {
+            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;
+        }
+
+        int64_t segment_idx = -9999999;
+        int64_t local_pos = token_idx;
+
+        for (int64_t v = 0; v < shared_vision_count; v++)
+        {
+            if (token_idx < shared_visions[v].token_pos)
+            {
+                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)
+            {
+                segment_idx = -(v + 1);
+                local_pos = token_idx - shared_visions[v].token_pos;
+                break;
+            }
+        }
+
+        if (segment_idx == -9999999)
+        {
+            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)
+        {
+            int64_t offset = shared_position_offsets[segment_idx];
+            pos_t = pos_h = pos_w = offset + local_pos;
+        }
+        else
+        {
+            int64_t vision_idx = -(segment_idx + 1);
+            const VisionDescriptor &desc = shared_visions[vision_idx];
+
+            int64_t t, h, w;
+            get_3d_coords(local_pos, desc.grid_h, desc.grid_w, t, h, w);
+
+            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;
+        }
+
+        position_ids[0 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_t;
+        position_ids[1 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_h;
+        position_ids[2 * batch_size * seq_len + batch_idx * seq_len + token_idx] = pos_w;
+
+        int64_t max_pos = max(pos_t, max(pos_h, pos_w));
+        thread_max_position = max(thread_max_position, max_pos);
+    }
+
+    shared_max_positions[thread_idx] = thread_max_position;
+    __syncthreads();
+
+    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();
+    }
+
+    if (thread_idx == 0)
+    {
+        int64_t global_max_position = shared_max_positions[0];
+        mrope_deltas[batch_idx] = global_max_position + 1 - seq_len;
+    }
+}
+
+
+void launch_optimized_3d_rope_kernel(
+    const int64_t *input_ids,
+    const int64_t *attention_mask,
+    const int64_t *image_grid_thw,
+    const int64_t *video_grid_thw,
+    const float *second_per_grid_ts,
+    int64_t *position_ids,
+    int64_t *mrope_deltas,
+    int64_t batch_size,
+    int64_t seq_len,
+    int64_t spatial_merge_size,
+    int64_t image_token_id,
+    int64_t video_token_id,
+    int64_t vision_start_token_id,
+    float tokens_per_second)
+{
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
+
+    torch::Device device(torch::kCUDA, at::cuda::current_device());
+
+    auto options = torch::TensorOptions().dtype(torch::kInt64).device(device);
+
+    auto vision_desc_tensor = torch::empty(
+        {batch_size * MAX_VISION_TOKENS * static_cast<int64_t>(sizeof(VisionDescriptor))},
+        torch::TensorOptions().dtype(torch::kUInt8).device(device)
+    );
+    VisionDescriptor *d_vision_desc = reinterpret_cast<VisionDescriptor*>(vision_desc_tensor.data_ptr());
+
+    auto vision_counts_tensor = torch::empty({batch_size}, options);
+    auto text_lengths_tensor = torch::empty({batch_size * (MAX_VISION_TOKENS + 1)}, options);
+    auto position_offsets_tensor = torch::empty({batch_size * (MAX_VISION_TOKENS + 1)}, options);
+    auto image_counts_tensor = torch::empty({batch_size}, options);
+    auto video_counts_tensor = torch::empty({batch_size}, options);
+
+    int64_t *d_vision_counts = vision_counts_tensor.data_ptr<int64_t>();
+    int64_t *d_text_lengths = text_lengths_tensor.data_ptr<int64_t>();
+    int64_t *d_position_offsets = position_offsets_tensor.data_ptr<int64_t>();
+    int64_t *d_image_counts = image_counts_tensor.data_ptr<int64_t>();
+    int64_t *d_video_counts = video_counts_tensor.data_ptr<int64_t>();
+
+    dim3 index_grid(static_cast<unsigned int>(batch_size));
+    dim3 index_block(256);
+
+    compute_vision_counts<<<index_grid, index_block, 0, stream>>>(
+        input_ids, attention_mask,
+        d_image_counts, d_video_counts,
+        batch_size, seq_len, image_token_id, video_token_id, vision_start_token_id);
+
+    int64_t threads_per_block = std::min(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, attention_mask, image_grid_thw, video_grid_thw,
+        second_per_grid_ts, 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);
+
+    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, attention_mask, d_vision_desc, d_vision_counts,
+        d_text_lengths, d_position_offsets, position_ids, mrope_deltas,
+        batch_size, seq_len, tokens_per_second);
+
+    AT_CUDA_CHECK(cudaGetLastError());
+}
+
+std::tuple<torch::Tensor, torch::Tensor> get_rope_index(
+    const torch::optional<torch::Tensor> &input_ids,
+    const torch::optional<torch::Tensor> &image_grid_thw,
+    const torch::optional<torch::Tensor> &video_grid_thw,
+    const torch::optional<torch::Tensor> &second_per_grid_ts,
+    const torch::optional<torch::Tensor> &attention_mask,
+    int spatial_merge_size,
+    int image_token_id,
+    int video_token_id,
+    int vision_start_token_id,
+    float tokens_per_second)
+{
+    TORCH_CHECK(input_ids.has_value(), "input_ids cannot be None");
+    TORCH_CHECK(input_ids->dim() == 2, "input_ids must be 2D tensor (batch_size, seq_len)");
+
+    const auto batch_size = input_ids->size(0);
+    const auto seq_len = input_ids->size(1);
+    const auto device = input_ids->device();
+
+    TORCH_CHECK(device.is_cuda(), "All tensors must be on CUDA device");
+
+    torch::Tensor input_ids_tensor = input_ids->contiguous();
+    const int64_t *input_ids_ptr = input_ids_tensor.data_ptr<int64_t>();
+
+    torch::Tensor attention_mask_tensor;
+    const int64_t *attention_mask_ptr = nullptr;
+
+    if (attention_mask.has_value()) {
+        attention_mask_tensor = attention_mask->contiguous();
+        attention_mask_ptr = attention_mask_tensor.data_ptr<int64_t>();
+    }
+
+    torch::Tensor image_grid_thw_tensor;
+    const int64_t *image_grid_thw_ptr = nullptr;
+
+    if (image_grid_thw.has_value()) {
+        TORCH_CHECK(image_grid_thw->dim() == 2 && image_grid_thw->size(1) == 3,
+                    "image_grid_thw must be shape (num_images, 3)");
+        image_grid_thw_tensor = image_grid_thw->contiguous();
+        image_grid_thw_ptr = image_grid_thw_tensor.data_ptr<int64_t>();
+    }
+
+    torch::Tensor video_grid_thw_tensor;
+    const int64_t *video_grid_thw_ptr = nullptr;
+
+    if (video_grid_thw.has_value()) {
+        TORCH_CHECK(video_grid_thw->dim() == 2 && video_grid_thw->size(1) == 3,
+                    "video_grid_thw must be shape (num_videos, 3)");
+        video_grid_thw_tensor = video_grid_thw->contiguous();
+        video_grid_thw_ptr = video_grid_thw_tensor.data_ptr<int64_t>();
+    }
+
+    torch::Tensor second_per_grid_ts_tensor;
+    const float *second_per_grid_ts_ptr = nullptr;
+
+    if (second_per_grid_ts.has_value()) {
+        TORCH_CHECK(second_per_grid_ts->dim() == 1,
+                    "second_per_grid_ts must be 1D tensor");
+        second_per_grid_ts_tensor = second_per_grid_ts->contiguous();
+        second_per_grid_ts_ptr = second_per_grid_ts_tensor.data_ptr<float>();
+    }
+
+    if (!image_grid_thw.has_value() && !video_grid_thw.has_value()) {
+        torch::Tensor position_ids;
+        torch::Tensor mrope_deltas;
+
+        if (attention_mask.has_value()) {
+            auto cumsum_result = attention_mask_tensor.to(torch::kInt64).cumsum(-1) - 1;
+            position_ids = cumsum_result.masked_fill_(attention_mask_tensor.eq(0), 1);
+            position_ids = position_ids.unsqueeze(0).expand({3, -1, -1});
+
+            auto max_position_ids = std::get<0>(std::get<0>(position_ids.max(0)).max(-1, true));
+            mrope_deltas = max_position_ids + 1 - attention_mask_tensor.size(-1);
+            mrope_deltas = mrope_deltas.view({batch_size, 1});
+        } else {
+            auto pos_range = torch::arange(seq_len, torch::TensorOptions().dtype(torch::kInt64).device(device));
+            position_ids = pos_range.view({1, 1, -1}).expand({3, batch_size, -1});
+            mrope_deltas = torch::zeros({batch_size, 1}, torch::TensorOptions().dtype(torch::kInt64).device(device));
+        }
+
+        return std::make_tuple(position_ids, mrope_deltas);
+    }
+
+    auto position_ids = torch::empty({3, batch_size, seq_len},
+                                     torch::TensorOptions().dtype(torch::kInt64).device(device));
+    auto mrope_deltas = torch::empty({batch_size, 1},
+                                     torch::TensorOptions().dtype(torch::kInt64).device(device));
+
+    launch_optimized_3d_rope_kernel(
+        input_ids_ptr,
+        attention_mask_ptr,
+        image_grid_thw_ptr,
+        video_grid_thw_ptr,
+        second_per_grid_ts_ptr,
+        position_ids.data_ptr<int64_t>(),
+        mrope_deltas.data_ptr<int64_t>(),
+        static_cast<int64_t>(batch_size),
+        static_cast<int64_t>(seq_len),
+        static_cast<int64_t>(spatial_merge_size),
+        static_cast<int64_t>(image_token_id),
+        static_cast<int64_t>(video_token_id),
+        static_cast<int64_t>(vision_start_token_id),
+        tokens_per_second);
+
+    return std::make_tuple(position_ids, mrope_deltas);
+}

csrc/rope_index.h

@@ -0,0 +1,13 @@
+#include <torch/extension.h>
+
+std::tuple<torch::Tensor, torch::Tensor> get_rope_index(
+    const torch::optional<torch::Tensor> &input_ids,
+    const torch::optional<torch::Tensor> &image_grid_thw,
+    const torch::optional<torch::Tensor> &video_grid_thw,
+    const torch::optional<torch::Tensor> &second_per_grid_ts,
+    const torch::optional<torch::Tensor> &attention_mask,
+    int spatial_merge_size,
+    int image_token_id,
+    int video_token_id,
+    int vision_start_token_id,
+    float tokens_per_second);

csrc/rot_pos.cu

@@ -0,0 +1,332 @@
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+
+// CUDA kernel for fused rotary position embedding computation - int32 version
+__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];
+
+    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 = inv_freq[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
+}
+
+// CUDA kernel for fused rotary position embedding computation - int64 version
+__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];
+
+    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 = inv_freq[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
+}
+
+// Parallel computation of token counts per grid - int32 version
+__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;
+}
+
+// Parallel computation of token counts per grid - int64 version
+__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;
+}
+
+// Implementation for int32
+torch::Tensor fused_rot_pos_emb_cuda_int32(
+    torch::Tensor inv_freq, // [dim/2]
+    torch::Tensor grid_thw, // [num_grids, 3]
+    int spatial_merge_size)
+{
+    TORCH_CHECK(inv_freq.dim() == 1, "inv_freq must be 1-dimensional");
+    TORCH_CHECK(inv_freq.is_cuda(), "inv_freq must be a CUDA tensor");
+    TORCH_CHECK(inv_freq.scalar_type() == torch::kFloat32, "inv_freq must be float32");
+
+    TORCH_CHECK(grid_thw.dim() == 2, "grid_thw must be 2-dimensional");
+    TORCH_CHECK(grid_thw.size(1) == 3, "grid_thw must have shape [num_grids, 3]");
+    TORCH_CHECK(grid_thw.is_cuda(), "grid_thw must be a CUDA tensor");
+    TORCH_CHECK(grid_thw.scalar_type() == torch::kInt32, "grid_thw must be int32");
+
+    TORCH_CHECK(spatial_merge_size > 0, "spatial_merge_size must be positive");
+
+    const int dim_half = inv_freq.size(0);
+    const int num_grids = grid_thw.size(0);
+
+    auto token_counts = torch::zeros({num_grids}, torch::TensorOptions().dtype(torch::kInt32).device(grid_thw.device()));
+    const int threads = 256;
+    const int blocks = (num_grids + threads - 1) / threads;
+
+    compute_token_counts_kernel_int32<<<blocks, threads>>>(
+        grid_thw.data_ptr<int32_t>(),
+        token_counts.data_ptr<int32_t>(),
+        spatial_merge_size,
+        num_grids);
+
+    auto cumsum_tokens = torch::cat({torch::zeros({1}, torch::TensorOptions().dtype(torch::kInt32).device(grid_thw.device())),
+                                     token_counts.cumsum(0).to(torch::kInt32)},
+                                    0);
+
+    cudaDeviceSynchronize();
+
+    int64_t total_tokens = cumsum_tokens[-1].item<int64_t>();
+    TORCH_CHECK(total_tokens > 0, "total_tokens must be positive");
+
+    auto output = torch::zeros({total_tokens, dim_half * 2},
+                               torch::TensorOptions().dtype(torch::kFloat32).device(inv_freq.device()));
+
+    const int threads_per_block = 256;
+    const int64_t num_elements = total_tokens * dim_half;
+    const int num_blocks = static_cast<int>((num_elements + threads_per_block - 1) / threads_per_block);
+
+    fused_rot_pos_emb_kernel_int32<<<num_blocks, threads_per_block>>>(
+        inv_freq.data_ptr<float>(),
+        grid_thw.data_ptr<int32_t>(),
+        output.data_ptr<float>(),
+        cumsum_tokens.data_ptr<int32_t>(),
+        dim_half,
+        spatial_merge_size,
+        num_grids);
+
+    cudaDeviceSynchronize();
+
+    TORCH_CHECK(output.scalar_type() == torch::kFloat32, "Output must be float32");
+    TORCH_CHECK(output.size(0) == total_tokens, "Output token count mismatch");
+    TORCH_CHECK(output.size(1) == dim_half * 2, "Output dimension mismatch");
+
+    return output;
+}
+
+// Implementation for int64
+torch::Tensor fused_rot_pos_emb_cuda_int64(
+    torch::Tensor inv_freq, // [dim/2]
+    torch::Tensor grid_thw, // [num_grids, 3]
+    int spatial_merge_size)
+{
+    TORCH_CHECK(inv_freq.dim() == 1, "inv_freq must be 1-dimensional");
+    TORCH_CHECK(inv_freq.is_cuda(), "inv_freq must be a CUDA tensor");
+    TORCH_CHECK(inv_freq.scalar_type() == torch::kFloat32, "inv_freq must be float32");
+
+    TORCH_CHECK(grid_thw.dim() == 2, "grid_thw must be 2-dimensional");
+    TORCH_CHECK(grid_thw.size(1) == 3, "grid_thw must have shape [num_grids, 3]");
+    TORCH_CHECK(grid_thw.is_cuda(), "grid_thw must be a CUDA tensor");
+    TORCH_CHECK(grid_thw.scalar_type() == torch::kInt64, "grid_thw must be int64");
+
+    TORCH_CHECK(spatial_merge_size > 0, "spatial_merge_size must be positive");
+
+    const int dim_half = inv_freq.size(0);
+    const int num_grids = grid_thw.size(0);
+
+    auto token_counts = torch::zeros({num_grids}, torch::TensorOptions().dtype(torch::kInt64).device(grid_thw.device()));
+    const int threads = 256;
+    const int blocks = (num_grids + threads - 1) / threads;
+
+    compute_token_counts_kernel_int64<<<blocks, threads>>>(
+        grid_thw.data_ptr<int64_t>(),
+        token_counts.data_ptr<int64_t>(),
+        spatial_merge_size,
+        num_grids);
+
+    auto cumsum_tokens = torch::cat({torch::zeros({1}, torch::TensorOptions().dtype(torch::kInt64).device(grid_thw.device())),
+                                     token_counts.cumsum(0).to(torch::kInt64)},
+                                    0);
+
+    cudaDeviceSynchronize();
+
+    int64_t total_tokens = cumsum_tokens[-1].item<int64_t>();
+    TORCH_CHECK(total_tokens > 0, "total_tokens must be positive");
+
+    auto output = torch::zeros({total_tokens, dim_half * 2},
+                               torch::TensorOptions().dtype(torch::kFloat32).device(inv_freq.device()));
+
+    const int threads_per_block = 256;
+    const int64_t num_elements = total_tokens * dim_half;
+    const int num_blocks = static_cast<int>((num_elements + threads_per_block - 1) / threads_per_block);
+
+    fused_rot_pos_emb_kernel_int64<<<num_blocks, threads_per_block>>>(
+        inv_freq.data_ptr<float>(),
+        grid_thw.data_ptr<int64_t>(),
+        output.data_ptr<float>(),
+        cumsum_tokens.data_ptr<int64_t>(),
+        dim_half,
+        spatial_merge_size,
+        num_grids);
+
+    cudaDeviceSynchronize();
+
+    TORCH_CHECK(output.scalar_type() == torch::kFloat32, "Output must be float32");
+    TORCH_CHECK(output.size(0) == total_tokens, "Output token count mismatch");
+    TORCH_CHECK(output.size(1) == dim_half * 2, "Output dimension mismatch");
+
+    return output;
+}
+
+// Main function that dispatches based on grid_thw scalar type
+torch::Tensor fused_rot_pos_emb_cuda(
+    torch::Tensor inv_freq,
+    torch::Tensor grid_thw,
+    int spatial_merge_size)
+{
+    if (grid_thw.scalar_type() == torch::kInt32)
+    {
+        return fused_rot_pos_emb_cuda_int32(inv_freq, grid_thw, spatial_merge_size);
+    }
+    else if (grid_thw.scalar_type() == torch::kInt64)
+    {
+        return fused_rot_pos_emb_cuda_int64(inv_freq, grid_thw, spatial_merge_size);
+    }
+    else
+    {
+        TORCH_CHECK(false, "Unsupported grid_thw scalar type: ", grid_thw.scalar_type());
+    }
+}

csrc/rot_pos.h

@@ -0,0 +1,6 @@
+#include <torch/extension.h>
+
+torch::Tensor fused_rot_pos_emb_cuda(
+    torch::Tensor inv_freq,
+    torch::Tensor grid_thw,
+    int spatial_merge_size);

csrc/window_index.cu

@@ -0,0 +1,286 @@
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+#include <device_launch_parameters.h>
+#include <cstdint>
+#include <algorithm>
+
+__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;
+        }
+    }
+}
+
+std::tuple<torch::Tensor, torch::Tensor> get_window_index_cuda(
+    torch::Tensor grid_thw,
+    int spatial_merge_size,
+    int vit_merger_window_size,
+    int patch_size,
+    int spatial_merge_unit)
+{
+    TORCH_CHECK(grid_thw.is_cuda(), "grid_thw must be a CUDA tensor");
+    TORCH_CHECK(grid_thw.dim() == 2 && grid_thw.size(1) == 3);
+    TORCH_CHECK(grid_thw.dtype() == torch::kInt32);
+
+    int num_grids = grid_thw.size(0);
+    if (num_grids == 0)
+    {
+        return std::make_tuple(
+            torch::empty({0}, grid_thw.options()),
+            torch::zeros({1}, grid_thw.options()));
+    }
+
+    const int *d_grid_thw = grid_thw.data_ptr<int>();
+    auto options = grid_thw.options();
+
+    auto grid_thw_cpu = grid_thw.cpu();
+    int max_grid_t = 0;
+    for (int i = 0; i < num_grids; i++)
+    {
+        max_grid_t = std::max(max_grid_t, grid_thw_cpu[i][0].item<int>());
+    }
+
+    auto grid_info_tensor = torch::empty({num_grids, 6}, options);
+    auto global_totals_tensor = torch::zeros({2}, options);
+
+    int *d_grid_info = grid_info_tensor.data_ptr<int>();
+    int *d_global_totals = global_totals_tensor.data_ptr<int>();
+
+    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);
+
+    auto totals_cpu = global_totals_tensor.cpu();
+    int total_elements = totals_cpu[0].item<int>();
+    int total_windows = totals_cpu[1].item<int>();
+
+    if (total_elements == 0 || total_windows == 0)
+    {
+        return std::make_tuple(
+            torch::empty({0}, options),
+            torch::zeros({1}, options));
+    }
+
+    torch::Tensor window_indices = torch::empty({total_elements}, options);
+    torch::Tensor cu_window_seqlens = torch::empty({total_windows + 1}, options);
+
+    int *d_window_indices = window_indices.data_ptr<int>();
+    int *d_cu_window_seqlens = cu_window_seqlens.data_ptr<int>();
+
+    auto window_counts_tensor = torch::empty({total_windows}, options);
+    int *d_window_counts = window_counts_tensor.data_ptr<int>();
+
+    dim3 blocks2(max_grid_t, num_grids);
+    dim3 threads2(256);
+
+    compute_window_counts<<<blocks2, threads2>>>(
+        d_grid_thw, d_grid_info, d_window_counts,
+        vit_merger_window_size, spatial_merge_unit, num_grids);
+
+    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);
+
+    return std::make_tuple(window_indices, cu_window_seqlens);
+}

csrc/window_index.h

@@ -0,0 +1,9 @@
+#include <torch/extension.h>
+#include <tuple>
+
+std::tuple<torch::Tensor, torch::Tensor> get_window_index_cuda(
+    torch::Tensor grid_thw,
+    int spatial_merge_size,
+    int vit_merger_window_size,
+    int patch_size,
+    int spatial_merge_unit);