Hybrid model upgrade - #3480
Conversation
|
note: the custom kernel is passing some testcases in isolation, but is not yet passing the mlc chat check (having some weird crash faults). Just wanted to open early for visibility, will be updating this as it progresses. |
![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 implements a chunked prefill path for GatedDeltaNet in the Qwen 3.5 model, introducing a Relax/TIR hybrid function and updating the model architecture to dispatch between chunked and recurrent kernels. It also enables single-batch optimizations for hybrid KV states in the C++ serving layer. Feedback highlights critical issues in the new TIR kernels, specifically a race condition in inter_chunk_recurrent due to missing thread bindings and a performance bottleneck in associative_scan caused by the lack of thread-level parallelism in its
| for b in T.thread_binding(batch_size, thread="blockIdx.x"): | ||
| for h in T.thread_binding(T.int64(num_value_heads), thread="blockIdx.y"): | ||
| for chunk_idx in T.serial(num_chunks): | ||
| for c1 in T.serial(T.int64(chunk_size)): | ||
| g_i = g_cumsum_buf[b, h, chunk_idx, c1] | ||
| for c2 in T.serial(T.int64(chunk_size)): | ||
| attn_buf[b, h, chunk_idx, c1, c2] = T.float32(0.0) | ||
| for kd in T.serial(T.int64(key_head_dim)): | ||
| attn_buf[b, h, chunk_idx, c1, c2] += ( | ||
| query_buf[b, h, chunk_idx, c1, kd] | ||
| * key_buf[b, h, chunk_idx, c2, kd] | ||
| ) | ||
| attn_buf[b, h, chunk_idx, c1, c2] = T.if_then_else( | ||
| c2 > c1, | ||
| T.float32(0.0), | ||
| attn_buf[b, h, chunk_idx, c1, c2] | ||
| * T.exp(g_i - g_cumsum_buf[b, h, chunk_idx, c2]), | ||
| ) | ||
|
|
||
| for c in T.serial(T.int64(chunk_size)): | ||
| g_curr = g_cumsum_buf[b, h, chunk_idx, c] | ||
| exp_curr = T.exp(g_curr) | ||
| for vd in T.serial(T.int64(value_head_dim)): | ||
| v_new_buf[b, h, chunk_idx, c, vd] = value_intra_buf[ | ||
| b, h, chunk_idx, c, vd | ||
| ] | ||
| for kd in T.serial(T.int64(key_head_dim)): | ||
| v_new_buf[b, h, chunk_idx, c, vd] -= ( | ||
| k_cumdecay_buf[b, h, chunk_idx, c, kd] | ||
| * recurrent_state_out_buf[b, h, kd, vd] | ||
| ) | ||
|
|
||
| for vd in T.serial(T.int64(value_head_dim)): | ||
| attn_inter_buf[b, h, chunk_idx, c, vd] = T.float32(0.0) | ||
| for kd in T.serial(T.int64(key_head_dim)): | ||
| attn_inter_buf[b, h, chunk_idx, c, vd] += ( | ||
| query_buf[b, h, chunk_idx, c, kd] | ||
| * exp_curr | ||
| * recurrent_state_out_buf[b, h, kd, vd] | ||
| ) | ||
|
|
||
| for vd in T.serial(T.int64(value_head_dim)): | ||
| core_attn_out_buf[b, h, chunk_idx, c, vd] = attn_inter_buf[ | ||
| b, h, chunk_idx, c, vd | ||
| ] | ||
| for r in T.serial(T.int64(chunk_size)): | ||
| core_attn_out_buf[b, h, chunk_idx, c, vd] += ( | ||
| attn_buf[b, h, chunk_idx, c, r] | ||
| * v_new_buf[b, h, chunk_idx, r, vd] | ||
| ) | ||
|
|
||
| for kd in T.serial(T.int64(key_head_dim)): | ||
| for vd in T.serial(T.int64(value_head_dim)): | ||
| g_last = g_cumsum_buf[b, h, chunk_idx, T.int64(chunk_size - 1)] | ||
| exp_last = T.exp(g_last) | ||
| recurrent_state_out_buf[b, h, kd, vd] = ( | ||
| recurrent_state_out_buf[b, h, kd, vd] * exp_last | ||
| ) | ||
| for c in T.serial(T.int64(chunk_size)): | ||
| recurrent_state_out_buf[b, h, kd, vd] += ( | ||
| key_buf[b, h, chunk_idx, c, kd] | ||
| * v_new_buf[b, h, chunk_idx, c, vd] | ||
| * T.exp(g_last - g_cumsum_buf[b, h, chunk_idx, c]) | ||
| ) |
There was a problem hiding this comment.
The inter_chunk_recurrent PrimFunc contains a critical race condition. While threadIdx.x is bound to value_head_dim in the initialization block (lines 624-628), it is not bound in the main computation block (lines 630-693). Consequently, every thread in the GPU block will redundantly execute the serial loops and perform unsynchronized writes to the same shared buffers (attn_buf, v_new_buf, recurrent_state_out_buf, etc.). This will lead to non-deterministic results and severe performance overhead. You should bind the loops iterating over vd to threadIdx.x and ensure that computations not depending on vd (like the attn_buf calculation) are executed by only one thread or properly parallelized.
| for i in T.serial(1, chunk_size): | ||
| g_i = g_cumsum_buf[b, h, c, i] | ||
| for j in T.serial(0, i): | ||
| decayed_lower_attn_ij = ( | ||
| -attn_mm_buf[b, h, c, i, j] | ||
| * T.exp(g_i - g_cumsum_buf[b, h, c, j]) | ||
| ) | ||
| attn_out_buf[b, h, c, i, j] = decayed_lower_attn_ij | ||
| for k_idx in T.serial(0, i): | ||
| decayed_lower_attn_ik = ( | ||
| -attn_mm_buf[b, h, c, i, k_idx] | ||
| * T.exp(g_i - g_cumsum_buf[b, h, c, k_idx]) | ||
| ) | ||
| attn_out_buf[b, h, c, i, j] += ( | ||
| decayed_lower_attn_ik * attn_out_buf[b, h, c, k_idx, j] | ||
| ) |
There was a problem hiding this comment.
The associative_scan PrimFunc implements an chunk_size) using nested serial loops without any thread-level parallelism. With chunk_size=64, this results in a large number of serial operations per chunk. Since this is a GPU kernel, it should be parallelized over i and j using threadIdx to utilize the GPU's compute resources effectively. As currently implemented, this kernel will likely become a bottleneck during prefill for large sequences.
be7ef72 to
c2464ec
Compare
this should help boost slow prefill performance by using smem and better parallelization
026926d to
7e29e72
Compare
|
now getting correct / sensible output from chat interface, still looking for more performance |
1 participant
Right now, we leverage the recurrent decode kernel for both prefill and decode in the qwen3.5 model. Huggingface / other implementations leverage a specialized prefill kernel that should be higher performance. This PR ports that logic via a custom implementation of chunked gated deltanet.
I also adjusted the engine / model to support batch size 1 methods, these were previously disabled for hybrid models universally.