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
wall_x/fusions/backend.py
+134
View File
@@ -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
+245
View File
@@ -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