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
This commit is contained in:
Starrick Liu
+332
@@ -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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user