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

.gitignore

@@ -196,7 +196,6 @@ bin/
 # Media files
 *.jpg
 *.jpeg
-*.png
 *.gif
 *.mp4
 *.avi
@@ -244,6 +243,7 @@ mlruns/
 # Cache directories
 .cache/
 cache/
+.ruff_cache/
 
 # =============================================================================
 # FONTS AND ASSETS (if not part of the project)

README.md

@@ -30,7 +30,7 @@ By creating a direct feedback loop between the model's decisions and the body's
 This repository provides the training and inference code that supports our WALL series open-source embodied foundation models. It includes end-to-end pipelines for data preparation (LeRobot), model configuration, flow-matching and FAST action branches, and evaluation utilities for real and simulated robots.
 
 ## News
-- We introduce [**WALL-OSS**](https://x2robot.com/en/research/68bc2cde8497d7f238dde690), an end-to-end embodied foundation model that leverages large-scale multimodal pretraining to achieve (1) embodiment-aware vision–language understanding, (2) strong language–action association, and (3) robust manipulation capability.
+- We introduce [**WALL-OSS: Igniting VLMs toward the Embodied Space**](https://x2robot.com/en/research/68bc2cde8497d7f238dde690), an end-to-end embodied foundation model that leverages large-scale multimodal pretraining to achieve (1) embodiment-aware vision–language understanding, (2) strong language–action association, and (3) robust manipulation capability.
 
 ## Models
 - WALL-OSS-FLOW: https://huggingface.co/x-square-robot/wall-oss-flow
@@ -83,6 +83,8 @@ bash ./workspace/lerobot_example/run.sh
 
 ## Inference
 
+### Basic Action Inference
+
 For model inference, please refer to:
 
 ```bash
@@ -95,28 +97,47 @@ This script demonstrates how to:
 - Run inference in validation mode with proper data types (bfloat16)
 - Validate model outputs and check for numerical stability
 
+### Open-Loop Evaluation
+
 To generate an open-loop comparison plot, please follow:
 
 ```bash
 python ./scripts/draw_openloop_plot.py
 ```
 
-To run VQA inference, please follow:
+### VQA Inference and Chain-of-Thought Testing
+
+To run VQA inference and test the model's Chain-of-Thought (COT) reasoning capabilities, please follow:
 
 ```bash
 python ./scripts/vqa_inference.py
 ```
 
+This script can be used to test the model's COT reasoning abilities for embodied tasks. Below is an example of COT testing:
+
+**Input Image:**
+
+![COT Example Frame](assets/cot_example_frame.png)
+
+**Input Text:**
+```
+To move the red block in the plate with same color, what should you do next? Think step by step.
+```
+
+**Model Output (COT Reasoning):**
+```
+To move the red block in the plate with the same color, you should first locate the red block. It is currently positioned on the table, not in the plate. Then, you should carefully grasp the red block using your fingers. Next, you should use your hand to lift the red block from the table and place it into the plate that is also red in color. Ensure that the red block is securely placed in the plate without slipping or falling.
+```
+
 ## 📚 Cite Us
 
 If you find WALL-OSS models useful, please cite:
 
 ```bibtex
-@misc{walloss_paper_2025,
+@article{zhai2025igniting,
-  title        = {WALL-OSS: Igniting VLMs toward the Embodied Space},
+  title   = {Igniting VLMs Toward the Embodied Space},
-  author       = {X Square Robot},
+  author  = {Zhai, Andy and Liu, Brae and Fang, Bruno and Cai, Chalse and Ma, Ellie and Yin, Ethan and Wang, Hao and Zhou, Hugo and Wang, James and Shi, Lights and Liang, Lucy and Wang, Make and Wang, Qian and Gan, Roy and Yu, Ryan and Li, Shalfun and Liu, Starrick and Chen, Sylas and Chen, Vincent and Xu, Zach},
-  year         = {2025},
+  journal = {arXiv preprint arXiv:2509.11766},
-  howpublished = {\url{https://x2robot.cn-wlcb.ufileos.com/wall_oss.pdf}},
+  year    = {2025}
-  note         = {White paper}
 }
 ```

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 +
+                                  batch_idx * seq_len * head_dim +
-                                    seq_idx * head_dim + cos_sin_d;
+                                  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 +
+                                 actual_head_idx * seq_len * head_dim +
-                                    seq_idx * head_dim + dim_idx;
+                                 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 +
+                                  batch_idx * seq_len * head_dim +
-                                    seq_idx * head_dim + cos_sin_d;
+                                  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 +
+                                 actual_head_idx * seq_len * head_dim +
-                                    seq_idx * head_dim + dim_idx;
+                                 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 +
+                                         batch_idx * seq_len * head_dim +
-                                            seq_idx * head_dim + paired_cos_sin_d;
+                                         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));
+    auto mrope_tensor = torch::from_blob(
-    cudaMemcpyAsync(d_mrope_section_doubled, mrope_section_doubled.data(), 3 * sizeof(int),
+        mrope_section_doubled.data(),
-                    cudaMemcpyHostToDevice, stream);
+        {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));
+    auto mrope_tensor = torch::from_blob(
-    cudaMemcpyAsync(d_mrope_section_doubled, mrope_section_doubled.data(), 3 * sizeof(int),
+        mrope_section_doubled.data(),
-                    cudaMemcpyHostToDevice, stream);
+        {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);

scripts/vqa_inference.py

@@ -68,13 +68,15 @@ if __name__ == "__main__":
     wrapper = VQAWrapper(model_path=MODEL_PATH_FOR_MODULE_TEST)
 
     try:
-        test_image_url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+        test_question = "To move the red block in the plate with same color, what should you do next? Think step by step."
-        test_question = "What is written on the sign?"
 
-        # img = Image.open("/path/to/your/local/image.jpg").convert("RGB")
+        # Local Image
-        import requests
+        img = Image.open("/path/to/wall-x/assets/cot_example_frame.png").convert("RGB")
+        # Internet Image
+        # import requests
+        # test_image_url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+        # img = Image.open(requests.get(test_image_url, stream=True).raw).convert("RGB")
 
-        img = Image.open(requests.get(test_image_url, stream=True).raw).convert("RGB")
         answer = wrapper.generate(img, test_question)
 
         print("model answer:", answer)

setup.py

@@ -36,6 +36,9 @@ ext_modules = [
             "csrc/dual_asym_grouped_gemm.cu",
             "csrc/permute.cu",
             "csrc/rope.cu",
+            "csrc/rope_index.cu",
+            "csrc/rot_pos.cu",
+            "csrc/window_index.cu",
         ],
         include_dirs=[f"{cwd}/3rdparty/cutlass/include/", f"{cwd}/csrc"],
         extra_compile_args={
@@ -47,7 +50,7 @@ ext_modules = [
 
 setup(
     name="wall_x",
-    version="1.0.0",
+    version="1.0.1",
     author="X2Robot Team",
     classifiers=[
         "Programming Language :: Python :: 3",

train_qact.py

@@ -3,6 +3,7 @@ import json
 import time
 import yaml
 import wandb
+import accelerate
 from argparse import ArgumentParser
 from accelerate import (
     Accelerator,
@@ -38,9 +39,32 @@ def setup_accelerator(config):
     ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
     accelerator_dataloader_config = DataLoaderConfiguration(dispatch_batches=False)
 
+    if config.get("FSDP2", False):
+        # Use Fully Sharded Data Parallel (FSDP) version 2
+        fsdp_plugin = accelerate.utils.dataclasses.FullyShardedDataParallelPlugin(
+            fsdp_version=2, reshard_after_forward=True
+        )
+        print("[INFO] Using FSDP version 2 for distributed training")
+    else:
+        fsdp_plugin = None
+
+    if config.get("torch_compile", False):
+        # Use Torch Dynamo for compilation
+        dynamo_plugin = accelerate.utils.TorchDynamoPlugin(
+            backend="inductor",
+            mode="default",
+            fullgraph=False,
+            dynamic=False,
+        )
+        print("[INFO] Using Torch Dynamo for compilation")
+    else:
+        dynamo_plugin = None
+
     accelerator = Accelerator(
         kwargs_handlers=[ddp_kwargs],
         mixed_precision="bf16",
+        fsdp_plugin=fsdp_plugin,
+        dynamo_plugin=dynamo_plugin,
         dataloader_config=accelerator_dataloader_config,
         gradient_accumulation_steps=config.get("gradient_accumulation_steps", 1),
     )

wall_x/data/load_lerobot_dataset.py

@@ -424,8 +424,9 @@ def load_lerobot_data(
 
     # repo_id = "lerobot/aloha_mobile_cabinet"
     repo_id = lerobot_config.get("repo_id", "lerobot/aloha_mobile_cabinet")
+    root = lerobot_config.get("root", None)
     dataset = LeRobotDataset(
-        repo_id, delta_timestamps=delta_timestamps, video_backend="pyav"
+        repo_id, root=root, delta_timestamps=delta_timestamps, video_backend="pyav"
     )
 
     if rank == 0:

wall_x/fusions/backend.py

@@ -297,3 +297,137 @@ def rope_bwd(
     return backend.rope_bwd(
         grad_q_out, grad_k_out, q, k, cos, sin, grad_q, grad_k, mrope_section_doubled
     )
+
+
+def get_rope_index(
+    input_ids: torch.Tensor,
+    image_grid_thw: Optional[torch.Tensor],
+    video_grid_thw: Optional[torch.Tensor],
+    second_per_grid_ts: Optional[torch.Tensor],
+    attention_mask: Optional[torch.Tensor],
+    spatial_merge_size: int,
+    image_token_id: int,
+    video_token_id: int,
+    vision_start_token_id: int,
+    tokens_per_second: float,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Generate position indices for multimodal RoPE (Rotary Position Embedding).
+
+    This function computes 3D position indices for text, image, and video tokens
+    to enable proper spatial-temporal position encoding in multimodal transformers.
+
+    Args:
+        input_ids (torch.Tensor): Input token IDs of shape [batch_size, seq_len]
+        image_grid_thw (torch.Tensor, optional): Image grid specifications of shape [num_images, 3] (T, H, W)
+        video_grid_thw (torch.Tensor, optional): Video grid specifications of shape [num_videos, 3] (T, H, W)
+        second_per_grid_ts (torch.Tensor, optional): Temporal scaling per video grid of shape [num_videos]
+        attention_mask (torch.Tensor, optional): Attention mask of shape [batch_size, seq_len]
+        spatial_merge_size (int): Spatial dimension merge factor for patch grouping
+        image_token_id (int): Token ID representing image patches
+        video_token_id (int): Token ID representing video frames
+        vision_start_token_id (int): Token ID marking vision sequence start
+        tokens_per_second (float): Temporal scaling factor for video sequences
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+            - position_ids: 3D position indices of shape [3, batch_size, seq_len]
+            - mrope_deltas: Position deltas for multimodal RoPE of shape [batch_size, 1]
+
+    Note:
+        When both image_grid_thw and video_grid_thw are None, returns standard
+        text-only position indices based on attention_mask or sequence order.
+    """
+    return backend.rope_index(
+        input_ids,
+        image_grid_thw,
+        video_grid_thw,
+        second_per_grid_ts,
+        attention_mask,
+        spatial_merge_size,
+        image_token_id,
+        video_token_id,
+        vision_start_token_id,
+        tokens_per_second,
+    )
+
+
+def rot_pos_emb(
+    inv_freq: torch.Tensor,
+    grid_thw: torch.Tensor,
+    spatial_merge_size: int,
+) -> torch.Tensor:
+    """
+    Compute fused rotary position embeddings for multimodal grids.
+
+    This function efficiently computes rotary position embeddings for spatial-temporal
+    grids using a fused CUDA kernel, supporting both int32 and int64 grid specifications.
+
+    Args:
+        inv_freq (torch.Tensor): Inverse frequencies for RoPE of shape [dim/2]
+                                Must be float32 dtype on CUDA device
+        grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W)
+                                Supports int32 or int64 dtype on CUDA device
+        spatial_merge_size (int): Merge factor for spatial dimensions (must be positive)
+
+    Returns:
+        torch.Tensor: Computed rotary embeddings of shape [total_tokens, dim]
+                     where total_tokens is determined by grid layouts and spatial_merge_size
+
+    Example:
+        >>> inv_freq = torch.randn(64, device='cuda', dtype=torch.float32)  # 128-dim model
+        >>> grids = torch.tensor([[8, 14, 14], [16, 7, 7]], device='cuda', dtype=torch.int32)
+        >>> embeddings = rot_pos_emb(inv_freq, grids, spatial_merge_size=2)
+        >>> print(embeddings.shape)  # [computed_tokens, 128]
+
+    Note:
+        The function automatically dispatches to int32 or int64 implementations
+        based on the dtype of grid_thw. Output is always float32.
+    """
+    return backend.rot_pos_emb(inv_freq, grid_thw, spatial_merge_size)
+
+
+def get_window_index(
+    grid_thw: torch.Tensor,
+    spatial_merge_size: int,
+    vit_merger_window_size: int,
+    patch_size: int,
+    spatial_merge_unit: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Generate window attention indices for Vision Transformer architectures.
+
+    Computes window-based attention indices for hierarchical processing of vision
+    tokens, enabling efficient sliding window attention patterns in ViT models.
+
+    Args:
+        grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W)
+                                Must be int32 dtype on CUDA device
+        spatial_merge_size (int): Spatial dimension merge factor
+        vit_merger_window_size (int): Size of attention windows for ViT processing
+        patch_size (int): Size of vision patches in pixels
+        spatial_merge_unit (int): Unit size for spatial merging operations
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+            - window_indices: Flattened window indices of shape [total_elements]
+            - cu_window_seqlens: Cumulative window sequence lengths of shape [num_windows + 1]
+
+    Example:
+        >>> grids = torch.tensor([[1, 14, 14]], device='cuda', dtype=torch.int32)
+        >>> indices, seqlens = get_window_index(
+        ...     grids, spatial_merge_size=2, vit_merger_window_size=7,
+        ...     patch_size=16, spatial_merge_unit=4
+        ... )
+
+    Note:
+        Returns empty tensors if input grid is empty or no valid windows can be formed.
+        The cu_window_seqlens tensor enables efficient batched attention computation.
+    """
+    return backend.get_window_index(
+        grid_thw,
+        spatial_merge_size,
+        vit_merger_window_size,
+        patch_size,
+        spatial_merge_unit,
+    )

wall_x/fusions/ops.py

@@ -495,3 +495,248 @@ class MultimodalRoPE(torch.autograd.Function):
 
 def multimodal_rope(q, k, cos, sin, mrope_section):
     return MultimodalRoPE.apply(q, k, cos, sin, mrope_section)
+
+
+################################################################################################
+##
+## RoPE Index 3D
+##
+################################################################################################
+
+
+def get_rope_index(
+    input_ids: torch.Tensor,
+    spatial_merge_size: int,
+    image_token_id: int,
+    video_token_id: int,
+    vision_start_token_id: int,
+    tokens_per_second: float,
+    image_grid_thw: torch.Tensor = None,
+    video_grid_thw: torch.Tensor = None,
+    second_per_grid_ts: torch.Tensor = None,
+    attention_mask: torch.Tensor = None,
+):
+    """
+    Generate 3D RoPE position indices for multimodal transformer inputs.
+
+    Computes position indices for text, image, and video tokens to enable proper
+    spatial-temporal position encoding in multimodal transformers with RoPE.
+
+    Args:
+        input_ids (torch.Tensor): Input token sequence of shape [batch_size, seq_len]
+                                 Must be LongTensor on CUDA device
+        spatial_merge_size (int): Spatial merge size for patch grouping (must be positive)
+        image_token_id (int): Token ID representing image patches
+        video_token_id (int): Token ID representing video frames
+        vision_start_token_id (int): Token ID marking start of vision sequences
+        tokens_per_second (float): Temporal scaling factor for video sequences (must be positive)
+        image_grid_thw (torch.Tensor, optional): Image grid dimensions of shape [num_images, 3] (T, H, W)
+        video_grid_thw (torch.Tensor, optional): Video grid dimensions of shape [num_videos, 3] (T, H, W)
+        second_per_grid_ts (torch.Tensor, optional): Video time intervals of shape [num_videos]
+        attention_mask (torch.Tensor, optional): Attention mask of shape [batch_size, seq_len]
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+            - position_ids: 3D position indices of shape [3, batch_size, seq_len]
+            - mrope_position_deltas: mRoPE position deltas of shape [batch_size, 1]
+
+    Raises:
+        TypeError: If input_ids is not a torch.Tensor
+        ValueError: If input dimensions are incorrect or tensors not on CUDA
+    """
+    # Input validation
+    if not isinstance(input_ids, torch.Tensor):
+        raise TypeError("input_ids must be a torch.Tensor")
+
+    if input_ids.dim() != 2:
+        raise ValueError("input_ids must be 2D tensor (batch_size, seq_len)")
+
+    if not input_ids.is_cuda:
+        raise ValueError("input_ids must be on CUDA device")
+
+    # Parameter validation
+    if not isinstance(spatial_merge_size, int) or spatial_merge_size <= 0:
+        raise ValueError(
+            f"spatial_merge_size must be positive integer, got {spatial_merge_size}"
+        )
+
+    if not isinstance(tokens_per_second, (int, float)) or tokens_per_second <= 0:
+        raise ValueError(
+            f"tokens_per_second must be positive number, got {tokens_per_second}"
+        )
+
+    return backend.get_rope_index(
+        input_ids,
+        image_grid_thw,
+        video_grid_thw,
+        second_per_grid_ts,
+        attention_mask,
+        spatial_merge_size,
+        image_token_id,
+        video_token_id,
+        vision_start_token_id,
+        float(tokens_per_second),
+    )
+
+
+################################################################################################
+##
+## Fused Rotary Position Embedding
+##
+################################################################################################
+
+
+def rot_pos_emb(
+    inv_freq: torch.Tensor,
+    grid_thw: torch.Tensor,
+    spatial_merge_size: int,
+) -> torch.Tensor:
+    """
+    Compute fused rotary position embeddings using optimized CUDA kernel.
+
+    This function fuses all rotary position embedding computations into a single
+    CUDA kernel for improved performance with spatial-temporal grids.
+
+    Args:
+        inv_freq (torch.Tensor): Inverse frequencies tensor of shape [dim/2]
+                                Contains precomputed 1.0 / (theta ** (torch.arange(0, dim, 2) / dim))
+                                Must be float32 on CUDA device
+        grid_thw (torch.Tensor): Grid dimensions tensor of shape [num_grids, 3]
+                                Each row contains (T, H, W) for temporal, height, width dimensions
+                                Supports int32 or int64 on CUDA device
+        spatial_merge_size (int): Spatial merge size for token grouping (must be positive)
+
+    Returns:
+        torch.Tensor: Rotary position embeddings of shape [total_tokens, dim]
+                     where dim = 2 * len(inv_freq)
+                     First half contains h_pos frequencies, second half contains w_pos frequencies
+
+    Raises:
+        TypeError: If inputs are not torch.Tensor or spatial_merge_size not int
+        ValueError: If tensor dimensions incorrect, not on CUDA, or devices mismatch
+        RuntimeError: If CUDA kernel execution fails
+    """
+    # Type checking
+    if not isinstance(inv_freq, torch.Tensor):
+        raise TypeError(f"inv_freq must be a torch.Tensor, got {type(inv_freq)}")
+
+    if not isinstance(grid_thw, torch.Tensor):
+        raise TypeError(f"grid_thw must be a torch.Tensor, got {type(grid_thw)}")
+
+    # Dimension checking
+    if inv_freq.dim() != 1:
+        raise ValueError(
+            f"inv_freq must be 1-dimensional, got {inv_freq.dim()}D tensor"
+        )
+
+    if grid_thw.dim() != 2:
+        raise ValueError(
+            f"grid_thw must be 2-dimensional, got {grid_thw.dim()}D tensor"
+        )
+
+    if grid_thw.size(1) != 3:
+        raise ValueError(
+            f"grid_thw must have shape [num_grids, 3], got shape {list(grid_thw.shape)}"
+        )
+
+    # Device checking
+    if not inv_freq.is_cuda:
+        raise ValueError("inv_freq must be on CUDA device")
+
+    if not grid_thw.is_cuda:
+        raise ValueError("grid_thw must be on CUDA device")
+
+    # Ensure both tensors are on the same device
+    if inv_freq.device != grid_thw.device:
+        raise ValueError(
+            f"inv_freq and grid_thw must be on the same device, "
+            f"got {inv_freq.device} and {grid_thw.device}"
+        )
+
+    # Parameter validation
+    if not isinstance(spatial_merge_size, int):
+        raise TypeError(
+            f"spatial_merge_size must be an integer, got {type(spatial_merge_size)}"
+        )
+
+    if spatial_merge_size <= 0:
+        raise ValueError(
+            f"spatial_merge_size must be positive, got {spatial_merge_size}"
+        )
+
+    # Ensure inv_freq is float32 (the kernel expects float)
+    if inv_freq.dtype != torch.float32:
+        inv_freq = inv_freq.to(torch.float32)
+
+    # Call the CUDA backend
+    try:
+        return backend.rot_pos_emb(inv_freq, grid_thw, spatial_merge_size)
+    except RuntimeError as e:
+        raise RuntimeError(f"CUDA kernel execution failed: {str(e)}")
+
+
+################################################################################################
+##
+## Fused Window Index Generation
+##
+################################################################################################
+
+
+def get_window_index(
+    grid_thw: torch.Tensor,
+    window_size: int,
+    spatial_merge_size: int,
+    patch_size: int,
+    spatial_merge_unit: int = 1,
+):
+    """
+    Generate window attention indices for Vision Transformer architectures.
+
+    Computes window-based attention indices for hierarchical processing of vision
+    tokens, enabling efficient sliding window attention patterns in ViT models.
+
+    Args:
+        grid_thw (torch.Tensor): Grid specifications of shape [num_grids, 3] (T, H, W)
+                                Must be or will be converted to int32 on CUDA device
+        window_size (int): Window size for attention computation
+        spatial_merge_size (int): Spatial merge size for patch grouping
+        patch_size (int): Size of vision patches in pixels
+        spatial_merge_unit (int, optional): Spatial merging unit size. Defaults to 1.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+            - window_index: Window indices tensor of shape [total_elements]
+            - cu_window_seqlens: Cumulative window sequence lengths of shape [num_windows + 1]
+
+    Raises:
+        AssertionError: If grid_thw dimensions are incorrect
+
+    Note:
+        Returns empty tensors if input grid is empty or no valid windows can be formed.
+        The function automatically converts input to CUDA int32 if needed.
+    """
+    # Input validation
+    assert (
+        grid_thw.dim() == 2 and grid_thw.size(1) == 3
+    ), f"grid_thw must have shape (num_grids, 3), got {grid_thw.shape}"
+
+    # Ensure input is on CUDA and int32 type
+    if not grid_thw.is_cuda:
+        grid_thw = grid_thw.cuda()
+
+    if grid_thw.dtype != torch.int32:
+        grid_thw = grid_thw.to(torch.int32)
+
+    # Calculate vit_merger_window_size
+    vit_merger_window_size = window_size // spatial_merge_size // patch_size
+
+    # Call CUDA backend
+    window_index, cu_window_seqlens = backend.get_window_index(
+        grid_thw,
+        spatial_merge_size,
+        vit_merger_window_size,
+        patch_size,
+        spatial_merge_unit,
+    )
+
+    return window_index, cu_window_seqlens

wall_x/model/qwen2_5_based/modeling_qwen2_5_vl.py

@@ -560,15 +560,17 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel):
             `torch.Tensor`: hidden_states.
         """
         hidden_states = self.patch_embed(hidden_states)
-        rotary_pos_emb = self.rot_pos_emb(grid_thw)
+        rotary_pos_emb = ops.rot_pos_emb(
-        window_index, cu_window_seqlens = self.get_window_index(grid_thw)
+            self.rotary_pos_emb.inv_freq, grid_thw, self.spatial_merge_size
-        window_index = window_index.to(hidden_states.device)
+        )
-        cu_window_seqlens = torch.tensor(
+
-            cu_window_seqlens,
+        window_index, cu_window_seqlens = ops.get_window_index(
-            device=hidden_states.device,
+            grid_thw=grid_thw,
-            dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+            window_size=self.window_size,
+            spatial_merge_size=self.spatial_merge_size,
+            patch_size=self.patch_size,
+            spatial_merge_unit=self.spatial_merge_unit,
         )
-        cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
 
         seq_len, _ = hidden_states.size()
         hidden_states = hidden_states.reshape(

wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py

@@ -1244,12 +1244,17 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
                 or self.rope_deltas is None
                 or (past_key_values is None or past_key_values.get_seq_length() == 0)
             ):
-                position_ids, rope_deltas = self.get_rope_index(
+                position_ids, rope_deltas = ops.get_rope_index(
-                    input_ids,
+                    input_ids=input_ids,
-                    image_grid_thw,
+                    image_grid_thw=image_grid_thw,
-                    video_grid_thw,
+                    video_grid_thw=video_grid_thw,
-                    second_per_grid_ts,
+                    second_per_grid_ts=second_per_grid_ts,
-                    attention_mask,
+                    attention_mask=attention_mask,
+                    spatial_merge_size=self.config.vision_config.spatial_merge_size,
+                    image_token_id=self.config.image_token_id,
+                    video_token_id=self.config.video_token_id,
+                    vision_start_token_id=self.config.vision_start_token_id,
+                    tokens_per_second=self.config.vision_config.tokens_per_second,
                 )
                 self.rope_deltas = rope_deltas
             # Use previously calculated rope deltas to get correct position IDs
@@ -1720,12 +1725,17 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
                 or self.rope_deltas is None
                 or (past_key_values is None or past_key_values.get_seq_length() == 0)
             ):
-                position_ids, rope_deltas = self.get_rope_index(
+                position_ids, rope_deltas = ops.get_rope_index(
-                    input_ids,
+                    input_ids=input_ids,
-                    image_grid_thw,
+                    image_grid_thw=image_grid_thw,
-                    video_grid_thw,
+                    video_grid_thw=video_grid_thw,
-                    second_per_grid_ts,
+                    second_per_grid_ts=second_per_grid_ts,
-                    attention_mask,
+                    attention_mask=attention_mask,
+                    spatial_merge_size=self.config.vision_config.spatial_merge_size,
+                    image_token_id=self.config.image_token_id,
+                    video_token_id=self.config.video_token_id,
+                    vision_start_token_id=self.config.vision_start_token_id,
+                    tokens_per_second=self.config.vision_config.tokens_per_second,
                 )
                 self.rope_deltas = rope_deltas
             # Use previously calculated rope deltas to get correct position IDs
@@ -1882,6 +1892,17 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
             )
             dof_mask = dof_mask.to(inputs_embeds.device).to(inputs_embeds.dtype)
 
+            # Calculate token distribution across MoE expert groups
+            group_size = torch.zeros(
+                self.config.num_experts, dtype=torch.long, device="cpu"
+            )
+            for i in range(self.config.num_experts):
+                group_size[i] = (moe_token_types == i).sum()
+
+            # Calculate start and end indices for each expert group
+            start_indices = torch.cumsum(group_size, dim=0) - group_size
+            end_indices = torch.cumsum(group_size, dim=0)
+
             def step(timestep, noisy_action):
                 """
                 Single denoising step for diffusion process.
@@ -1915,6 +1936,8 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
                     past_key_values=past_key_values,
                     inputs_embeds=temp_inputs_embeds,
                     moe_token_types=moe_token_types,
+                    start_indices=start_indices,
+                    end_indices=end_indices,
                     use_cache=True,
                     output_attentions=False,
                     output_hidden_states=False,

wall_x/trainer/qwen_vl_act_trainer.py

@@ -5,11 +5,13 @@ import torch
 import random
 import numpy as np
 import torch.nn as nn
+import torch.distributed as dist
 
 from tqdm import tqdm
 from functools import wraps
 from datetime import datetime
 from torch.optim import AdamW
+from torch.distributed.tensor import distribute_tensor
 from accelerate import Accelerator
 from safetensors.torch import load_file
 from transformers.optimization import get_cosine_with_min_lr_schedule_with_warmup
@@ -769,7 +771,9 @@ class QwenVlAct_Trainer:
         """
         checkpoint_path = self.config["resume"]["ckpt"]
 
-        if self.config.get("resume", {}).get("load_ckpt_only", False):
+        if self.config.get("FSDP2", False):
+            self._load_fsdp_state_dict_with_distribute_tensor()
+        elif self.config.get("resume", {}).get("load_ckpt_only", False):
             # Load only model weights
             ckpt_path = self.config["resume"]["ckpt"] + "/model.safetensors"
             state_dict = load_file(ckpt_path, device="cpu")
@@ -788,6 +792,48 @@ class QwenVlAct_Trainer:
 
         self.print_rank0(f"Resumed from checkpoint: {checkpoint_path}")
 
+    def _load_fsdp_state_dict_with_distribute_tensor(self):
+
+        rank = dist.get_rank() if dist.is_initialized() else 0
+
+        full_sd = load_file(
+            self.config["resume"]["ckpt"] + "/model.safetensors", device="cpu"
+        )
+        meta_sharded_sd = self.model.state_dict()
+        sharded_sd = {}
+
+        def find_matching_key(target_key, available_keys):
+            if target_key in available_keys:
+                return target_key
+            prefixed_key = f"_orig_mod.{target_key}"
+            if prefixed_key in available_keys:
+                return prefixed_key
+            if target_key.startswith("_orig_mod."):
+                unprefixed_key = target_key[len("_orig_mod.") :]
+                if unprefixed_key in available_keys:
+                    return unprefixed_key
+            return None
+
+        for param_name, full_tensor in full_sd.items():
+            matching_key = find_matching_key(param_name, meta_sharded_sd.keys())
+            if matching_key is None:
+                if rank == 0:
+                    print(
+                        f"[Rank {rank}] Warning: Parameter not found:",
+                        param_name,
+                        flush=True,
+                    )
+                continue
+            sharded_meta_param = meta_sharded_sd[matching_key]
+            sharded_tensor = distribute_tensor(
+                full_tensor,
+                sharded_meta_param.device_mesh,
+                sharded_meta_param.placements,
+            )
+            sharded_sd[matching_key] = nn.Parameter(sharded_tensor)
+
+        self.model.load_state_dict(sharded_sd, assign=True, strict=False)
+
     def log_l1_details(self, all_label, all_pred, all_task, all_dof_mask):
         """
         Log detailed L1 loss metrics by degrees of freedom.

workspace/README.md

@@ -1,15 +1,28 @@
-# Training Configuration Guide
+# Training Guide
 
-This document explains the key configuration parameters that can be modified for Wall-X training.
+This document explains the key configuration parameters and memory requirements for Wall-X training.
 
 ## Quick Start Checklist
-1. **Update run.sh**: Set `code_dir` and `config_path` to your actual paths
+
-2. **Configure GPUs**: Set `CUDA_VISIBLE_DEVICES` for your available GPUs
+### 🚀 **Step 1: Download Pre-trained Model**
-3. **Update config paths**: Replace all `/path/to/` placeholders in `config_qact.yml` with actual paths
+Choose one of the available models:
-4. **Configure robot**: Set `dof_config` and `agent_pos_config` for your robot
+- **WALL-OSS-FLOW**: https://huggingface.co/x-square-robot/wall-oss-flow
-5. **Set dataset**: Choose appropriate `repo_id` for your dataset
+- **WALL-OSS-FAST**: https://huggingface.co/x-square-robot/wall-oss-fast
-6. **Adjust batch size**: Set `batch_size_per_gpu` based on GPU memory
+
-7. **Run training**: Execute `bash ./workspace/lerobot_example/run.sh`
+### ⚙️ **Step 2: Configure Environment**
+- Update `run.sh`: Set `code_dir` and `config_path` to your actual paths
+- Set `CUDA_VISIBLE_DEVICES` for your available GPUs
+
+### 📝 **Step 3: Update Configuration Files**
+- Replace all `/path/to/` placeholders in `config_qact.yml` with actual paths
+- Configure robot settings: `dof_config` and `agent_pos_config`
+- Set dataset: Choose appropriate `repo_id`
+- Adjust `batch_size_per_gpu` based on your GPU memory
+
+### ▶️ **Step 4: Start Training**
+```bash
+bash ./workspace/lerobot_example/run.sh
+```
 
 ## Enable FAST tokenizer
 To fine-tune using the FAST tokenizer, please download the repository and update the `action_tokenizer_path`. Make sure to set `use_fast_tokenizer` to `true`:
@@ -38,6 +51,15 @@ action_tokenizer_path: "/path/to/fast/"             # Must set if use_fast_token
 - `num_training_steps`: Total training steps
 - `num_epoch`: Number of training epochs
 
+### Training Optimization Settings
+- `FSDP2`: Enable FSDP2 for distributed training (default: True) - **Recommended for multi-GPU**
+- `torch_compile`: Enable PyTorch compilation optimization (default: False)
+
+**⚠️ Important Note on torch_compile:**
+- **Benefits**: Enabling `torch_compile` can significantly improve training efficiency
+- **Requirements**: Requires that the data input shape is always consistent throughout training
+- **Caution**: If you don't have sufficient understanding of torch compile, please **DO NOT** enable it as it may cause unexpected issues with dynamic input shapes
+
 ## Robot Configuration (Modify for Your Robot)
 
 ### DOF Configuration
@@ -67,6 +89,20 @@ Keep `agent_pos_config` consistent with `dof_config`.
 - `resume.ckpt`: Path to checkpoint for resuming training
 - `resume.load_ckpt_only`: Only load model weights, not optimizer state
 
-## Performance Settings (Optional)
+## Memory Usage
-- `profile`: Enable PyTorch profiling (true/false)
+
-- `padding_side`: Token padding side (left/right)
+Below are the memory consumption benchmarks for different training configurations using the `lerobot/aloha_mobile_cabinet` dataset:
+
+| Dataset | Batch Size | FSDP2 | Torch Compile | Num GPUs | Max Allocated Memory |
+|---------|------------|--------|---------------|----------|---------------------|
+| lerobot/aloha_mobile_cabinet | 1 | ❌ | ❌ | 1 | 40.11G |
+| lerobot/aloha_mobile_cabinet | 1 | ❌ | ❌ | 8 | 48.02G |
+| lerobot/aloha_mobile_cabinet | 1 | ✅ | ❌ | 2 | 43.70G |
+| lerobot/aloha_mobile_cabinet | 1 | ✅ | ❌ | 8 | 24.96G |
+| lerobot/aloha_mobile_cabinet | 1 | ✅ | ✅ | 8 | 24.21G |
+
+
+**Hardware Recommendations:**
+
+- For single GPU training: Ensure at least 48GB VRAM (e.g., RTX 6000 Ada, A6000)
+- For multi-GPU training: Enable FSDP2 for optimal memory distribution

workspace/lerobot_example/config_qact.yml

@@ -28,6 +28,10 @@ batch_size_per_gpu: 8
 padding_side: left
 epoch_save_interval: 10
 
+# Training optimization settings
+FSDP2: True
+torch_compile: False
+
 # Robot configuration - Define degrees of freedom for each component
 dof_config:
   follow_left_ee_cartesian_pos: 3      # Left end-effector Cartesian position

workspace/lerobot_example/run.sh

@@ -1,5 +1,4 @@
 #!/bin/bash
-# export CUDA_VISIBLE_DEVICES=4,5,6,7
 export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
 NUM_GPUS=$(echo $CUDA_VISIBLE_DEVICES | tr ',' '\n' | wc -l)