diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py index e7430f4f9..28fc166f4 100644 --- a/python/sglang/srt/layers/attention/fla/chunk.py +++ b/python/sglang/srt/layers/attention/fla/chunk.py @@ -141,11 +141,11 @@ def chunk_gated_delta_rule( Scale factor for the RetNet attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. + Initial state of shape `[N, H, V, K]` for `N` input sequences. For equal-length input sequences, `N` equals the batch size `B`. Default: `None`. output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + Whether to output the final state of shape `[N, H, V, K]`. Default: `False`. cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. @@ -157,7 +157,7 @@ def chunk_gated_delta_rule( o (torch.Tensor): Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`. final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`. Examples:: >>> import torch diff --git a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py index 38a7c8f29..7e5fd53fc 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py +++ b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py @@ -70,24 +70,24 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( NT = tl.cdiv(T, BT) boh = i_n * NT - # [BK, BV] - b_h1 = tl.zeros([64, BV], dtype=tl.float32) + # [BV, BK] + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) if K > 64: - b_h2 = tl.zeros([64, BV], dtype=tl.float32) + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) if K > 128: - b_h3 = tl.zeros([64, BV], dtype=tl.float32) + b_h3 = tl.zeros([BV, 64], dtype=tl.float32) if K > 192: - b_h4 = tl.zeros([64, BV], dtype=tl.float32) + b_h4 = tl.zeros([BV, 64], dtype=tl.float32) # calculate offset - h += ((boh * H + i_h) * K * V).to(tl.int64) + h += ((boh * H + i_h) * V * K).to(tl.int64) v += ((bos * H + i_h) * V).to(tl.int64) k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64) w += ((bos * H + i_h) * K).to(tl.int64) if SAVE_NEW_VALUE: v_new += ((bos * H + i_h) * V).to(tl.int64) stride_v = H * V - stride_h = H * K * V + stride_h = H * V * K stride_k = Hg * K stride_w = H * K @@ -95,49 +95,49 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( h0 = initial_state + index * stride_h ht = initial_state + index * stride_h if USE_INITIAL_STATE: - h0 = h0 + i_h * K * V + h0 = h0 + i_h * V * K if INPLACE_UPDATE: - ht = ht + i_h * K * V + ht = ht + i_h * V * K # load initial state if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) if K > 64: p_h0_2 = tl.make_block_ptr( - h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) ) b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) if K > 128: p_h0_3 = tl.make_block_ptr( - h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) ) b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) if K > 192: p_h0_4 = tl.make_block_ptr( - h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) ) b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) # main recurrence for i_t in range(NT): p_h1 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0) + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0) ) tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) if K > 64: p_h2 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) ) tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) if K > 128: p_h3 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) ) tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) if K > 192: p_h4 = tl.make_block_ptr( - h + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) ) tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) @@ -145,25 +145,25 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) ) b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) if K > 64: p_w = tl.make_block_ptr( w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) ) b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) + b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) if K > 128: p_w = tl.make_block_ptr( w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) ) b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) + b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) if K > 192: p_w = tl.make_block_ptr( w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) ) b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) + b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) p_v = tl.make_block_ptr( v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) ) @@ -199,7 +199,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k1 < K), other=0.0, ) - b_h1 *= exp(b_gk_last1)[:, None] + b_h1 *= exp(b_gk_last1)[None, :] if K > 64: o_k2 = 64 + o_k1 b_gk_last2 = tl.load( @@ -207,7 +207,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k2 < K), other=0.0, ) - b_h2 *= exp(b_gk_last2)[:, None] + b_h2 *= exp(b_gk_last2)[None, :] if K > 128: o_k3 = 128 + o_k1 b_gk_last3 = tl.load( @@ -215,7 +215,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k3 < K), other=0.0, ) - b_h3 *= exp(b_gk_last3)[:, None] + b_h3 *= exp(b_gk_last3)[None, :] if K > 192: o_k4 = 192 + o_k1 b_gk_last4 = tl.load( @@ -223,50 +223,50 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( mask=(o_k4 < K), other=0.0, ) - b_h4 *= exp(b_gk_last4)[:, None] + b_h4 *= exp(b_gk_last4)[None, :] b_v = b_v.to(k.dtype.element_ty) p_k = tl.make_block_ptr( k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) ) b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h1 += tl.dot(b_k, b_v) + b_h1 += tl.trans(tl.dot(b_k, b_v)) if K > 64: p_k = tl.make_block_ptr( k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) ) b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h2 += tl.dot(b_k, b_v) + b_h2 += tl.trans(tl.dot(b_k, b_v)) if K > 128: p_k = tl.make_block_ptr( k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) ) b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h3 += tl.dot(b_k, b_v) + b_h3 += tl.trans(tl.dot(b_k, b_v)) if K > 192: p_k = tl.make_block_ptr( k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) ) b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h4 += tl.dot(b_k, b_v) + b_h4 += tl.trans(tl.dot(b_k, b_v)) # epilogue if INPLACE_UPDATE: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) + p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) if K > 64: p_ht = tl.make_block_ptr( - ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0) + ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) ) tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) if K > 128: p_ht = tl.make_block_ptr( - ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0) + ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) ) tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) if K > 192: p_ht = tl.make_block_ptr( - ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0) + ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) ) tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) @@ -302,7 +302,7 @@ def chunk_gated_delta_rule_fwd_h( ) assert K <= 256, "current kernel does not support head dimension larger than 256." - h = k.new_empty(B, NT, H, K, V) + h = k.new_empty(B, NT, H, V, K) v_new = torch.empty_like(u) if save_new_value else None diff --git a/python/sglang/srt/layers/attention/fla/chunk_o.py b/python/sglang/srt/layers/attention/fla/chunk_o.py index bb89421eb..bac5e93a5 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_o.py +++ b/python/sglang/srt/layers/attention/fla/chunk_o.py @@ -71,7 +71,7 @@ def chunk_fwd_kernel_o( k += (bos * Hg + i_h // (H // Hg)) * K v += (bos * H + i_h) * V o += (bos * H + i_h) * V - h += (i_tg * H + i_h).to(tl.int64) * K * V + h += (i_tg * H + i_h).to(tl.int64) * V * K b_o = tl.zeros([BT, BV], dtype=tl.float32) b_A = tl.zeros([BT, BT], dtype=tl.float32) @@ -84,17 +84,17 @@ def chunk_fwd_kernel_o( k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) ) p_h = tl.make_block_ptr( - h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) + h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0) ) # [BT, BK] b_q = tl.load(p_q, boundary_check=(0, 1)) # [BK, BT] b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BK, BV] + # [BV, BK] b_h = tl.load(p_h, boundary_check=(0, 1)) # [BT, BK] @ [BK, BV] -> [BT, BV] - b_o += tl.dot(b_q, b_h) + b_o += tl.dot(b_q, tl.trans(b_h)) # [BT, BK] @ [BK, BT] -> [BT, BT] b_A += tl.dot(b_q, b_k) diff --git a/python/sglang/srt/layers/attention/fla/fused_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_recurrent.py index 8f4130c10..2fe5e4244 100644 --- a/python/sglang/srt/layers/attention/fla/fused_recurrent.py +++ b/python/sglang/srt/layers/attention/fla/fused_recurrent.py @@ -64,17 +64,17 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( if not IS_KDA: p_g = g + bos * HV + i_hv else: - p_gk = g + (bos * HV + i_hv) * K + o_k + p_gk = g + (bos * H + i_h) * K + o_k p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v mask_k = o_k < K mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] + mask_h = mask_v[:, None] & mask_k[None, :] - b_h = tl.zeros([BK, BV], dtype=tl.float32) + b_h = tl.zeros([BV, BK], dtype=tl.float32) if USE_INITIAL_STATE: - p_h0 = h0 + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + p_h0 = h0 + i_nh * V * K + o_v[:, None] * K + o_k[None, :] b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) for _ in range(0, T): @@ -86,24 +86,24 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6)) b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6)) b_q = b_q * scale - # [BK, BV] + # [BV, BK] if not IS_KDA: b_g = tl.load(p_g).to(tl.float32) b_h *= exp(b_g) else: - b_gk = tl.load(p_gk).to(tl.float32) - b_h *= exp(b_gk[:, None]) + b_gk = tl.load(p_gk, mask=mask_k, other=0).to(tl.float32) + b_h *= exp(b_gk[None, :]) # [BV] - b_v -= tl.sum(b_h * b_k[:, None], 0) + b_v -= tl.sum(b_h * b_k[None, :], 1) if IS_BETA_HEADWISE: b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) else: b_beta = tl.load(p_beta).to(tl.float32) b_v *= b_beta - # [BK, BV] - b_h += b_k[:, None] * b_v[None, :] + # [BV, BK] + b_h += b_v[:, None] * b_k[None, :] # [BV] - b_o = tl.sum(b_h * b_q[:, None], 0) + b_o = tl.sum(b_h * b_q[None, :], 1) tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) p_q += H * K @@ -113,11 +113,11 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( if not IS_KDA: p_g += HV else: - p_gk += HV * K + p_gk += H * K p_beta += HV * (V if IS_BETA_HEADWISE else 1) if STORE_FINAL_STATE: - p_ht = ht + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + p_ht = ht + i_nh * V * K + o_v[:, None] * K + o_k[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) @@ -144,7 +144,7 @@ def fused_recurrent_gated_delta_rule_fwd( o = q.new_empty(NK, *v.shape) if output_final_state: - final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) + final_state = q.new_empty(N, HV, V, K, dtype=torch.float32) else: final_state = None @@ -252,11 +252,11 @@ def fused_recurrent_gated_delta_rule( Scale factor for the RetNet attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, HV, K, V]` for `N` input sequences. + Initial state of shape `[N, HV, V, K]` for `N` input sequences. For equal-length input sequences, `N` equals the batch size `B`. Default: `None`. output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + Whether to output the final state of shape `[N, HV, V, K]`. Default: `False`. cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. @@ -264,7 +264,7 @@ def fused_recurrent_gated_delta_rule( o (torch.Tensor): Outputs of shape `[B, T, HV, V]`. final_state (torch.Tensor): - Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + Final state of shape `[N, HV, V, K]` if `output_final_state=True` else `None`. Examples:: >>> import torch >>> import torch.nn.functional as F @@ -277,7 +277,7 @@ def fused_recurrent_gated_delta_rule( >>> v = torch.randn(B, T, HV, V, device='cuda') >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda')) >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() - >>> h0 = torch.randn(B, HV, K, V, device='cuda') + >>> h0 = torch.randn(B, HV, V, K, device='cuda') >>> o, ht = fused_gated_recurrent_delta_rule( q, k, v, g, beta, initial_state=h0, @@ -413,9 +413,9 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( mask_k = o_k < K mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] + mask_h = mask_v[:, None] & mask_k[None, :] - b_h = tl.zeros([BK, BV], dtype=tl.float32) + b_h = tl.zeros([BV, BK], dtype=tl.float32) if USE_INITIAL_STATE: idx = tl.load(h0_indices + i_n) # Add bounds checking for idx @@ -424,8 +424,8 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( h0_source + idx * HV * K * V + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] + + o_v[:, None] * K + + o_k[None, :] ) b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) @@ -449,8 +449,8 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( + cache_idx * cache_steps * HV * K * V + step_offset + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] + + o_v[:, None] * K + + o_k[None, :] ) b_h = tl.load(cache_ptr, mask=mask_h, other=0).to(tl.float32) @@ -466,17 +466,17 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( # [BK, BV] b_h *= exp(b_g) # [BV] - b_v -= tl.sum(b_h * b_k[:, None], 0) + b_v -= tl.sum(b_h * b_k[None, :], 1) if IS_BETA_HEADWISE: b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) else: b_beta = tl.load(p_beta).to(tl.float32) b_v *= b_beta - # [BK, BV] - b_h += b_k[:, None] * b_v[None, :] + # [BV, BK] + b_h += b_v[:, None] * b_k[None, :] # [BV] if not DISABLE_OUTPUT_CALCULATION: - b_o = tl.sum(b_h * b_q[:, None], 0) + b_o = tl.sum(b_h * b_q[None, :], 1) # core attn output tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) @@ -490,8 +490,8 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( + cache_idx * cache_steps * HV * K * V + step_offset + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] + + o_v[:, None] * K + + o_k[None, :] ) tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h) @@ -513,8 +513,8 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( h0_source + idx * HV * K * V + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] + + o_v[:, None] * K + + o_k[None, :] ) tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) diff --git a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py index f140ccae4..f3e94a035 100644 --- a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py +++ b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py @@ -99,8 +99,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel( h0_source + idx * HV * K * V + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] + + o_v[None, :] * K + + o_k[:, None] ) b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) @@ -137,8 +137,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel( + cache_idx * cache_steps * HV * K * V + step_offset + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] + + o_v[None, :] * K + + o_k[:, None] ) b_h = tl.load(cache_ptr, mask=mask_h, other=0).to(tl.float32) @@ -207,8 +207,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel( + cache_idx * cache_steps * HV * K * V + step_offset + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] + + o_v[None, :] * K + + o_k[:, None] ) tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h) @@ -234,8 +234,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel( h0_source + idx * HV * K * V + i_hv * K * V - + o_k[:, None] * V - + o_v[None, :] + + o_v[None, :] * K + + o_k[:, None] ) tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h) diff --git a/python/sglang/srt/layers/attention/fla/kda.py b/python/sglang/srt/layers/attention/fla/kda.py index 286746332..3f17b21cc 100644 --- a/python/sglang/srt/layers/attention/fla/kda.py +++ b/python/sglang/srt/layers/attention/fla/kda.py @@ -59,11 +59,11 @@ def fused_recurrent_kda_fwd( num_stages = 3 num_warps = 1 - o = torch.empty_like(k) + o = q.new_empty(NK, *v.shape) if inplace_final_state: final_state = initial_state else: - final_state = q.new_empty(T, HV, K, V, dtype=initial_state.dtype) + final_state = q.new_empty(N, HV, V, K, dtype=initial_state.dtype) stride_init_state_token = initial_state.stride(0) stride_final_state_token = final_state.stride(0) @@ -113,6 +113,7 @@ def fused_recurrent_kda_fwd( num_stages=num_stages, ) + o = o.squeeze(0) return o, final_state @@ -756,11 +757,11 @@ def chunk_gla_fwd_kernel_o( (1, 0), ) p_h = tl.make_block_ptr( - h + (i_tg * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, i_v * BV), - (BK, BV), + h + (i_tg * H + i_h) * V * K, + (V, K), + (K, 1), + (i_v * BV, i_k * BK), + (BV, BK), (1, 0), ) @@ -776,7 +777,7 @@ def chunk_gla_fwd_kernel_o( # works but dkw, owing to divine benevolence # [BT, BV] if i_k >= 0: - b_o += tl.dot(b_qg, b_h.to(b_qg.dtype)) + b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) p_v = tl.make_block_ptr( v + (bos * H + i_h) * V, (T, V), diff --git a/test/registered/attention/test_chunk_gated_delta_rule.py b/test/registered/attention/test_chunk_gated_delta_rule.py new file mode 100644 index 000000000..8fb71e2c8 --- /dev/null +++ b/test/registered/attention/test_chunk_gated_delta_rule.py @@ -0,0 +1,274 @@ +import unittest + +import torch + +from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule +from sglang.srt.layers.attention.fla.fused_recurrent import ( + fused_recurrent_gated_delta_rule, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, suite="stage-b-test-large-1-gpu") + + +@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") +class TestChunkGatedDeltaRule(unittest.TestCase): + """Test chunk_gated_delta_rule against token-by-token fused_recurrent reference.""" + + ATOL = 2e-2 + RTOL = 1e-2 + + def _run_reference(self, pool_init, cache_indices, q, k, v, g, beta): + """Per-batch token-by-token reference using fused_recurrent_gated_delta_rule. + + initial_state shape: [N, H, V, K] (native layout on this branch). + """ + B = cache_indices.shape[0] + T_per_seq = q.shape[1] // B + pool = pool_init.clone() + h_cur = pool[cache_indices].contiguous().clone() + + o_list = [] + for b in range(B): + sl = slice(b * T_per_seq, (b + 1) * T_per_seq) + o_b, h_b = fused_recurrent_gated_delta_rule( + q=q[0, sl].unsqueeze(0), + k=k[0, sl].unsqueeze(0), + v=v[0, sl].unsqueeze(0), + g=g[0, sl].unsqueeze(0), + beta=beta[0, sl].unsqueeze(0), + initial_state=h_cur[b : b + 1], + output_final_state=True, + use_qk_l2norm_in_kernel=True, + ) + o_list.append(o_b) + h_cur[b] = h_b[0] + + pool[cache_indices] = h_cur + return torch.cat(o_list, dim=1), pool + + def _run_chunk(self, pool_init, cache_indices, q, k, v, g, beta, cu_seqlens): + """Run chunk_gated_delta_rule with native [V, K] pool.""" + pool = pool_init.clone() + o, _, _ = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=pool, + initial_state_indices=cache_indices, + cu_seqlens=cu_seqlens, + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + return o, pool + + def _check_shape( + self, B, T_per_seq, H, K, V, pool_size, sequential_indices=False, seed=42 + ): + """Run correctness check for one (B, T_per_seq, H, K, V, pool_size) config.""" + device = "cuda" + dtype = torch.bfloat16 + T = B * T_per_seq + + torch.manual_seed(seed) + + if sequential_indices: + cache_indices = torch.arange(B, dtype=torch.int32, device=device) + else: + perm = torch.randperm(pool_size, device=device)[:B] + cache_indices = perm.to(torch.int32) + + pool_init = ( + torch.randn(pool_size, H, V, K, dtype=torch.float32, device=device) * 0.1 + ) + cu_seqlens = torch.zeros(B + 1, dtype=torch.long, device=device) + cu_seqlens[1:] = ( + torch.arange(1, B + 1, dtype=torch.long, device=device) * T_per_seq + ) + + q = torch.randn(1, T, H, K, dtype=dtype, device=device) + k = torch.randn(1, T, H, K, dtype=dtype, device=device) + v = torch.randn(1, T, H, V, dtype=dtype, device=device) + g = torch.nn.functional.logsigmoid( + torch.randn(1, T, H, dtype=dtype, device=device) + ) + beta = torch.sigmoid(torch.randn(1, T, H, dtype=dtype, device=device)) + + o_ref, pool_ref = self._run_reference( + pool_init, cache_indices, q, k, v, g, beta + ) + o_new, pool_new = self._run_chunk( + pool_init, cache_indices, q, k, v, g, beta, cu_seqlens + ) + + self.assertTrue( + torch.allclose( + o_ref.float(), o_new.float(), atol=self.ATOL, rtol=self.RTOL + ), + f"Output mismatch: max_diff=" + f"{(o_ref.float() - o_new.float()).abs().max().item():.2e}", + ) + + ref_slots = pool_ref[cache_indices].contiguous() + new_slots = pool_new[cache_indices].contiguous() + self.assertTrue( + torch.allclose( + ref_slots.float(), new_slots.float(), atol=self.ATOL, rtol=self.RTOL + ), + f"State mismatch: max_diff=" + f"{(ref_slots.float() - new_slots.float()).abs().max().item():.2e}", + ) + + # ------------------------------------------------------------------ + # Production-style configs (Qwen3-Next) + # ------------------------------------------------------------------ + def test_production_nt1(self): + self._check_shape(B=4, T_per_seq=64, H=16, K=128, V=128, pool_size=32) + + def test_production_nt2(self): + self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=128, pool_size=32) + + def test_production_nt4(self): + self._check_shape(B=4, T_per_seq=256, H=16, K=128, V=128, pool_size=32) + + # ------------------------------------------------------------------ + # Batch size sweep + # ------------------------------------------------------------------ + def test_batch_1(self): + self._check_shape(B=1, T_per_seq=128, H=16, K=128, V=128, pool_size=32) + + def test_batch_2(self): + self._check_shape(B=2, T_per_seq=128, H=16, K=128, V=128, pool_size=32) + + def test_batch_8(self): + self._check_shape(B=8, T_per_seq=128, H=16, K=128, V=128, pool_size=64) + + def test_batch_16(self): + self._check_shape(B=16, T_per_seq=64, H=16, K=128, V=128, pool_size=128) + + def test_batch_32(self): + self._check_shape(B=32, T_per_seq=32, H=16, K=128, V=128, pool_size=256) + + # ------------------------------------------------------------------ + # Head count sweep + # ------------------------------------------------------------------ + def test_heads_4(self): + self._check_shape(B=4, T_per_seq=128, H=4, K=128, V=128, pool_size=32) + + def test_heads_8(self): + self._check_shape(B=4, T_per_seq=128, H=8, K=128, V=128, pool_size=32) + + def test_heads_32(self): + self._check_shape(B=4, T_per_seq=128, H=32, K=128, V=128, pool_size=32) + + def test_heads_64(self): + self._check_shape(B=4, T_per_seq=128, H=64, K=128, V=128, pool_size=32) + + # ------------------------------------------------------------------ + # K != V (exercises that [V,K] != [K,V] byte-order matters) + # ------------------------------------------------------------------ + def test_dim_64x64(self): + self._check_shape(B=4, T_per_seq=128, H=16, K=64, V=64, pool_size=32) + + def test_dim_k_lt_v(self): + self._check_shape(B=4, T_per_seq=128, H=16, K=64, V=128, pool_size=32) + + def test_dim_k_gt_v(self): + self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=64, pool_size=32) + + def test_dim_256x256(self): + self._check_shape(B=4, T_per_seq=128, H=16, K=256, V=256, pool_size=32) + + # ------------------------------------------------------------------ + # Short sequences (T < chunk_size=64) + # ------------------------------------------------------------------ + def test_seqlen_1(self): + self._check_shape(B=4, T_per_seq=1, H=16, K=128, V=128, pool_size=32) + + def test_seqlen_7(self): + self._check_shape(B=4, T_per_seq=7, H=16, K=128, V=128, pool_size=32) + + def test_seqlen_16(self): + self._check_shape(B=4, T_per_seq=16, H=16, K=128, V=128, pool_size=32) + + def test_seqlen_32(self): + self._check_shape(B=4, T_per_seq=32, H=16, K=128, V=128, pool_size=32) + + # ------------------------------------------------------------------ + # Multi-chunk and large pool + # ------------------------------------------------------------------ + def test_multi_chunk_nt8(self): + self._check_shape(B=4, T_per_seq=512, H=16, K=128, V=128, pool_size=32) + + def test_large_pool(self): + self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=128, pool_size=512) + + # ------------------------------------------------------------------ + # Combined stress + # ------------------------------------------------------------------ + def test_stress(self): + self._check_shape(B=32, T_per_seq=128, H=32, K=128, V=128, pool_size=256) + + # ------------------------------------------------------------------ + # Sequential-index variants (pool_size == B, indices = 0..B-1) + # ------------------------------------------------------------------ + def test_seq_idx_b4(self): + self._check_shape( + B=4, + T_per_seq=128, + H=16, + K=128, + V=128, + pool_size=4, + sequential_indices=True, + ) + + def test_seq_idx_b8(self): + self._check_shape( + B=8, + T_per_seq=128, + H=16, + K=128, + V=128, + pool_size=8, + sequential_indices=True, + ) + + def test_seq_idx_h32(self): + self._check_shape( + B=4, + T_per_seq=128, + H=32, + K=128, + V=128, + pool_size=4, + sequential_indices=True, + ) + + def test_seq_idx_h64(self): + self._check_shape( + B=4, + T_per_seq=128, + H=64, + K=128, + V=128, + pool_size=4, + sequential_indices=True, + ) + + def test_seq_idx_stress(self): + self._check_shape( + B=32, + T_per_seq=128, + H=32, + K=128, + V=128, + pool_size=32, + sequential_indices=True, + ) + + +if __name__ == "__main__": + unittest.main()