diff --git a/python/tokenspeed/runtime/layers/attention/backends/hybrid_kda.py b/python/tokenspeed/runtime/layers/attention/backends/hybrid_kda.py index a2c225feb0..d85e6e5959 100644 --- a/python/tokenspeed/runtime/layers/attention/backends/hybrid_kda.py +++ b/python/tokenspeed/runtime/layers/attention/backends/hybrid_kda.py @@ -27,6 +27,7 @@ from typing import TYPE_CHECKING import torch +from tokenspeed_kernel.ops.activation.triton import rmsnorm_gated_sigmoid from tokenspeed_kernel.ops.attention import ( kda_paged_decode, kda_paged_prefill, @@ -134,29 +135,50 @@ def _decode( attn_tp_size: int, head_v_dim: int, lower_bound: float | None, + output_gate: torch.Tensor | None, + norm_weight: torch.Tensor | None, + norm_eps: float | None, ) -> torch.Tensor | None: - + if output_gate is not None and (norm_weight is None or norm_eps is None): + raise ValueError( + "norm_weight and norm_eps are required with a KDA output gate" + ) if f_a_out is None: return None - else: - num_value_heads = value_dim // attn_tp_size // head_v_dim - return try_kda_fused_paged_decode( - mixed_qkv, - conv_weights, - conv_states, - f_a_out, - f_b_weight, - beta_raw, - A_log, - dt_bias, - state_pool=ssm_states, - read_indices=read_indices, - write_indices=write_indices, - num_heads=num_value_heads, - head_dim=head_v_dim, - cu_seqlens=self.forward_metadata.query_start_loc, - lower_bound=lower_bound, - ) + + num_value_heads = value_dim // attn_tp_size // head_v_dim + result = try_kda_fused_paged_decode( + mixed_qkv, + conv_weights, + conv_states, + f_a_out, + f_b_weight, + beta_raw, + A_log, + dt_bias, + state_pool=ssm_states, + read_indices=read_indices, + write_indices=write_indices, + num_heads=num_value_heads, + head_dim=head_v_dim, + cu_seqlens=self.forward_metadata.query_start_loc, + lower_bound=lower_bound, + output_gate=output_gate, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + if result is None: + return None + if result.output_norm_applied or output_gate is None: + return result.out + return rmsnorm_gated_sigmoid( + result.out.reshape(-1, num_value_heads * head_v_dim).contiguous(), + output_gate.contiguous(), + norm_weight, + norm_eps, + num_value_heads, + head_v_dim, + ).view(1, -1, num_value_heads, head_v_dim) @override def _decode_scan( @@ -177,8 +199,10 @@ def _decode_scan( f_b_weight: torch.Tensor | None, beta_raw: torch.Tensor | None, lower_bound: float | None, + output_gate: torch.Tensor | None, + norm_weight: torch.Tensor | None, + norm_eps: float | None, ) -> torch.Tensor: - seq_len = query.shape[0] num_heads = query.shape[2] head_k_dim = query.shape[3] @@ -193,7 +217,7 @@ def _decode_scan( g_kda = g_raw.view(1, seq_len, num_value_heads, head_k_dim) beta_kda = beta_raw.view(1, seq_len, num_value_heads) - return kda_paged_decode( + core_attn_out = kda_paged_decode( query, key, value, @@ -206,7 +230,17 @@ def _decode_scan( write_indices=write_indices, cu_seqlens=query_start_loc, lower_bound=lower_bound, - ).squeeze(0) + ) + if output_gate is not None: + core_attn_out = rmsnorm_gated_sigmoid( + core_attn_out.reshape(-1, num_value_heads * head_v_dim).contiguous(), + output_gate.contiguous(), + norm_weight, + norm_eps, + num_value_heads, + head_v_dim, + ).view(1, -1, num_value_heads, head_v_dim) + return core_attn_out.squeeze(0) @override def _verify( diff --git a/python/tokenspeed/runtime/layers/attention/backends/hybrid_linear_attn.py b/python/tokenspeed/runtime/layers/attention/backends/hybrid_linear_attn.py index dacd87547d..5f3cf30c49 100644 --- a/python/tokenspeed/runtime/layers/attention/backends/hybrid_linear_attn.py +++ b/python/tokenspeed/runtime/layers/attention/backends/hybrid_linear_attn.py @@ -1292,6 +1292,9 @@ def forward_decode( f_b_weight = kwargs.get("f_b_weight") g_raw = kwargs.get("g_raw") beta_raw = kwargs.get("beta_raw") + output_gate = kwargs.get("output_gate") + norm_weight = kwargs.get("norm_weight") + norm_eps = kwargs.get("norm_eps") gate_lower_bound = kwargs.get("lower_bound") A_log = kwargs["A_log"] dt_bias = kwargs["dt_bias"] @@ -1320,6 +1323,9 @@ def forward_decode( attn_tp_size=attn_tp_size, head_v_dim=head_v_dim, lower_bound=gate_lower_bound, + output_gate=output_gate, + norm_weight=norm_weight, + norm_eps=norm_eps, ) if fused_out is not None: return fused_out @@ -1368,6 +1374,9 @@ def forward_decode( f_b_weight=f_b_weight, beta_raw=beta_raw, lower_bound=gate_lower_bound, + output_gate=output_gate, + norm_weight=norm_weight, + norm_eps=norm_eps, ) def _decode( @@ -1388,6 +1397,9 @@ def _decode( attn_tp_size: int, head_v_dim: int, lower_bound: float | None, + output_gate: torch.Tensor | None, + norm_weight: torch.Tensor | None, + norm_eps: float | None, ) -> torch.Tensor | None: """Whole-step decode attempt; ``None`` falls through to the shared flow. @@ -1415,6 +1427,9 @@ def _decode( attn_tp_size: Attention tensor-parallel size. head_v_dim: Value head dimension. lower_bound: KDA decay clamp. + output_gate: Optional KDA gated-norm logits. + norm_weight: Optional KDA output RMSNorm weight. + norm_eps: Optional KDA output RMSNorm epsilon. Returns: The layer output when a fused kernel ran, else None. @@ -1439,6 +1454,9 @@ def _decode_scan( f_b_weight: torch.Tensor | None, beta_raw: torch.Tensor | None, lower_bound: float | None, + output_gate: torch.Tensor | None, + norm_weight: torch.Tensor | None, + norm_eps: float | None, ) -> torch.Tensor: """Single-token recurrent scan over the split, conv'd projections. @@ -1464,6 +1482,9 @@ def _decode_scan( f_b_weight: KDA second gate projection. beta_raw: KDA raw per-head beta logits. lower_bound: KDA decay clamp. + output_gate: Optional KDA gated-norm logits. + norm_weight: Optional KDA output RMSNorm weight. + norm_eps: Optional KDA output RMSNorm epsilon. Returns: ``[1, B, Hv, V]`` layer output. diff --git a/python/tokenspeed/runtime/models/kimi_k3.py b/python/tokenspeed/runtime/models/kimi_k3.py index 51f3814682..8ecacbd317 100644 --- a/python/tokenspeed/runtime/models/kimi_k3.py +++ b/python/tokenspeed/runtime/models/kimi_k3.py @@ -918,6 +918,7 @@ def forward( # plain GEMV on the prefill path. # Fused [3*proj, k] conv kernel bank, built once in post_load_weights. conv_weights = self.conv_weights + fuse_decode_output_norm = ctx.forward_mode.is_decode() and num_tokens == ctx.bs core_out = ctx.attn_backend.forward( q=None, @@ -943,19 +944,25 @@ def forward( A_log=self.A_log, dt_bias=self.dt_bias, lower_bound=self.gate_lower_bound, + output_gate=out_gate if fuse_decode_output_norm else None, + norm_weight=self.o_norm.weight if fuse_decode_output_norm else None, + norm_eps=self.o_norm.variance_epsilon if fuse_decode_output_norm else None, layer_id=self.layer_id, seq_len=num_tokens, ) - # Per-head gated RMSNorm + sigmoid output gate in one kernel. - core_out = rmsnorm_gated_sigmoid( - core_out.reshape(num_tokens, hn * hd).contiguous(), - out_gate.contiguous(), - self.o_norm.weight, - self.o_norm.variance_epsilon, - hn, - hd, - ) + core_out = core_out.reshape(num_tokens, hn * hd) + if not fuse_decode_output_norm: + # Decode kernels may fuse this epilogue; prefill retains the shared + # per-head norm implementation. + core_out = rmsnorm_gated_sigmoid( + core_out.contiguous(), + out_gate.contiguous(), + self.o_norm.weight, + self.o_norm.variance_epsilon, + hn, + hd, + ) output, _ = self.o_proj(core_out) return output diff --git a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/attention/kda/decode.py b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/attention/kda/decode.py index 4d854beee1..bcf0a20a9d 100644 --- a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/attention/kda/decode.py +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/gfx950/attention/kda/decode.py @@ -180,6 +180,242 @@ def _kda_recurrent_decode_kernel( ) +@gluon.jit +def _kda_fused_decode_kernel( + mixed_qkv, + conv_weights, + conv_states, + raw_g, + beta_logits, + output_gate, + norm_weight, + state_pool, + read_indices, + write_indices, + output, + cu_seqlens, + a_log, + dt_bias, + H: gl.constexpr, + D: gl.constexpr, + MIXED_ROW_STRIDE: gl.constexpr, + CONV_WEIGHT_ROW_STRIDE: gl.constexpr, + CONV_WEIGHT_COL_STRIDE: gl.constexpr, + CONV_PAGE_STRIDE: gl.constexpr, + CONV_CHANNEL_STRIDE: gl.constexpr, + CONV_HISTORY_STRIDE: gl.constexpr, + GATE_ROW_STRIDE: gl.constexpr, + BETA_ROW_STRIDE: gl.constexpr, + OUTPUT_GATE_ROW_STRIDE: gl.constexpr, + STATE_PAGE_STRIDE: gl.constexpr, + NUM_SLOTS: gl.constexpr, + HAS_LOWER_BOUND: gl.constexpr, + LOWER_BOUND: gl.constexpr, + NORM_EPS: gl.constexpr, +): + """Fuse the K3 decode convolution, recurrence, and gated RMSNorm.""" + sequence_head = gl.program_id(0) + sequence_idx = sequence_head // H + head_idx = sequence_head % H + + state_layout: gl.constexpr = gl.BlockedLayout( + [8, 8], + [16, 4], + [1, 4], + [1, 0], + ) + key_layout: gl.constexpr = gl.SliceLayout(1, state_layout) + value_layout: gl.constexpr = gl.SliceLayout(0, state_layout) + key_offsets = gl.arange(0, D, layout=key_layout) + value_offsets = gl.arange(0, D, layout=value_layout) + + begin = gl.load(cu_seqlens + sequence_idx) + end = gl.load(cu_seqlens + sequence_idx + 1) + output_offsets = (sequence_idx * H + head_idx) * D + value_offsets + if begin == end: + gl.store(output + output_offsets, 0.0) + return + + read_idx = gl.load(read_indices + sequence_idx) + write_idx = gl.load(write_indices + sequence_idx) + valid_read = (read_idx >= 0) & (read_idx < NUM_SLOTS) + if not valid_read: + gl.store(output + output_offsets, 0.0) + return + read_page_offset = read_idx.to(gl.int64) + + token_idx = begin + projection_width: gl.constexpr = H * D + + q_channel = head_idx * D + key_offsets + k_channel = projection_width + q_channel + v_channel = 2 * projection_width + head_idx * D + value_offsets + q_input = gl.load(mixed_qkv + token_idx * MIXED_ROW_STRIDE + q_channel).to( + gl.float32 + ) + k_input = gl.load(mixed_qkv + token_idx * MIXED_ROW_STRIDE + k_channel).to( + gl.float32 + ) + v_input = gl.load(mixed_qkv + token_idx * MIXED_ROW_STRIDE + v_channel).to( + gl.float32 + ) + + q_history0 = gl.load( + conv_states + + read_page_offset * CONV_PAGE_STRIDE + + q_channel * CONV_CHANNEL_STRIDE + ).to(gl.float32) + q_history1 = gl.load( + conv_states + + read_page_offset * CONV_PAGE_STRIDE + + q_channel * CONV_CHANNEL_STRIDE + + CONV_HISTORY_STRIDE + ).to(gl.float32) + q_history2 = gl.load( + conv_states + + read_page_offset * CONV_PAGE_STRIDE + + q_channel * CONV_CHANNEL_STRIDE + + 2 * CONV_HISTORY_STRIDE + ).to(gl.float32) + k_history0 = gl.load( + conv_states + + read_page_offset * CONV_PAGE_STRIDE + + k_channel * CONV_CHANNEL_STRIDE + ).to(gl.float32) + k_history1 = gl.load( + conv_states + + read_page_offset * CONV_PAGE_STRIDE + + k_channel * CONV_CHANNEL_STRIDE + + CONV_HISTORY_STRIDE + ).to(gl.float32) + k_history2 = gl.load( + conv_states + + read_page_offset * CONV_PAGE_STRIDE + + k_channel * CONV_CHANNEL_STRIDE + + 2 * CONV_HISTORY_STRIDE + ).to(gl.float32) + v_history0 = gl.load( + conv_states + + read_page_offset * CONV_PAGE_STRIDE + + v_channel * CONV_CHANNEL_STRIDE + ).to(gl.float32) + v_history1 = gl.load( + conv_states + + read_page_offset * CONV_PAGE_STRIDE + + v_channel * CONV_CHANNEL_STRIDE + + CONV_HISTORY_STRIDE + ).to(gl.float32) + v_history2 = gl.load( + conv_states + + read_page_offset * CONV_PAGE_STRIDE + + v_channel * CONV_CHANNEL_STRIDE + + 2 * CONV_HISTORY_STRIDE + ).to(gl.float32) + + q_weight_base = conv_weights + q_channel * CONV_WEIGHT_ROW_STRIDE + k_weight_base = conv_weights + k_channel * CONV_WEIGHT_ROW_STRIDE + v_weight_base = conv_weights + v_channel * CONV_WEIGHT_ROW_STRIDE + q_value = ( + q_history0 * gl.load(q_weight_base) + + q_history1 * gl.load(q_weight_base + CONV_WEIGHT_COL_STRIDE) + + q_history2 * gl.load(q_weight_base + 2 * CONV_WEIGHT_COL_STRIDE) + + q_input * gl.load(q_weight_base + 3 * CONV_WEIGHT_COL_STRIDE) + ).to(gl.float32) + k_value = ( + k_history0 * gl.load(k_weight_base) + + k_history1 * gl.load(k_weight_base + CONV_WEIGHT_COL_STRIDE) + + k_history2 * gl.load(k_weight_base + 2 * CONV_WEIGHT_COL_STRIDE) + + k_input * gl.load(k_weight_base + 3 * CONV_WEIGHT_COL_STRIDE) + ).to(gl.float32) + v_value = ( + v_history0 * gl.load(v_weight_base) + + v_history1 * gl.load(v_weight_base + CONV_WEIGHT_COL_STRIDE) + + v_history2 * gl.load(v_weight_base + 2 * CONV_WEIGHT_COL_STRIDE) + + v_input * gl.load(v_weight_base + 3 * CONV_WEIGHT_COL_STRIDE) + ).to(gl.float32) + q_value *= 1.0 / (1.0 + gl.exp(-q_value)) + k_value *= 1.0 / (1.0 + gl.exp(-k_value)) + v_value *= 1.0 / (1.0 + gl.exp(-v_value)) + + valid_write = (write_idx >= 0) & (write_idx < NUM_SLOTS) + safe_write_idx = gl.where(valid_write, write_idx, 0) + write_page_offset = safe_write_idx.to(gl.int64) + q_write_base = ( + conv_states + + write_page_offset * CONV_PAGE_STRIDE + + q_channel * CONV_CHANNEL_STRIDE + ) + k_write_base = ( + conv_states + + write_page_offset * CONV_PAGE_STRIDE + + k_channel * CONV_CHANNEL_STRIDE + ) + v_write_base = ( + conv_states + + write_page_offset * CONV_PAGE_STRIDE + + v_channel * CONV_CHANNEL_STRIDE + ) + gl.store(q_write_base, q_history1, mask=valid_write) + gl.store(q_write_base + CONV_HISTORY_STRIDE, q_history2, mask=valid_write) + gl.store(q_write_base + 2 * CONV_HISTORY_STRIDE, q_input, mask=valid_write) + gl.store(k_write_base, k_history1, mask=valid_write) + gl.store(k_write_base + CONV_HISTORY_STRIDE, k_history2, mask=valid_write) + gl.store(k_write_base + 2 * CONV_HISTORY_STRIDE, k_input, mask=valid_write) + gl.store(v_write_base, v_history1, mask=valid_write) + gl.store(v_write_base + CONV_HISTORY_STRIDE, v_history2, mask=valid_write) + gl.store(v_write_base + 2 * CONV_HISTORY_STRIDE, v_input, mask=valid_write) + gate_value = gl.load( + raw_g + token_idx * GATE_ROW_STRIDE + head_idx * D + key_offsets + ).to(gl.float32) + gate_value += gl.load(dt_bias + head_idx * D + key_offsets).to(gl.float32) + a_value = gl.exp(gl.load(a_log + head_idx).to(gl.float32)) + if HAS_LOWER_BOUND: + log_decay = LOWER_BOUND / (1.0 + gl.exp(-(a_value * gate_value))) + else: + softplus = gl.maximum(gate_value, 0.0) + gl.log( + 1.0 + gl.exp(-gl.abs(gate_value)) + ) + log_decay = -a_value * softplus + beta_value = gl.load(beta_logits + token_idx * BETA_ROW_STRIDE + head_idx).to( + gl.float32 + ) + beta_value = 1.0 / (1.0 + gl.exp(-beta_value)) + + q_value *= gl.rsqrt(gl.sum(q_value * q_value, axis=0) + 1e-6) * (D**-0.5) + k_value *= gl.rsqrt(gl.sum(k_value * k_value, axis=0) + 1e-6) + state_offsets = key_offsets[:, None] * D + value_offsets[None, :] + read_base = read_page_offset * STATE_PAGE_STRIDE + head_idx * D * D + running = cdna4.buffer_load( + state_pool + read_base, + state_offsets.to(gl.int32), + cache=".cs", + ).to(gl.float32) + running *= gl.exp(log_decay)[:, None] + prediction = gl.sum(running * k_value[:, None], axis=0) + prior_output = gl.sum(running * q_value[:, None], axis=0) + key_query = gl.sum(k_value * q_value, axis=0) + delta = beta_value * (v_value - prediction) + running += k_value[:, None] * delta[None, :] + out_value = prior_output + delta * key_query + + write_base = write_page_offset * STATE_PAGE_STRIDE + head_idx * D * D + cdna4.buffer_store( + running, + state_pool + write_base, + state_offsets.to(gl.int32), + mask=valid_write, + cache=".cs", + ) + + inverse_rms = gl.rsqrt(gl.sum(out_value * out_value, axis=0) / D + NORM_EPS) + gate = gl.load( + output_gate + token_idx * OUTPUT_GATE_ROW_STRIDE + head_idx * D + value_offsets + ).to(gl.float32) + weight = gl.load(norm_weight + value_offsets).to(gl.float32) + out_value *= inverse_rms * weight * (1.0 / (1.0 + gl.exp(-gate))) + gl.store(output + output_offsets, out_value.to(output.dtype.element_ty)) + + def gluon_kda_recurrent_decode_gfx950( q: torch.Tensor, k: torch.Tensor, @@ -287,4 +523,181 @@ def gluon_kda_recurrent_decode_gfx950( return output -__all__ = ["gluon_kda_recurrent_decode_gfx950"] +def gluon_kda_fused_decode_gfx950( + mixed_qkv: torch.Tensor, + conv_weights: torch.Tensor, + conv_states: torch.Tensor, + raw_g: torch.Tensor, + beta_logits: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + output_gate: torch.Tensor, + norm_weight: torch.Tensor, + norm_eps: float, + *, + state_pool: torch.Tensor, + read_indices: torch.Tensor, + write_indices: torch.Tensor, + num_heads: int, + head_dim: int, + cu_seqlens: torch.Tensor, + lower_bound: float | None, +) -> torch.Tensor: + """Run fused single-token K3 decode against paged convolution/KDA state. + + Args: + mixed_qkv: BF16 pre-convolution Q/K/V rows with shape + ``[batch, 3 * num_heads * head_dim]``. + conv_weights: Contiguous four-tap depthwise convolution weights with + shape ``[3 * num_heads * head_dim, 4]``. + conv_states: Mutable BF16 convolution state pool with shape + ``[pages, 3 * num_heads * head_dim, 3]``. + raw_g: Projected decay-gate values with shape + ``[batch, num_heads * head_dim]``. + beta_logits: Per-token, per-head update logits with shape + ``[batch, num_heads]``. + A_log: Per-head FP32 decay parameters with shape ``[num_heads]``. + dt_bias: Per-head, per-channel FP32 decay bias with shape + ``[num_heads * head_dim]``. + output_gate: Gated-RMSNorm logits with shape + ``[batch, num_heads * head_dim]``. + norm_weight: Gated-RMSNorm weights with shape ``[head_dim]``. + norm_eps: Gated-RMSNorm epsilon. + state_pool: Mutable FP32 recurrent state in canonical K-major layout, + with shape ``[pages, num_heads, head_dim, head_dim]``. + read_indices: Source state page per batch row. Negative entries mark + graph-padding rows. + write_indices: Destination state page per batch row. Negative entries + suppress state updates. + num_heads: Number of local KDA heads; this specialization requires 12. + head_dim: Per-head key/value width; this specialization requires 128. + cu_seqlens: Packed row boundaries with shape ``[batch + 1]``. Active + rows contain one token and graph-padding rows contain zero tokens. + lower_bound: Optional safe lower bound for the log-decay gate. + + Returns: + Gated-RMSNorm KDA output with shape + ``[1, batch, num_heads, head_dim]`` and ``mixed_qkv`` dtype. The + convolution and recurrent states at valid ``write_indices`` are + updated as part of the call. + """ + tensors = ( + mixed_qkv, + conv_weights, + conv_states, + raw_g, + beta_logits, + A_log, + dt_bias, + output_gate, + norm_weight, + state_pool, + read_indices, + write_indices, + cu_seqlens, + ) + if not all(tensor.is_cuda for tensor in tensors): + raise ValueError("gfx950 fused KDA decode requires GPU tensors") + if num_heads != 12 or head_dim != 128: + raise ValueError("gfx950 fused KDA decode requires 12 heads of width 128") + + tokens = read_indices.numel() + projection_width = num_heads * head_dim + if mixed_qkv.ndim != 2 or mixed_qkv.shape != (tokens, 3 * projection_width): + raise ValueError("mixed_qkv must have shape [batch, 3 * heads * head_dim]") + if mixed_qkv.stride(1) != 1: + raise ValueError("mixed_qkv channels must be contiguous") + if ( + conv_weights.shape != (3 * projection_width, 4) + or not conv_weights.is_contiguous() + ): + raise ValueError("conv_weights must be contiguous [3 * heads * head_dim, 4]") + if conv_states.ndim != 3 or conv_states.shape[1:] != (3 * projection_width, 3): + raise ValueError("conv_states must have shape [pages, 3 * heads * head_dim, 3]") + if conv_states.stride()[1:] != (3, 1): + raise ValueError( + "conv_states inner channel/history dimensions must be contiguous" + ) + if raw_g.shape != (tokens, projection_width) or raw_g.stride(1) != 1: + raise ValueError("raw_g must have shape [batch, heads * head_dim]") + if beta_logits.shape != (tokens, num_heads) or beta_logits.stride(1) != 1: + raise ValueError("beta_logits must have shape [batch, heads]") + if output_gate.shape != (tokens, projection_width) or output_gate.stride(1) != 1: + raise ValueError("output_gate must have shape [batch, heads * head_dim]") + if A_log.shape != (num_heads,) or not A_log.is_contiguous(): + raise ValueError("A_log must have shape [heads]") + if dt_bias.shape != (projection_width,) or not dt_bias.is_contiguous(): + raise ValueError("dt_bias must have shape [heads * head_dim]") + if norm_weight.shape != (head_dim,) or not norm_weight.is_contiguous(): + raise ValueError("norm_weight must have shape [head_dim]") + if state_pool.ndim != 4 or state_pool.shape[1:] != ( + num_heads, + head_dim, + head_dim, + ): + raise ValueError( + "state_pool must have shape [pages, heads, head_dim, head_dim]" + ) + if state_pool.stride()[1:] != (head_dim * head_dim, head_dim, 1): + raise ValueError("state_pool inner dimensions must be contiguous") + if conv_states.shape[0] != state_pool.shape[0]: + raise ValueError( + "convolution and recurrent state pools must have equal capacity" + ) + if read_indices.shape != (tokens,) or write_indices.shape != (tokens,): + raise ValueError("read_indices and write_indices must match the decode batch") + if read_indices.dtype != torch.int32 or write_indices.dtype != torch.int32: + raise ValueError("read_indices and write_indices must be int32") + if not read_indices.is_contiguous() or not write_indices.is_contiguous(): + raise ValueError("read_indices and write_indices must be contiguous") + if cu_seqlens.shape != (tokens + 1,) or cu_seqlens.dtype != torch.int32: + raise ValueError("cu_seqlens must be an int32 boundary vector") + if not cu_seqlens.is_contiguous(): + raise ValueError("cu_seqlens must be contiguous") + + output = torch.empty( + (1, tokens, num_heads, head_dim), + dtype=mixed_qkv.dtype, + device=mixed_qkv.device, + ) + _kda_fused_decode_kernel[(tokens * num_heads,)]( + mixed_qkv, + conv_weights, + conv_states, + raw_g, + beta_logits, + output_gate, + norm_weight, + state_pool, + read_indices, + write_indices, + output, + cu_seqlens, + A_log, + dt_bias, + H=num_heads, + D=head_dim, + MIXED_ROW_STRIDE=mixed_qkv.stride(0), + CONV_WEIGHT_ROW_STRIDE=conv_weights.stride(0), + CONV_WEIGHT_COL_STRIDE=conv_weights.stride(1), + CONV_PAGE_STRIDE=conv_states.stride(0), + CONV_CHANNEL_STRIDE=conv_states.stride(1), + CONV_HISTORY_STRIDE=conv_states.stride(2), + GATE_ROW_STRIDE=raw_g.stride(0), + BETA_ROW_STRIDE=beta_logits.stride(0), + OUTPUT_GATE_ROW_STRIDE=output_gate.stride(0), + STATE_PAGE_STRIDE=state_pool.stride(0), + NUM_SLOTS=state_pool.shape[0], + HAS_LOWER_BOUND=lower_bound is not None, + LOWER_BOUND=0.0 if lower_bound is None else lower_bound, + NORM_EPS=norm_eps, + num_warps=4, + num_stages=2, + ) + return output + + +__all__ = [ + "gluon_kda_fused_decode_gfx950", + "gluon_kda_recurrent_decode_gfx950", +] diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/__init__.py index 8394afded6..d23ec92812 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/__init__.py @@ -38,7 +38,10 @@ GdnCheckpointLayout, GdnChunkPrefillResult, ) -from tokenspeed_kernel.ops.attention.kda_utils import KdaPrefillResult +from tokenspeed_kernel.ops.attention.kda_utils import ( + KdaFusedDecodeResult, + KdaPrefillResult, +) from tokenspeed_kernel.platform import current_platform from tokenspeed_kernel.profiling import ShapeCapture, kernel_scope from tokenspeed_kernel.registry import KernelRegistry, Priority @@ -132,6 +135,7 @@ def _blockscaled_signature_and_scales( "kda_paged_decode", "try_kda_fused_paged_decode", "try_kda_fused_paged_verify", + "KdaFusedDecodeResult", "KdaPrefillResult", "GdnCheckpointLayout", "GdnChunkPrefillResult", @@ -1326,14 +1330,28 @@ def try_kda_fused_paged_decode( head_dim: int, cu_seqlens: torch.Tensor, lower_bound: float | None = -5.0, + output_gate: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float | None = None, override: str | None = None, solution: str | None = None, -) -> torch.Tensor | None: +) -> KdaFusedDecodeResult | None: """Try a registered pre-convolution KDA decode fusion. + ``output_gate``, ``norm_weight``, and ``norm_eps`` request a fused gated + RMSNorm epilogue. If the selected backend only supports the original core + fusion, the returned result reports that the caller must apply the + epilogue. + Returns ``None`` only when no implementation supports the current - platform. Invalid inputs and execution failures remain visible. + platform. Otherwise, returns the output and whether output normalization + was applied. Invalid inputs and execution failures remain visible. """ + if (output_gate is None) != (norm_weight is None): + raise ValueError("output_gate and norm_weight must be provided together") + if output_gate is not None and norm_eps is None: + raise ValueError("norm_eps is required with fused KDA output normalization") + signature = _attention_format_signature( q=mixed_qkv, k=mixed_qkv, @@ -1344,13 +1362,49 @@ def try_kda_fused_paged_decode( "attention", "kda_fused_paged_decode", signature, - traits={"paged_state": True}, + traits={ + "paged_state": True, + "fused_output_norm": output_gate is not None, + "num_heads": num_heads, + "head_dim": head_dim, + "conv_kernel_size": conv_weights.shape[-1], + }, solution=solution, override=override, ) except NoKernelFoundError: - return None - return kernel( + if output_gate is None: + return None + try: + kernel = select_kernel( + "attention", + "kda_fused_paged_decode", + signature, + traits={ + "paged_state": True, + "fused_output_norm": False, + "num_heads": num_heads, + "head_dim": head_dim, + "conv_kernel_size": conv_weights.shape[-1], + }, + solution=solution, + override=override, + ) + except NoKernelFoundError: + return None + + selected_spec = KernelRegistry.get().get_by_name(kernel.name) + output_norm_applied = ( + output_gate is not None + and selected_spec is not None + and spec_matches_traits( + selected_spec, + {"fused_output_norm": True}, + require_all_traits=True, + ) + ) + + out = kernel( mixed_qkv=mixed_qkv, conv_weights=conv_weights, conv_states=conv_states, @@ -1366,7 +1420,11 @@ def try_kda_fused_paged_decode( head_dim=head_dim, cu_seqlens=cu_seqlens, lower_bound=lower_bound, + output_gate=output_gate if output_norm_applied else None, + norm_weight=norm_weight if output_norm_applied else None, + norm_eps=norm_eps if output_norm_applied else None, ) + return KdaFusedDecodeResult(out=out, output_norm_applied=output_norm_applied) def try_kda_fused_paged_verify( diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py index f222943b83..226b0d2ee4 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py @@ -52,6 +52,9 @@ from tokenspeed_kernel_amd.ops.gfx950.attention.dsa.sparse_mla import ( gluon_dsa_prefill_topk_fp8_gfx950 as _dsa_prefill_topk_impl, ) + from tokenspeed_kernel_amd.ops.gfx950.attention.kda.decode import ( + gluon_kda_fused_decode_gfx950 as _kda_fused_decode_impl, + ) from tokenspeed_kernel_amd.ops.gfx950.attention.kda.decode import ( gluon_kda_recurrent_decode_gfx950 as _kda_decode_impl, ) @@ -221,6 +224,76 @@ def gluon_kda_paged_decode_gfx950(**kwargs): """Run specialized gfx950 KDA decode against the canonical K-major pool.""" return _kda_decode_impl(**kwargs) + @register_kernel( + "attention", + "kda_fused_paged_decode", + name="gluon_kda_fused_paged_decode_gfx950", + solution="gluon", + capability=CapabilityRequirement( + min_arch_version=ArchVersion(9, 5), + max_arch_version=ArchVersion(9, 5), + vendors=frozenset({"amd"}), + ), + signatures=format_signatures( + ("q", "k", "v"), + "dense", + {torch.bfloat16}, + ), + priority=Priority.SPECIALIZED, + traits={ + "paged_state": frozenset({True}), + "fused_output_norm": frozenset({True}), + "num_heads": frozenset({12}), + "head_dim": frozenset({128}), + "conv_kernel_size": frozenset({4}), + }, + tags={"amd", "gfx950", "paged_cache", "cuda_graph", "fusion"}, + ) + def gluon_kda_fused_paged_decode_gfx950( + mixed_qkv: torch.Tensor, + conv_weights: torch.Tensor, + conv_states: torch.Tensor, + f_a_out: torch.Tensor, + f_b_weight: torch.Tensor, + beta_logits: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + *, + state_pool: torch.Tensor, + read_indices: torch.Tensor, + write_indices: torch.Tensor, + num_heads: int, + head_dim: int, + cu_seqlens: torch.Tensor, + lower_bound: float | None, + output_gate: torch.Tensor | None, + norm_weight: torch.Tensor | None, + norm_eps: float | None, + ): + """Run the decay projection and fused gfx950 KDA decode epilogue.""" + if output_gate is None or norm_weight is None or norm_eps is None: + raise ValueError("gfx950 fused KDA decode requires output normalization") + raw_g = torch.nn.functional.linear(f_a_out, f_b_weight) + return _kda_fused_decode_impl( + mixed_qkv=mixed_qkv, + conv_weights=conv_weights, + conv_states=conv_states, + raw_g=raw_g, + beta_logits=beta_logits, + A_log=A_log, + dt_bias=dt_bias, + output_gate=output_gate, + norm_weight=norm_weight, + norm_eps=norm_eps, + state_pool=state_pool, + read_indices=read_indices, + write_indices=write_indices, + num_heads=num_heads, + head_dim=head_dim, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + ) + @register_kernel( "attention", "kda_paged_decode", diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/kda_utils.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/kda_utils.py index bccc2a1690..d95b8ddd2e 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/kda_utils.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/kda_utils.py @@ -38,3 +38,17 @@ class KdaPrefillResult: out: torch.Tensor final_state: torch.Tensor + + +@dataclass(frozen=True) +class KdaFusedDecodeResult: + """Result from an optional pre-convolution KDA decode fusion. + + Attributes: + out: Packed decode output ``[1, batch, heads, value_dim]``. + output_norm_applied: Whether the selected kernel applied the output + gate and RMSNorm, so the caller must not apply them again. + """ + + out: torch.Tensor + output_norm_applied: bool diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/triton/kda_dispatch.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/triton/kda_dispatch.py index 1995a4fa9f..5a17edc503 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/triton/kda_dispatch.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/triton/kda_dispatch.py @@ -30,7 +30,10 @@ capability=CapabilityRequirement(vendors=frozenset({"nvidia"})), signatures=_DENSE_HALF_SIGNATURES, priority=Priority.SPECIALIZED, - traits={"paged_state": frozenset({True})}, + traits={ + "paged_state": frozenset({True}), + "fused_output_norm": frozenset({False}), + }, tags={"nvidia", "paged_cache", "cuda_graph", "fusion"}, ) def triton_nvidia_kda_fused_paged_decode( @@ -50,8 +53,12 @@ def triton_nvidia_kda_fused_paged_decode( head_dim: int, cu_seqlens: torch.Tensor, lower_bound: float | None, + output_gate: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float | None = None, ) -> torch.Tensor: """Adapt dev's NVIDIA conv/GEMV/recurrent megafusion.""" + del output_gate, norm_weight, norm_eps from tokenspeed_kernel.thirdparty.triton.fla_kda_recurrent import ( fused_recurrent_kda_megafuse, ) diff --git a/tokenspeed-kernel/test/ops/test_kda_recurrent.py b/tokenspeed-kernel/test/ops/test_kda_recurrent.py index d5976ee387..4158f7ad47 100644 --- a/tokenspeed-kernel/test/ops/test_kda_recurrent.py +++ b/tokenspeed-kernel/test/ops/test_kda_recurrent.py @@ -6,12 +6,15 @@ import torch from kimi3_reference import kda_gate from kimi3_reference import kda_recurrent as reference_kda_recurrent +from tokenspeed_kernel.ops import attention as attention_ops from tokenspeed_kernel.ops.attention import ( _attention_format_signature, kda_paged_decode, kda_paged_prefill, + try_kda_fused_paged_decode, ) from tokenspeed_kernel.platform import current_platform +from tokenspeed_kernel.registry import KernelRegistry from tokenspeed_kernel.selection import NoKernelFoundError, select_kernel @@ -234,6 +237,227 @@ def test_kda_paged_decode_graph_padding_and_page_stride() -> None: ) +def test_kda_fused_paged_decode_matches_reference() -> None: + """The K3 megafusion preserves state paging and its fused norm epilogue.""" + if not current_platform().is_cdna4: + pytest.skip("gfx950 KDA fusion test") + + torch.manual_seed(31) + batch, active, heads, head_dim, pages = 4, 2, 12, 128, 6 + projection_width = heads * head_dim + mixed_qkv = torch.randn( + batch, + 3 * projection_width, + device="cuda", + dtype=torch.bfloat16, + ) + conv_weights = 0.1 * torch.randn( + 3 * projection_width, + 4, + device="cuda", + dtype=torch.bfloat16, + ) + conv_states = 0.1 * torch.randn( + pages, + 3 * projection_width, + 3, + device="cuda", + dtype=torch.bfloat16, + ) + initial_conv_states = conv_states.clone() + expected_conv_states = conv_states.clone() + f_a_out = torch.randn( + batch, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + f_b_weight = 0.1 * torch.randn( + projection_width, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + beta_logits = torch.randn( + batch, + heads, + device="cuda", + dtype=torch.bfloat16, + ) + a_log = torch.randn(heads, device="cuda", dtype=torch.float32) + dt_bias = torch.randn(projection_width, device="cuda", dtype=torch.float32) + output_gate = torch.randn( + batch, + projection_width, + device="cuda", + dtype=torch.bfloat16, + ) + norm_weight = torch.randn(head_dim, device="cuda", dtype=torch.bfloat16) + norm_eps = 1e-6 + state_pool = 0.01 * torch.randn( + pages, + heads, + head_dim, + head_dim, + device="cuda", + dtype=torch.float32, + ) + initial_state_pool = state_pool.clone() + expected_state_pool = state_pool.clone() + read_indices = torch.tensor([0, 1, -1, -1], device="cuda", dtype=torch.int32) + write_indices = torch.tensor([2, 3, -1, -1], device="cuda", dtype=torch.int32) + cu_seqlens = torch.tensor([0, 1, 2, 2, 2], device="cuda", dtype=torch.int32) + + raw_g = torch.nn.functional.linear(f_a_out, f_b_weight) + expected_out = torch.zeros( + batch, + heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + for row in range(active): + read_idx = read_indices[row].item() + write_idx = write_indices[row].item() + current = mixed_qkv[row].view(3, heads, head_dim).float() + history = initial_conv_states[read_idx].view(3, heads, head_dim, 3).float() + weights = conv_weights.view(3, heads, head_dim, 4).float() + convolved = torch.nn.functional.silu( + (history * weights[..., :3]).sum(dim=-1) + current * weights[..., 3] + ) + q, k, v = convolved.unbind(dim=0) + core_out, final_state = reference_kda_recurrent( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + raw_g[row].view(1, heads, head_dim), + beta_logits[row].unsqueeze(0), + initial_state_pool[read_idx].transpose(-1, -2), + a_log, + dt_bias.view(heads, head_dim), + ) + core_out = core_out[0].float() + inverse_rms = torch.rsqrt(core_out.square().mean(dim=-1) + norm_eps) + expected_out[row] = ( + core_out + * inverse_rms[:, None] + * norm_weight.float()[None, :] + * torch.sigmoid(output_gate[row].view(heads, head_dim).float()) + ).to(torch.bfloat16) + expected_state_pool[write_idx] = final_state.transpose(-1, -2) + expected_conv_states[write_idx] = torch.stack( + (history[..., 1], history[..., 2], current), dim=-1 + ).reshape(3 * projection_width, 3) + + result = try_kda_fused_paged_decode( + mixed_qkv, + conv_weights, + conv_states, + f_a_out, + f_b_weight, + beta_logits, + a_log, + dt_bias, + state_pool=state_pool, + read_indices=read_indices, + write_indices=write_indices, + num_heads=heads, + head_dim=head_dim, + cu_seqlens=cu_seqlens, + output_gate=output_gate, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + + assert result is not None + assert result.output_norm_applied + torch.testing.assert_close( + result.out[0].float(), expected_out.float(), atol=5e-2, rtol=5e-2 + ) + torch.testing.assert_close(state_pool, expected_state_pool, atol=2e-4, rtol=2e-4) + torch.testing.assert_close(conv_states, expected_conv_states) + + +def test_kda_fused_decode_override_preserves_external_output_norm(monkeypatch) -> None: + """A core-only override must not claim that it applied output normalization.""" + kernel_name = "triton_nvidia_kda_fused_paged_decode" + spec = KernelRegistry.get().get_by_name(kernel_name) + assert spec is not None + assert spec.traits["fused_output_norm"] == frozenset({False}) + + captured_kwargs = {} + + class CoreOnlyKernel: + name = kernel_name + + def __call__(self, **kwargs): + captured_kwargs.update(kwargs) + return kwargs["mixed_qkv"] + + monkeypatch.setattr( + attention_ops, + "select_kernel", + lambda *_args, **_kwargs: CoreOnlyKernel(), + ) + tensor = torch.empty(1, dtype=torch.bfloat16) + result = try_kda_fused_paged_decode( + tensor, + tensor, + tensor, + tensor, + tensor, + tensor, + tensor, + tensor, + state_pool=tensor, + read_indices=tensor, + write_indices=tensor, + num_heads=12, + head_dim=128, + cu_seqlens=tensor, + output_gate=tensor, + norm_weight=tensor, + norm_eps=1e-6, + override=kernel_name, + ) + + assert result is not None + assert not result.output_norm_applied + assert captured_kwargs["output_gate"] is None + assert captured_kwargs["norm_weight"] is None + assert captured_kwargs["norm_eps"] is None + + +def test_kda_fused_decode_rejects_unsupported_conv_width() -> None: + """Unsupported convolution widths must fall back before kernel execution.""" + if not current_platform().is_cdna4: + pytest.skip("gfx950 KDA fusion dispatch test") + + tensor = torch.empty(1, dtype=torch.bfloat16) + conv_weights = torch.empty(1, 5, dtype=torch.bfloat16) + result = try_kda_fused_paged_decode( + tensor, + conv_weights, + tensor, + tensor, + tensor, + tensor, + tensor, + tensor, + state_pool=tensor, + read_indices=tensor, + write_indices=tensor, + num_heads=12, + head_dim=128, + cu_seqlens=tensor, + output_gate=tensor, + norm_weight=tensor, + norm_eps=1e-6, + ) + + assert result is None + + def test_kda_paged_decode_does_not_select_nvidia_kernel_on_amd() -> None: """The NVIDIA portable adapter remains outside the AMD dispatch surface.""" if not current_platform().is_amd: