diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py index 41d228b735a4..5914a687b261 100644 --- a/python/sglang/srt/layers/attention/fla/chunk.py +++ b/python/sglang/srt/layers/attention/fla/chunk.py @@ -2,6 +2,7 @@ # -*- coding: utf-8 -*- # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +import os from typing import Optional import torch @@ -32,6 +33,306 @@ _is_hip = is_hip() +try: + from aiter.ops.chunk_gated_delta_rule_fwd_h import ( + chunk_gated_delta_rule_fwd_h_hip_fn as aiter_chunk_gated_delta_rule_fwd_h_hip_fn, + ) + from aiter.ops.triton.gated_delta_net.gated_delta_rule import ( + chunk_gated_delta_rule_opt_vk as aiter_chunk_gated_delta_rule_opt_vk, + ) + from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h_opt_vk as aiter_chunk_gated_delta_rule_fwd_h_triton_opt_vk, + ) + from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.chunk_o import ( + chunk_fwd_o_opt_vk as aiter_chunk_fwd_o_opt_vk, + ) + from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.fused_cumsum_kkt import ( + fused_chunk_local_cumsum_scaled_dot_kkt_fwd as aiter_fused_cumsum_kkt_opt_vk, + ) + from aiter.ops.triton._triton_kernels.gated_delta_rule.prefill.fused_solve_tril_recompute import ( + fused_solve_tril_recompute_w_u as aiter_fused_solve_tril_recompute_w_u, + ) + + _aiter_prefill_opt_vk_available = True +except Exception: + aiter_chunk_gated_delta_rule_fwd_h_hip_fn = None + aiter_chunk_gated_delta_rule_opt_vk = None + aiter_chunk_gated_delta_rule_fwd_h_triton_opt_vk = None + aiter_chunk_fwd_o_opt_vk = None + aiter_fused_cumsum_kkt_opt_vk = None + aiter_fused_solve_tril_recompute_w_u = None + _aiter_prefill_opt_vk_available = False + + +def _env_enabled(name: str) -> bool: + return os.getenv(name, "0").lower() in ("1", "true", "yes", "on") + + +def should_use_aiter_prefill_opt_vk_for_seq_len(seq_len: int) -> bool: + return aiter_prefill_opt_vk_enabled() + + +def _can_use_aiter_prefill_opt_vk( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: Optional[torch.Tensor], + initial_state_indices: Optional[torch.Tensor], + cu_seqlens: Optional[torch.Tensor], +) -> bool: + if not should_use_aiter_prefill_opt_vk_for_seq_len(q.shape[1]): + return False + if q.dtype != torch.bfloat16 or k.dtype != torch.bfloat16 or v.dtype != torch.bfloat16: + return False + if k.shape[-1] != 128 or v.shape[-1] != 128: + return False + if g.dtype != torch.float32 or beta.dtype != torch.float32: + return False + if initial_state is None or initial_state_indices is None or cu_seqlens is None: + return False + if initial_state.dtype not in (torch.float32, torch.bfloat16): + return False + if initial_state.shape[-2:] != (128, 128): + return False + if initial_state.shape[1] != v.shape[2]: + return False + return True + + +def aiter_prefill_opt_vk_enabled() -> bool: + return ( + _is_hip + and _env_enabled("SGLANG_GDN_PREFILL_OPT_VK") + and _aiter_prefill_opt_vk_available + ) + + +def selected_aiter_prefill_opt_vk_k5_backend() -> str: + # K5 (chunk-delta-h) backend: default "hip"; override with "triton"/"auto". + backend = os.getenv("SGLANG_GDN_PREFILL_OPT_VK_K5", "hip").lower() + if backend not in ("auto", "hip", "triton"): + raise ValueError( + "SGLANG_GDN_PREFILL_OPT_VK_K5 must be 'auto', 'hip', or 'triton', " + f"got {backend!r}." + ) + return backend + + +def _resolved_aiter_prefill_opt_vk_k5_backend(seq_len: int) -> str: + backend = selected_aiter_prefill_opt_vk_k5_backend() + if backend != "auto": + return backend + hip_min_t = int(os.getenv("SGLANG_GDN_PREFILL_OPT_VK_HIP_MIN_T", "8192")) + return "hip" if seq_len >= hip_min_t else "triton" + + +def _chunk_gated_delta_rule_fwd_aiter_opt_vk( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + initial_state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + initial_state_layout: int = 0, + output_state_layout: int = 0, + return_intermediate_h: bool = True, +): + """Experimental env-gated opt-vk prefill core.""" + state_indices = initial_state_indices.to(torch.long) + k5_backend = _resolved_aiter_prefill_opt_vk_k5_backend(k.shape[1]) + pool_indexed_vk_state = initial_state_layout == 1 and output_state_layout == 1 + if pool_indexed_vk_state: + state_vk = initial_state + elif initial_state_layout == 1: + state_vk = initial_state[state_indices].contiguous() + else: + state_vk = initial_state[state_indices].transpose(-1, -2).contiguous() + if not return_intermediate_h: + o, final_state_vk = aiter_chunk_gated_delta_rule_opt_vk( + q=q, + k=k, + v=v, + o=v.new_empty(v.shape), + g=g, + beta=beta, + scale=scale, + initial_state=state_vk, + initial_state_indices=initial_state_indices if pool_indexed_vk_state else None, + output_final_state=True, + inplace_final_state=True if pool_indexed_vk_state else None, + cu_seqlens=cu_seqlens, + use_chunk_hip=k5_backend == "hip", + state_dtype=state_vk.dtype, + use_exp2=True, + ) + if pool_indexed_vk_state: + pass + elif output_state_layout == 1: + initial_state.index_copy_( + 0, state_indices, final_state_vk.to(initial_state.dtype) + ) + else: + final_state_kv = final_state_vk.transpose(-1, -2).contiguous() + initial_state.index_copy_( + 0, state_indices, final_state_kv.to(initial_state.dtype) + ) + empty_h = q.new_empty(0) + return g.new_empty(0), o, None, None, empty_h, None + + g_cumsum, A = aiter_fused_cumsum_kkt_opt_vk( + k=k, + beta=beta, + g=g, + cu_seqlens=cu_seqlens, + use_exp2=True, + ) + w, u = aiter_fused_solve_tril_recompute_w_u( + A_raw=A, + k=k, + v=v, + beta=beta, + g_cumsum=g_cumsum, + cu_seqlens=cu_seqlens, + use_exp2=True, + ) + + if k5_backend == "triton" or pool_indexed_vk_state: + h_vk, v_new, final_state_vk = aiter_chunk_gated_delta_rule_fwd_h_triton_opt_vk( + k=k, + w=w, + u=u, + g=g_cumsum, + initial_state=state_vk, + initial_state_indices=initial_state_indices if pool_indexed_vk_state else None, + output_final_state=True, + inplace_final_state=True if pool_indexed_vk_state else None, + cu_seqlens=cu_seqlens, + state_dtype=state_vk.dtype, + use_exp2=True, + ) + else: + h_vk, v_new, final_state_vk = aiter_chunk_gated_delta_rule_fwd_h_hip_fn( + k=k, + w=w, + u=u, + g=g_cumsum, + initial_state=state_vk, + output_final_state=True, + cu_seqlens=cu_seqlens, + state_dtype=state_vk.dtype, + use_exp2=True, + g_head_major=True, + ) + + o = v.new_empty(v.shape) + o = aiter_chunk_fwd_o_opt_vk( + q=q, + k=k, + v=v_new, + o=o, + h=h_vk, + g=g_cumsum, + scale=scale, + cu_seqlens=cu_seqlens, + use_exp2=True, + ) + + if pool_indexed_vk_state: + h = h_vk + elif output_state_layout == 1: + initial_state.index_copy_(0, state_indices, final_state_vk.to(initial_state.dtype)) + h = h_vk + else: + final_state_kv = final_state_vk.transpose(-1, -2).contiguous() + initial_state.index_copy_(0, state_indices, final_state_kv.to(initial_state.dtype)) + h = h_vk.transpose(-1, -2).contiguous() + + return g_cumsum.transpose(1, 2).contiguous(), o, A, w, h, v_new + + +def chunk_gated_delta_rule_prefill_opt_vk_no_h( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + initial_state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + use_qk_l2norm_in_kernel: bool = True, +): + """Fast inference-only opt-vk prefill path for canonical VK state pools.""" + if use_qk_l2norm_in_kernel: + if _is_hip: + q, k = fused_l2norm_qk(q, k) + else: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + k5_backend = _resolved_aiter_prefill_opt_vk_k5_backend(k.shape[1]) + g_cumsum, A = aiter_fused_cumsum_kkt_opt_vk( + k=k, + beta=beta, + g=g, + cu_seqlens=cu_seqlens, + use_exp2=True, + ) + w, u = aiter_fused_solve_tril_recompute_w_u( + A_raw=A, + k=k, + v=v, + beta=beta, + g_cumsum=g_cumsum, + cu_seqlens=cu_seqlens, + use_exp2=True, + ) + if k5_backend == "hip": + h, v_new, _ = aiter_chunk_gated_delta_rule_fwd_h_hip_fn( + k=k, + w=w, + u=u, + g=g_cumsum, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + output_final_state=True, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + state_dtype=initial_state.dtype, + use_exp2=True, + g_head_major=True, + ) + else: + h, v_new, _ = aiter_chunk_gated_delta_rule_fwd_h_triton_opt_vk( + k=k, + w=w, + u=u, + g=g_cumsum, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + output_final_state=True, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + state_dtype=initial_state.dtype, + use_exp2=True, + ) + o = aiter_chunk_fwd_o_opt_vk( + q=q, + k=k, + v=v_new, + o=v.new_empty(v.shape), + h=h, + g=g_cumsum, + scale=scale, + cu_seqlens=cu_seqlens, + use_exp2=True, + ) + return o.to(q.dtype), None, q.new_empty(0) + def chunk_gated_delta_rule_fwd( q: torch.Tensor, @@ -43,10 +344,38 @@ def chunk_gated_delta_rule_fwd( initial_state: torch.Tensor, initial_state_indices: torch.Tensor, cu_seqlens: Optional[torch.LongTensor] = None, + initial_state_layout: int = 0, + output_state_layout: int = 0, + return_intermediate_h: bool = True, ): B, T = q.shape[0], q.shape[1] Hv = g.shape[2] + if _can_use_aiter_prefill_opt_vk( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + cu_seqlens=cu_seqlens, + ): + return _chunk_gated_delta_rule_fwd_aiter_opt_vk( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + cu_seqlens=cu_seqlens, + initial_state_layout=initial_state_layout, + output_state_layout=output_state_layout, + return_intermediate_h=return_intermediate_h, + ) + if _is_hip and T >= 64: g, A = fused_cumsum_kkt(g, k, beta, chunk_size=64, cu_seqlens=cu_seqlens) chunk_indices_16 = ( @@ -127,6 +456,9 @@ def forward( initial_state_indices: torch.Tensor, cu_seqlens: Optional[torch.LongTensor] = None, use_qk_l2norm_in_kernel: bool = False, + initial_state_layout: int = 0, + output_state_layout: int = 0, + return_intermediate_h: bool = True, ): q_orig = q k_orig = k @@ -148,6 +480,9 @@ def forward( initial_state=initial_state, initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, + initial_state_layout=initial_state_layout, + output_state_layout=output_state_layout, + return_intermediate_h=return_intermediate_h, ) return o.to(q.dtype), h @@ -165,6 +500,9 @@ def chunk_gated_delta_rule( cu_seqlens: Optional[torch.LongTensor] = None, head_first: bool = False, use_qk_l2norm_in_kernel: bool = False, + initial_state_layout: int = 0, + output_state_layout: int = 0, + return_intermediate_h: bool = True, ): r""" Args: @@ -268,6 +606,57 @@ def chunk_gated_delta_rule( ) if scale is None: scale = k.shape[-1] ** -0.5 + if not return_intermediate_h and _can_use_aiter_prefill_opt_vk( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + cu_seqlens=cu_seqlens, + ): + if use_qk_l2norm_in_kernel: + if _is_hip: + q, k = fused_l2norm_qk(q, k) + else: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + if initial_state_layout == 1 and output_state_layout == 1: + k5_backend = _resolved_aiter_prefill_opt_vk_k5_backend(k.shape[1]) + o, _ = aiter_chunk_gated_delta_rule_opt_vk( + q=q, + k=k, + v=v, + o=v.new_empty(v.shape), + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + output_final_state=True, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + use_chunk_hip=k5_backend == "hip", + state_dtype=initial_state.dtype, + use_exp2=True, + ) + return o.to(q.dtype), None, q.new_empty(0) + _, o, _, _, h, _ = _chunk_gated_delta_rule_fwd_aiter_opt_vk( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + cu_seqlens=cu_seqlens, + initial_state_layout=initial_state_layout, + output_state_layout=output_state_layout, + return_intermediate_h=False, + ) + return o.to(q.dtype), None, h o, h = ChunkGatedDeltaRuleFunction.apply( q, k, @@ -279,6 +668,9 @@ def chunk_gated_delta_rule( initial_state_indices, cu_seqlens, use_qk_l2norm_in_kernel, + initial_state_layout, + output_state_layout, + return_intermediate_h, ) if head_first: o = rearrange(o, "b t h ... -> b h t ...") diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 0611c1f730d5..44273d8fa91b 100755 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -72,7 +72,11 @@ # CuTe DSL path requires cuda-python (cuda.bindings.*). Keep runtime usable # by falling back to non-CuTe kernels when it's unavailable. cutedsl_fused_sigmoid_gating_delta_rule_update = None - from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule + from sglang.srt.layers.attention.fla.chunk import ( + aiter_prefill_opt_vk_enabled, + chunk_gated_delta_rule, + chunk_gated_delta_rule_prefill_opt_vk_no_h, + ) from sglang.srt.layers.attention.fla.chunk_delta_h import ( CHUNK_SIZE as FLA_CHUNK_SIZE, ) @@ -267,7 +271,7 @@ def copy_h_to_ssm_track_kernel( h_stride_0, ssm_stride_0, row_numel: tl.constexpr, - layout_kv: tl.constexpr, + producer_layout: tl.constexpr, HAS_LAYOUT: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): @@ -283,7 +287,7 @@ def copy_h_to_ssm_track_kernel( tl.store(ssm_states_ptr + dst_idx * ssm_stride_0 + offsets, data, mask=mask) if HAS_LAYOUT and tile_idx == 0: - tl.store(slot_layout_ptr + dst_idx, layout_kv) + tl.store(slot_layout_ptr + dst_idx, producer_layout) @triton.jit @@ -319,7 +323,7 @@ def copy_h_to_ssm_track( src_indices: torch.Tensor, dst_indices: torch.Tensor, slot_layout: Optional[torch.Tensor], - layout_kv: int, + producer_layout: int, ) -> None: if src_indices.numel() == 0: return @@ -335,7 +339,7 @@ def copy_h_to_ssm_track( h.stride(0), ssm_states.stride(0), row_numel, - layout_kv, + producer_layout, slot_layout is not None, block_size, ) @@ -800,6 +804,7 @@ def _track_mamba_state_extend( h: torch.Tensor, ssm_states: torch.Tensor, forward_metadata: ForwardMetadata, + h_layout: Optional[int] = None, ): """ Track and copy SSM states during extend for prefix caching. @@ -822,7 +827,7 @@ def _track_mamba_state_extend( forward_metadata.track_ssm_h_src, forward_metadata.track_ssm_h_dst, self._slot_layout, - self._layout_kv, + self._layout_kv if h_layout is None else h_layout, ) if forward_metadata.track_ssm_final_src.numel() > 0: copy_ssm_to_ssm_track( @@ -1032,6 +1037,9 @@ def __init__(self, model_runner: ModelRunner): self._num_v_heads_per_layer: int = 0 self._layout_kv: int = 0 self._layout_vk: int = 1 + # Device-side scalar mirror of _layout_vk (avoids a per-layer H2D sync on scatter). + self._layout_vk_scalar: Optional[torch.Tensor] = None + self._prefill_opt_vk_enabled: bool = False if _use_hip_linear_attn and _is_hip: local_num_k_heads = None local_num_v_heads = None @@ -1147,6 +1155,7 @@ def __init__(self, model_runner: ModelRunner): self._full_temporal_state = full_temporal self._num_mamba_layers = full_temporal.shape[0] self._num_v_heads_per_layer = full_temporal.shape[2] + self._prefill_opt_vk_enabled = aiter_prefill_opt_vk_enabled() total_slots = full_temporal.shape[1] # Share the pool-owned bitmap (single source of truth) so # prefix-cache reuse can't desync the KV<->VK layout. @@ -1159,7 +1168,8 @@ def __init__(self, model_runner: ModelRunner): f"{selected_vk_backend.upper()} GDN decode loaded " f"(local heads {local_num_k_heads}/{local_num_v_heads}, " f"{self._num_mamba_layers} layers, batched multi-layer " - "transpose, per-slot gate)." + "transpose, per-slot gate, " + f"prefill opt-vk={self._prefill_opt_vk_enabled})." ) else: rank0_log( @@ -1209,6 +1219,8 @@ def _prepare_gdn_layout_for_forward_mode(self, forward_mode: ForwardMode): layout_kv=self._layout_kv, layout_vk=self._layout_vk, ) + if forward_mode.is_extend() and self._prefill_opt_vk_enabled: + target_layout = self._layout_vk if target_layout is not None: self._eager_apply_layout_transition(target_layout) if ( @@ -1602,18 +1614,62 @@ def forward_extend( if is_npu() or is_cpu(): recurrent_state = ssm_states[cache_indices] recurrent_state_indices_args = {} - core_attn_out, last_recurrent_state, h = chunk_gated_delta_rule( - q=query, - k=key, - v=value, - g=g, - beta=beta, - initial_state=recurrent_state, - cu_seqlens=query_start_loc, - head_first=False, - use_qk_l2norm_in_kernel=True, - **recurrent_state_indices_args, + prefill_state_layout = ( + self._layout_vk if self._prefill_opt_vk_enabled else self._layout_kv ) + need_intermediate_h = not self._prefill_opt_vk_enabled or ( + forward_metadata.has_mamba_track_mask + and forward_metadata.track_ssm_h_src.numel() > 0 + ) + if self._prefill_opt_vk_enabled and not need_intermediate_h: + core_attn_out, last_recurrent_state, h = ( + chunk_gated_delta_rule_prefill_opt_vk_no_h( + q=query, + k=key, + v=value, + g=g, + beta=beta, + scale=key.shape[-1] ** -0.5, + initial_state=recurrent_state, + initial_state_indices=cache_indices, + cu_seqlens=query_start_loc, + use_qk_l2norm_in_kernel=True, + ) + ) + else: + core_attn_out, last_recurrent_state, h = chunk_gated_delta_rule( + q=query, + k=key, + v=value, + g=g, + beta=beta, + initial_state=recurrent_state, + cu_seqlens=query_start_loc, + head_first=False, + use_qk_l2norm_in_kernel=True, + initial_state_layout=prefill_state_layout, + output_state_layout=prefill_state_layout, + return_intermediate_h=need_intermediate_h, + **recurrent_state_indices_args, + ) + if ( + self._prefill_opt_vk_enabled + and self._slot_layout is not None + and cache_indices is not None + and cache_indices.numel() > 0 + ): + # Scatter via the cached device scalar to keep it on-device (no H2D sync). + if ( + self._layout_vk_scalar is None + or self._layout_vk_scalar.device != self._slot_layout.device + or self._layout_vk_scalar.dtype != self._slot_layout.dtype + ): + self._layout_vk_scalar = torch.tensor( + self._layout_vk, + dtype=self._slot_layout.dtype, + device=self._slot_layout.device, + ) + self._slot_layout[cache_indices] = self._layout_vk_scalar if is_npu() or is_cpu(): last_recurrent_state = last_recurrent_state.to( ssm_states.dtype, copy=False @@ -1621,11 +1677,14 @@ def forward_extend( ssm_states[cache_indices] = last_recurrent_state self._track_mamba_state_extend( - forward_batch, h, ssm_states, forward_metadata + forward_batch, + h, + ssm_states, + forward_metadata, + h_layout=prefill_state_layout, ) - # KV→VK is deferred to the next decode pass entry - # (_eager_apply_layout_transition). + # opt-vk leaves slots in VK; otherwise KV->VK defers to the next decode entry. return core_attn_out diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py index 3df6539f924c..7aec4377a9d7 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py @@ -1186,7 +1186,7 @@ def _fused_append_shared_experts_gated_kernel( tl.store(out_ids_ptr + out_ids_row_ptr + offs_k, ids, mask=mask_k) tl.store(out_weights_ptr + out_w_row_ptr + offs_k, ws, mask=mask_k) - shared_w = tl.load(shared_weights_ptr + pid) + shared_w = tl.load(shared_weights_ptr + pid).to(ws.dtype) offs_s = tl.arange(0, BLOCK_S) mask_s = offs_s < S shared_ids = tl.cast(N_BASE + offs_s, ids.dtype) @@ -1219,8 +1219,6 @@ def fused_append_shared_experts_gated( (m, k + s), dtype=topk_weights.dtype, device=topk_weights.device ) - shared_weights = shared_weights.to(topk_weights.dtype) - _fused_append_shared_experts_gated_kernel[(m,)]( topk_ids, topk_weights, diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 234b4e827e03..6473bcb912a2 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -613,7 +613,7 @@ def _forward( core_attn_out = core_attn_out_pad core_attn_out = self.norm(core_attn_out, z) - core_attn_out = core_attn_out.reshape(z_shape_og[0], -1) + core_attn_out = core_attn_out.view(z_shape_og[0], -1) output, _ = self.out_proj(core_attn_out) return output