causal_conv1d_update_kernel_bdt_fwd optimized version - #388
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the causal_conv1d_update_kernel_bdt_fwd kernel in convolution.py by changing how weights and input blocks are loaded, updating the state-saving logic, and manually calculating the SiLU activation. The review feedback highlights several critical areas for improvement: adding a boundary mask to weight loading to prevent out-of-bounds memory access, restricting redundant state-copying operations to ti == 0 to optimize memory traffic, and reverting to the more idiomatic and optimized tl.sigmoid function.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| w_off_d = di * D_CHK_SIZE + tl.arange(0, D_CHK_SIZE)[:, None] | ||
| w_off_w = tl.arange(0, width)[None, :] | ||
| w = tl.load(weight_ptr + w_off_d * width + w_off_w).to(tl.float32) |
There was a problem hiding this comment.
The weight loading is performed using flat pointer arithmetic without any boundary checks. If dim is not a multiple of D_CHK_SIZE, this will result in out-of-bounds memory access. A mask should be applied to ensure safe loading.
| w_off_d = di * D_CHK_SIZE + tl.arange(0, D_CHK_SIZE)[:, None] | |
| w_off_w = tl.arange(0, width)[None, :] | |
| w = tl.load(weight_ptr + w_off_d * width + w_off_w).to(tl.float32) | |
| w_off_d = di * D_CHK_SIZE + tl.arange(0, D_CHK_SIZE)[:, None] | |
| w_off_w = tl.arange(0, width)[None, :] | |
| w_mask = w_off_d < dim | |
| w = tl.load(weight_ptr + w_off_d * width + w_off_w, mask=w_mask, other=0.0).to(tl.float32) |
| nst_mask = (nst_off_y0 < dim) & (nst_off_y1 < state_len) | ||
| block_ptr_st = bi * dim * state_len + nst_off_y0 * state_len + nst_off_y1 | ||
|
|
||
| if new_state_start_off < 0: |
There was a problem hiding this comment.
When new_state_start_off < 0, the old state elements are copied from conv_state_ptr to conv_state_update_ptr. Since cs_mask and cs_val do not depend on ti, every single ti block will redundantly load and store the exact same state elements. Restricting this operation to ti == 0 avoids redundant global memory traffic and potential race conditions.
| if new_state_start_off < 0: | |
| if new_state_start_off < 0 and ti == 0: |
| if SILU_ACTIVATION: | ||
| out_block = out_block * tl.sigmoid(out_block) | ||
| tl.store( | ||
| tl.make_block_ptr( | ||
| out_ptr, | ||
| shape=(batch, dim, out_len), | ||
| strides=(dim * out_len, out_len, 1), | ||
| offsets=(bi, di * D_CHK_SIZE, ti * T_CHK_SIZE), | ||
| block_shape=(1, D_CHK_SIZE, T_CHK_SIZE), | ||
| order=(2, 1, 0), | ||
| ), | ||
| out_block[None, :, :], | ||
| boundary_check=(0, 1, 2), | ||
| ) | ||
| out_block = out_block / (1.0 + tl.exp(-out_block)) |
There was a problem hiding this comment.
Using tl.sigmoid is more idiomatic, cleaner, and potentially more numerically stable and optimized than manually computing the sigmoid denominator via tl.exp.
| if SILU_ACTIVATION: | |
| out_block = out_block * tl.sigmoid(out_block) | |
| tl.store( | |
| tl.make_block_ptr( | |
| out_ptr, | |
| shape=(batch, dim, out_len), | |
| strides=(dim * out_len, out_len, 1), | |
| offsets=(bi, di * D_CHK_SIZE, ti * T_CHK_SIZE), | |
| block_shape=(1, D_CHK_SIZE, T_CHK_SIZE), | |
| order=(2, 1, 0), | |
| ), | |
| out_block[None, :, :], | |
| boundary_check=(0, 1, 2), | |
| ) | |
| out_block = out_block / (1.0 + tl.exp(-out_block)) | |
| if SILU_ACTIVATION: | |
| out_block = out_block * tl.sigmoid(out_block) |
1 participant
causal_conv1d_update_kernel_bdt_fwd optimized version