# Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Fused Triton kernel for DFlash KV materialization. Combines: KV projection + RMSNorm + RoPE, then pool-managed KV writes. """ from typing import Callable, List, Optional import torch import triton import triton.language as tl @triton.jit def _fused_norm_rope_kernel_stacked( kv_ptr, # [total_ctx, n_layers, kv_size * 2] k_norm_weight_ptr, # [n_layers, head_dim] eps_ptr, # [n_layers] cos_sin_cache_ptr, # [max_pos, rotary_dim] positions_ptr, # [total_ctx] k_out_ptr, # [n_layers, total_ctx, num_kv_heads, head_dim] v_out_ptr, # [n_layers, total_ctx, num_kv_heads, head_dim] kv_stride_ctx, kv_stride_layer, k_norm_weight_stride_layer, cos_sin_stride_pos, k_out_stride_layer, k_out_stride_ctx, k_out_stride_head, v_out_stride_layer, v_out_stride_ctx, v_out_stride_head, total_ctx, n_layers: tl.constexpr, num_kv_heads: tl.constexpr, head_dim: tl.constexpr, kv_size: tl.constexpr, rotary_dim: tl.constexpr, half_rotary_dim: tl.constexpr, BLOCK_HD: tl.constexpr, ): """Fused RMSNorm(K) + RoPE(K) materialization. Grid: (total_ctx, num_kv_heads, n_layers).""" ctx_id = tl.program_id(0) head_id = tl.program_id(1) layer_id = tl.program_id(2) if ctx_id >= total_ctx or layer_id >= n_layers: return position = tl.load(positions_ptr + ctx_id) eps = tl.load(eps_ptr + layer_id).to(tl.float32) kv_base = kv_ptr + ctx_id * kv_stride_ctx + layer_id * kv_stride_layer k_base = kv_base + head_id * head_dim v_base = kv_base + kv_size + head_id * head_dim k_write = ( k_out_ptr + layer_id * k_out_stride_layer + ctx_id * k_out_stride_ctx + head_id * k_out_stride_head ) v_write = ( v_out_ptr + layer_id * v_out_stride_layer + ctx_id * v_out_stride_ctx + head_id * v_out_stride_head ) offs = tl.arange(0, BLOCK_HD) mask_hd = offs < head_dim mask_half = offs < half_rotary_dim k_raw = tl.load(k_base + offs, mask=mask_hd, other=0.0).to(tl.float32) v_raw = tl.load(v_base + offs, mask=mask_hd, other=0.0) inv_rms = tl.rsqrt(tl.sum(k_raw * k_raw) / head_dim + eps) norm_w = tl.load( k_norm_weight_ptr + layer_id * k_norm_weight_stride_layer + offs, mask=mask_hd, other=1.0, ).to(tl.float32) k_normed = k_raw * inv_rms * norm_w cos_sin_base = cos_sin_cache_ptr + position * cos_sin_stride_pos cos_v = tl.load(cos_sin_base + offs, mask=mask_half, other=1.0).to(tl.float32) sin_v = tl.load( cos_sin_base + half_rotary_dim + offs, mask=mask_half, other=0.0 ).to(tl.float32) k_first = tl.where(mask_half, k_normed, 0.0) k_second_raw = tl.load( k_base + half_rotary_dim + offs, mask=mask_half, other=0.0 ).to(tl.float32) norm_w_second = tl.load( k_norm_weight_ptr + layer_id * k_norm_weight_stride_layer + half_rotary_dim + offs, mask=mask_half, other=1.0, ).to(tl.float32) k_second = k_second_raw * inv_rms * norm_w_second k_rot_first = k_first * cos_v - k_second * sin_v k_rot_second = k_second * cos_v + k_first * sin_v tl.store(v_write + offs, v_raw, mask=mask_hd) tl.store(k_write + offs, k_rot_first.to(v_raw.dtype), mask=mask_half) tl.store( k_write + half_rotary_dim + offs, k_rot_second.to(v_raw.dtype), mask=mask_half ) mask_pass = (offs >= rotary_dim) & (offs < head_dim) tl.store(k_write + offs, k_normed.to(v_raw.dtype), mask=mask_pass) def _fused_norm_rope_stacked( kv: torch.Tensor, # [total_ctx, n_layers, kv_size*2] k_norm_weight: torch.Tensor, # [n_layers, head_dim] eps: torch.Tensor, # [n_layers] cos_sin_cache: torch.Tensor, # [max_pos, rotary_dim] positions: torch.Tensor, # [total_ctx] num_kv_heads: int, head_dim: int, rotary_dim: int, k_out: Optional[torch.Tensor] = None, v_out: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Fused RMSNorm + RoPE materialization for all layers.""" if kv.ndim != 3: raise ValueError( "Invalid stacked fused KV projection shape: " f"got {tuple(kv.shape)}, expected 3D [total_ctx, n_layers, kv_size*2]." ) total_ctx, n_layers, kv_dim = kv.shape if total_ctx == 0: empty = torch.empty( (n_layers, 0, num_kv_heads, head_dim), dtype=kv.dtype, device=kv.device ) return empty, empty kv_size = num_kv_heads * head_dim if kv_dim != kv_size * 2: raise ValueError( "Invalid fused KV projection shape: " f"got {tuple(kv.shape)}, expected trailing dim {kv_size * 2}." ) if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2 != 0: raise ValueError( "Invalid fused KV rotary/head dim pair: " f"rotary_dim={rotary_dim}, head_dim={head_dim}." ) if k_norm_weight.shape != (n_layers, head_dim): raise ValueError( "Invalid stacked k_norm_weight shape for fused KV materialization: " f"got {tuple(k_norm_weight.shape)}, expected {(n_layers, head_dim)}." ) if eps.shape != (n_layers,): raise ValueError( "Invalid stacked eps shape for fused KV materialization: " f"got {tuple(eps.shape)}, expected {(n_layers,)}." ) half_rotary_dim = rotary_dim // 2 BLOCK_HD = triton.next_power_of_2(head_dim) if positions.device != kv.device: positions = positions.to(device=kv.device, dtype=torch.int64) elif positions.dtype != torch.int64: positions = positions.to(torch.int64) expected_shape = (n_layers, total_ctx, num_kv_heads, head_dim) if k_out is None: k_out = torch.empty(expected_shape, dtype=kv.dtype, device=kv.device) else: if k_out.shape != expected_shape: raise ValueError( "Invalid k_out shape for fused KV materialization: " f"got {tuple(k_out.shape)}, expected {expected_shape}." ) if k_out.device != kv.device or k_out.dtype != kv.dtype: raise ValueError( "Invalid k_out device/dtype for fused KV materialization: " f"got device={k_out.device}, dtype={k_out.dtype}, " f"expected device={kv.device}, dtype={kv.dtype}." ) if v_out is None: v_out = torch.empty_like(k_out) else: if v_out.shape != expected_shape: raise ValueError( "Invalid v_out shape for fused KV materialization: " f"got {tuple(v_out.shape)}, expected {expected_shape}." ) if v_out.device != kv.device or v_out.dtype != kv.dtype: raise ValueError( "Invalid v_out device/dtype for fused KV materialization: " f"got device={v_out.device}, dtype={v_out.dtype}, " f"expected device={kv.device}, dtype={kv.dtype}." ) _fused_norm_rope_kernel_stacked[(total_ctx, num_kv_heads, n_layers)]( kv, k_norm_weight, eps, cos_sin_cache, positions, k_out, v_out, kv.stride(0), kv.stride(1), k_norm_weight.stride(0), cos_sin_cache.stride(0), k_out.stride(0), k_out.stride(1), k_out.stride(2), v_out.stride(0), v_out.stride(1), v_out.stride(2), total_ctx, n_layers, num_kv_heads, head_dim, kv_size, rotary_dim, half_rotary_dim, BLOCK_HD, ) return k_out, v_out class FusedKVMaterializeHelper: """Fused KV materialization helper using batched projection. Uses a single large GEMM across all layers, then a Triton kernel for fused RMSNorm + RoPE materialization across all layers. """ def __init__( self, layers: List, rotary_emb, num_kv_heads: int, head_dim: int, device: torch.device, max_position_hint: Optional[int] = None, ): self.num_kv_heads = num_kv_heads self.head_dim = head_dim self.rotary_emb = rotary_emb self.n_layers = len(layers) self.device = device self.kv_size = self.num_kv_heads * self.head_dim self.layer_out_dim = 2 * self.kv_size self.rotary_dim = int(getattr(rotary_emb, "rotary_dim", head_dim)) self.is_neox_style = bool(getattr(rotary_emb, "is_neox_style", True)) if not self.is_neox_style: raise NotImplementedError("Only neox-style RoPE is supported.") if self.rotary_dim <= 0 or self.rotary_dim > self.head_dim: raise ValueError( "Invalid fused KV rotary/head dim pair: " f"rotary_dim={self.rotary_dim}, head_dim={self.head_dim}." ) self.max_position_hint = ( max(int(max_position_hint) - 1, 0) if max_position_hint is not None else None ) self._reserved_rope_cache_len = int( getattr(self.rotary_emb, "cos_sin_cache", torch.empty((0,))).shape[0] ) self._mm_out_supported = True self._workspace_capacity = 0 self._workspace_dtype: Optional[torch.dtype] = None self._proj_workspace: Optional[torch.Tensor] = None self._k_workspace: Optional[torch.Tensor] = None self._v_workspace: Optional[torch.Tensor] = None kv_weights = [] k_norm_weights = [] eps_values = [] for layer_id, layer in enumerate(layers): attn = layer.self_attn if int(attn.num_kv_heads) != self.num_kv_heads: raise ValueError( "num_kv_heads mismatch across layers for fused KV path: " f"expected {self.num_kv_heads}, got {int(attn.num_kv_heads)} at layer {layer_id}." ) if int(attn.head_dim) != self.head_dim: raise ValueError( "head_dim mismatch across layers for fused KV path: " f"expected {self.head_dim}, got {int(attn.head_dim)} at layer {layer_id}." ) layer_rotary_dim = int( getattr(attn.rotary_emb, "rotary_dim", self.head_dim) ) layer_is_neox = bool(getattr(attn.rotary_emb, "is_neox_style", True)) if ( layer_rotary_dim != self.rotary_dim or layer_is_neox != self.is_neox_style ): raise ValueError( "RoPE config mismatch across layers for fused KV path: " f"expected (rotary_dim={self.rotary_dim}, neox={self.is_neox_style}), " f"got (rotary_dim={layer_rotary_dim}, neox={layer_is_neox}) at layer {layer_id}." ) qkv_w = attn.qkv_proj.weight kv_weight = qkv_w[attn.q_size : attn.q_size + 2 * attn.kv_size] kv_weights.append(kv_weight) k_norm_weights.append(attn.k_norm.weight) eps_values.append(float(attn.k_norm.variance_epsilon)) flat_kv_weight = torch.stack(kv_weights).reshape( self.n_layers * self.layer_out_dim, -1 ) self.flat_kv_weight_t = flat_kv_weight.transpose(0, 1).contiguous() self.k_norm_weights = torch.stack(k_norm_weights).contiguous() self.eps_values = torch.tensor( eps_values, dtype=torch.float32, device=self.device ) if self.max_position_hint is not None: self._ensure_rope_cache(self.max_position_hint) def _ensure_rope_cache(self, max_position: int) -> torch.Tensor: if max_position + 1 > self._reserved_rope_cache_len: ensure_cos_sin_cache_length = getattr( self.rotary_emb, "_ensure_cos_sin_cache_length", None ) if callable(ensure_cos_sin_cache_length): ensure_cos_sin_cache_length(max_position) self._reserved_rope_cache_len = int( self.rotary_emb.cos_sin_cache.shape[0] ) cos_sin_cache = self.rotary_emb.cos_sin_cache if max_position >= int(cos_sin_cache.shape[0]): raise RuntimeError( "RoPE cos/sin cache is too short for fused KV materialization: " f"max_position={max_position}, cache_len={int(cos_sin_cache.shape[0])}." ) if cos_sin_cache.device != self.device: cos_sin_cache = cos_sin_cache.to(self.device) return cos_sin_cache def _ensure_workspace(self, total_ctx: int, dtype: torch.dtype) -> None: if ( self._workspace_capacity >= total_ctx and self._workspace_dtype == dtype and self._proj_workspace is not None and self._k_workspace is not None and self._v_workspace is not None ): return new_capacity = max(1, total_ctx) if self._workspace_capacity > 0: new_capacity = max(new_capacity, self._workspace_capacity * 2) self._proj_workspace = torch.empty( (new_capacity, self.n_layers * self.layer_out_dim), dtype=dtype, device=self.device, ) self._k_workspace = torch.empty( (self.n_layers, new_capacity, self.num_kv_heads, self.head_dim), dtype=dtype, device=self.device, ) self._v_workspace = torch.empty_like(self._k_workspace) self._workspace_capacity = new_capacity self._workspace_dtype = dtype def materialize( self, ctx_hidden: torch.Tensor, positions: torch.Tensor, write_layer_kv: Callable[[int, torch.Tensor, torch.Tensor], None], ) -> None: """Materialize KV cache for all layers using batched projection.""" total_ctx = ctx_hidden.shape[0] if total_ctx == 0: return if positions.ndim != 1: positions = positions.reshape(-1) if positions.numel() != total_ctx: raise ValueError( "positions must match ctx_hidden token count for fused KV materialization: " f"positions={positions.numel()}, total_ctx={total_ctx}." ) if ctx_hidden.device != self.device: ctx_hidden = ctx_hidden.to(self.device, non_blocking=True) if ctx_hidden.dtype != self.flat_kv_weight_t.dtype: ctx_hidden = ctx_hidden.to(self.flat_kv_weight_t.dtype) if positions.device != self.device: positions = positions.to( device=self.device, dtype=torch.int64, non_blocking=True ) elif positions.dtype != torch.int64: positions = positions.to(torch.int64) max_position = ( self.max_position_hint if self.max_position_hint is not None else int(positions.max().item()) ) cos_sin_cache = self._ensure_rope_cache(max_position) self._ensure_workspace(total_ctx, ctx_hidden.dtype) assert self._proj_workspace is not None assert self._k_workspace is not None assert self._v_workspace is not None proj_out_2d = self._proj_workspace[:total_ctx] if self._mm_out_supported: try: torch.mm(ctx_hidden, self.flat_kv_weight_t, out=proj_out_2d) except Exception: self._mm_out_supported = False proj_out_2d = torch.mm(ctx_hidden, self.flat_kv_weight_t) else: proj_out_2d = torch.mm(ctx_hidden, self.flat_kv_weight_t) proj_out = proj_out_2d.view(total_ctx, self.n_layers, self.layer_out_dim) tmp_k = self._k_workspace[:, :total_ctx] tmp_v = self._v_workspace[:, :total_ctx] cache_k, cache_v = _fused_norm_rope_stacked( proj_out, self.k_norm_weights, self.eps_values, cos_sin_cache, positions, self.num_kv_heads, self.head_dim, self.rotary_dim, k_out=tmp_k, v_out=tmp_v, ) for layer_idx in range(self.n_layers): write_layer_kv(layer_idx, cache_k[layer_idx], cache_v[layer_idx])