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

feat: Major optimization and robustness improvements (#31)

Browse Source 
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
This commit is contained in:
Starrick Liu
2025-09-17 23:09:20 +08:00
committed by GitHub
parent 86f70a3b08
commit 421db17d53
24 changed files with 1862 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/rope.cu
+44 -24
View File
@@ -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)
{