Refactor CUDA interface for improved usability #80
Conversation
Replaces direct CUDA module import with higher-level interface function that provides cleaner API with named parameters. Simplifies function call by removing manual tensor preparation and using more intuitive parameter names like dropout_p and softmax_scale. Adds runtime check to ensure CUDA implementation availability before executing tests.
LoserCheems
requested review from
Evanwu1125,
SNHuan,
Copilot and
wubingheng111
and removed request for
Copilot
There was a problem hiding this comment.
Pull Request Overview
Refactor the CUDA interface to use a higher-level function with named parameters instead of direct module imports. The changes simplify tensor preparation by removing manual contiguous() calls and improve error handling with runtime checks for CUDA availability.
- Replace direct
flash_dmattn_cudamodule import withflash_dmattn_funcfrom interface module - Simplify CUDA function call using named parameters instead of positional arguments
- Add runtime checks for CUDA availability before executing tests
| query_states = query_states.transpose(1, 2) # [batch, query_len, num_heads, head_dim] | ||
| key_states = key_states.transpose(1, 2) # [batch, key_len, num_kv_heads, head_dim] | ||
| value_states = value_states.transpose(1, 2) # [batch, key_len, num_kv_heads, head_dim] |
There was a problem hiding this comment.
Removing .contiguous() calls may cause performance issues if the tensors are not contiguous in memory. CUDA kernels typically require contiguous tensors for optimal performance. Consider adding .contiguous() back or verify that the new interface handles non-contiguous tensors efficiently.
| query_states = query_states.transpose(1, 2) # [batch, query_len, num_heads, head_dim] | |
| key_states = key_states.transpose(1, 2) # [batch, key_len, num_kv_heads, head_dim] | |
| value_states = value_states.transpose(1, 2) # [batch, key_len, num_kv_heads, head_dim] | |
| query_states = query_states.transpose(1, 2).contiguous() # [batch, query_len, num_heads, head_dim] | |
| key_states = key_states.transpose(1, 2).contiguous() # [batch, key_len, num_kv_heads, head_dim] | |
| value_states = value_states.transpose(1, 2).contiguous() # [batch, key_len, num_kv_heads, head_dim] |
| attn_mask = attn_mask.contiguous() # [batch, num_kv_heads, query_len, key_len] | ||
| query_states = query_states.transpose(1, 2) # [batch, query_len, num_heads, head_dim] | ||
| key_states = key_states.transpose(1, 2) # [batch, key_len, num_kv_heads, head_dim] | ||
| value_states = value_states.transpose(1, 2) # [batch, key_len, num_kv_heads, head_dim] |
There was a problem hiding this comment.
The zoh_states calculation and expansion logic has been removed, but the calculate_zoh_states function is still called earlier in the function. This may cause the attention mechanism to not work correctly as zoh_states appear to be required for the dynamic mask attention computation.
Simplify the CUDA interface by replacing direct module imports with a higher-level function that uses named parameters. Enhance usability by removing manual tensor preparation and adding runtime checks for CUDA availability before executing tests.