# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """GLM-5.3-Flash KDA layer with separate convolutions and a bounded safe gate.""" import os from dataclasses import replace import torch from torch import nn from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import divide from vllm.forward_context import get_forward_context from vllm.model_executor.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, RowParallelLinear, ) from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateDtypeCalculator, MambaStateShapeCalculator, is_conv_state_dim_first, ) from vllm.model_executor.layers.mamba.ops.causal_conv1d import ( causal_conv1d_fn, causal_conv1d_update, ) from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( gather_initial_states, ) from vllm.model_executor.layers.mamba.ops.scatter_states import scatter_states from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader from vllm.model_executor.utils import ( maybe_disable_graph_partition, set_weight_attrs, ) from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig from vllm.v1.attention.backend import AttentionBackend from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata from vllm.v1.kv_cache_interface import KVCacheSpec from vllm.v1.worker.workspace import current_workspace_manager # Spec verify from one checkpoint state per request (recoverssm.py). _RECOVERSSM = os.environ.get("VLLM_GLM5NEXT_RECOVERSSM") == "1" if _RECOVERSSM: from vllm.models.glm5next.common.recoverssm import ( Glm5NextRecoverSSMBackend, kda_recoverssm_verify, ) if current_platform.is_rocm(): from vllm.models.glm5next.amd.ops.third_party.kda import ( chunk_kda_with_fused_gate, fused_recurrent_kda, ) else: from vllm.models.glm5next.nvidia.ops.third_party.kda import ( chunk_kda_with_fused_gate, fused_recurrent_kda, ) class _Glm5NextMergedColumnParallelLinear(MergedColumnParallelLinear): """Merged projection with multiple replicated output shards. Extends K3's ``_KimiGDNMergedColumnParallelLinear`` to support two replicated shards (f_a, g_a) instead of one. Pre-multiplies each replicated entry's output_size by tp_size so the per-rank shard divides back to the full size, and forces tp_rank=0 during weight loading for replicated shards. """ def __init__( self, input_size: int, output_sizes: list[int], replicated_shard_ids: tuple[int, ...], tp_size: int, **kwargs, ) -> None: self.replicated_shard_ids = set(replicated_shard_ids) output_sizes = output_sizes.copy() for sid in self.replicated_shard_ids: output_sizes[sid] *= tp_size super().__init__(input_size, output_sizes, **kwargs) def weight_loader( self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: tuple[int, ...] | int | None = None, ) -> None: tp_rank = self.tp_rank param_tp_rank = getattr(param, "tp_rank", None) if loaded_shard_id in self.replicated_shard_ids: self.tp_rank = 0 if param_tp_rank is not None: param.tp_rank = 0 try: super().weight_loader(param, loaded_weight, loaded_shard_id) finally: self.tp_rank = tp_rank if param_tp_rank is not None: param.tp_rank = param_tp_rank def weight_loader_v2( self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: tuple[int, ...] | int | None = None, ) -> None: tp_rank = self.tp_rank param_tp_rank = getattr(param, "tp_rank", None) if loaded_shard_id in self.replicated_shard_ids: self.tp_rank = 0 if param_tp_rank is not None: param.tp_rank = 0 try: super().weight_loader_v2(param, loaded_weight, loaded_shard_id) finally: self.tp_rank = tp_rank if param_tp_rank is not None: param.tp_rank = param_tp_rank @torch.compile( dynamic=True, backend=current_platform.simple_compile_backend, options=maybe_disable_graph_partition(current_platform.simple_compile_backend), ) def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor: """Fuse the fp32 cast + sigmoid into one Inductor kernel.""" return x.float().sigmoid() def _resolve_kda_prefill_backend( backend: str, head_dim: int, dtype: torch.dtype, lower_bound: float | None ) -> str: """Pick the chunked-prefill kernel: FlashKDA (fused CUDA, ~2-4x faster on SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate) or the Triton ``chunk_kda_with_fused_gate`` path. ``backend`` comes from ``additional_config.kda_prefill_backend`` (auto / triton / flashkda).""" if backend not in ("auto", "triton", "flashkda"): raise ValueError(f"Unsupported KDA prefill backend: {backend}") capability = current_platform.get_device_capability() supported = ( current_platform.is_cuda() and capability is not None and capability.major in (9, 10, 12) and head_dim == 128 and dtype == torch.bfloat16 and lower_bound is not None ) if backend == "flashkda" and not supported: raise RuntimeError( "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, head_dim=128 " "and a bounded KDA gate." ) return "flashkda" if supported and backend != "triton" else "triton" class Glm5NextLinearAttention(GatedDeltaNetAttention): head_dim: int num_heads: int conv_size: int def get_state_dtype( self, ) -> tuple[torch.dtype, torch.dtype]: if self.model_config is None or self.cache_config is None: raise ValueError("model_config and cache_config must be set") return MambaStateDtypeCalculator.kda_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype ) def get_state_shape( self, ) -> tuple[tuple[int, ...], tuple[int, ...]]: # conv_state width must include num_spec so the spec-decode conv update # (causal_conv1d_update with num_accepted_tokens + max_query_len) can # slide the window across the draft-verify tokens without reading past # the allocated width. Matches qwen_gdn_linear_attn.get_state_shape. return MambaStateShapeCalculator.kda_state_shape( self.tp_size, self.num_heads, self.head_dim, conv_kernel_size=self.conv_size, num_spec=self.num_spec, ) def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: spec = super().get_kv_cache_spec(vllm_config) if self._recoverssm: # One checkpoint block per request. The commit writes only the final state. spec = replace(spec, num_speculative_blocks=0) return spec def get_attn_backend(self) -> type[AttentionBackend]: if self._recoverssm: return Glm5NextRecoverSSMBackend return super().get_attn_backend() def __init__( self, config: Glm5NextConfig, vllm_config: VllmConfig, prefix: str = "", ) -> None: # KDA projections remain BF16 because fp8 checkpoints omit their scales. saved_quant_config = vllm_config.quant_config try: vllm_config.quant_config = None super().__init__(config, vllm_config, prefix) finally: vllm_config.quant_config = saved_quant_config self._recoverssm = _RECOVERSSM and self.num_spec > 0 # The RecoverSSM metadata builder allocates these (recoverssm.Records). self.recoverssm_records = None self.head_dim = config.linear_head_dim self.num_heads = config.linear_num_heads self.conv_size = config.linear_conv_kernel_dim assert self.num_heads % self.tp_size == 0 self.local_num_heads = divide(self.num_heads, self.tp_size) projection_size = self.head_dim * self.num_heads self.local_projection_size = divide(projection_size, self.tp_size) _, recurrent_state_shape = self.get_state_shape() _, recurrent_state_dtype = self.get_state_dtype() scatter_states.register_warmup( state_shape=recurrent_state_shape, state_dtype=recurrent_state_dtype, indices_dtype=torch.int32, max_num_tokens=vllm_config.scheduler_config.max_num_seqs, ) # Merge q, k, v, b, f_a, g_a projections into one GEMM (6→1 launches). # Order matches checkpoint's fused_qkvbfg_a_proj convention. # Shards 4 (f_a) and 5 (g_a) are replicated across TP ranks. self.in_proj_qkvbfg_a = _Glm5NextMergedColumnParallelLinear( self.hidden_size, [ projection_size, # q (shard 0) projection_size, # k (shard 1) projection_size, # v (shard 2) self.num_heads, # b (shard 3) self.head_dim, # f_a (shard 4, replicated) self.head_dim, # g_a (shard 5, replicated) ], replicated_shard_ids=(4, 5), tp_size=self.tp_size, bias=False, quant_config=self.quant_config, prefix=f"{prefix}.in_proj_qkvbfg_a", ) self.f_b_proj = ColumnParallelLinear( self.head_dim, projection_size, bias=False, quant_config=self.quant_config, prefix=f"{prefix}.f_b_proj", ) self.dt_bias = nn.Parameter( torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32) ) set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)}) self.q_conv1d = ColumnParallelLinear( input_size=self.conv_size, output_size=projection_size, bias=False, params_dtype=torch.float32, prefix=f"{prefix}.q_conv1d", ) self.k_conv1d = ColumnParallelLinear( input_size=self.conv_size, output_size=projection_size, bias=False, params_dtype=torch.float32, prefix=f"{prefix}.k_conv1d", ) self.v_conv1d = ColumnParallelLinear( input_size=self.conv_size, output_size=projection_size, bias=False, params_dtype=torch.float32, prefix=f"{prefix}.v_conv1d", ) # unsqueeze to fit conv1d weights shape into the linear weights shape. # Can't do this in `weight_loader` since it already exists in # `ColumnParallelLinear` and `set_weight_attrs` # doesn't allow to override it self.q_conv1d.weight.data = self.q_conv1d.weight.data.unsqueeze(1) self.k_conv1d.weight.data = self.k_conv1d.weight.data.unsqueeze(1) self.v_conv1d.weight.data = self.v_conv1d.weight.data.unsqueeze(1) # Lazily-built merged q|k|v conv weight (built on first forward, after # weights are loaded). See _forward. self._merged_conv_weight: torch.Tensor | None = None self.A_log = nn.Parameter( torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32) ) set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(2)}) self.g_b_proj = ColumnParallelLinear( self.head_dim, projection_size, bias=False, quant_config=self.quant_config, prefix=f"{prefix}.g_b_proj", ) self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid") self.o_proj = RowParallelLinear( projection_size, self.hidden_size, bias=False, quant_config=self.quant_config, prefix=f"{prefix}.o_proj", ) compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self # Checkpoints store A_log as 1-D; the model parameter is 4-D. def _a_log_weight_loader(param, loaded_weight): if loaded_weight.dim() == 1: loaded_weight = loaded_weight.view([1, 1, -1, 1]) return sharded_weight_loader(2)(param, loaded_weight) self.A_log.weight_loader = _a_log_weight_loader # GLM-5.3-Flash uses a bounded sigmoid gate instead of the default # unbounded softplus gate. self.kda_safe_gate = True self.kda_lower_bound = config.linear_lower_bound # Process-global conv-state layout, resolved once here instead of on # every _forward call (it reads an env-derived flag each time). self._conv_state_dim_first = is_conv_state_dim_first() additional_config = vllm_config.additional_config self.kda_prefill_backend = _resolve_kda_prefill_backend( additional_config.get("kda_prefill_backend", "auto") if isinstance(additional_config, dict) else "auto", self.head_dim, vllm_config.model_config.dtype, self.kda_lower_bound, ) self._flashkda_buffer_specs: ( tuple[tuple[tuple[int, ...], torch.dtype], ...] | None ) = None if self.kda_prefill_backend == "flashkda": import vllm._flashkda_C # noqa: F401 max_tokens = vllm_config.scheduler_config.max_num_batched_tokens max_seqs = vllm_config.scheduler_config.max_num_seqs workspace_size = torch.ops._flashkda_C.get_workspace_size( max_tokens, self.local_num_heads, max_seqs ) self._flashkda_buffer_specs = ( ( (max_seqs, self.local_num_heads, self.head_dim, self.head_dim), self.get_state_dtype()[1], ), ((workspace_size,), torch.uint8), # Output buffer for steps that also carry spec-decode tokens: # the non-spec tokens are then scattered by non_spec_token_indx. ( (1, max_tokens, self.local_num_heads, self.head_dim), vllm_config.model_config.dtype, ), ) def _flashkda_prefill( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, initial_state: torch.Tensor, cu_seqlens: torch.Tensor, out: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: """Fused KDA chunked prefill (FlashKDA). Takes the raw gate logits ``g`` and raw ``beta`` logits, l2-normalizes q/k in-kernel and applies the bounded gate ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, matching ``chunk_kda_with_fused_gate(..., safe_gate=True)``. Writes the attention output into ``out`` (a workspace buffer when ``None``) and returns ``(out, final_state)``.""" assert self._flashkda_buffer_specs is not None final_state, workspace, workspace_out = ( current_workspace_manager().get_simultaneous(*self._flashkda_buffer_specs) ) final_state = final_state[: initial_state.shape[0]] if out is None: out = workspace_out[:, : q.shape[1]] # FlashKDA hardcodes dense q/k/v/g strides; beta may be row-strided. torch.ops._flashkda_C.fwd( q.contiguous(), k.contiguous(), v.contiguous(), g.contiguous(), beta, self.head_dim**-0.5, out, workspace, self.A_log.view(-1), self.dt_bias.view(-1, self.head_dim), self.kda_lower_bound, initial_state.contiguous(), final_state, cu_seqlens.contiguous(), None, None, ) return out, final_state def forward( self, hidden_states: torch.Tensor, positions: torch.Tensor, ) -> torch.Tensor: num_tokens = hidden_states.size(0) # One merged GEMM for q, k, v, b, f_a, g_a (replaces 6 separate GEMMs). projected = self.in_proj_qkvbfg_a(hidden_states)[0] qkv, beta_raw, f_a, g_a = projected.split( [ 3 * self.local_projection_size, self.local_num_heads, self.head_dim, self.head_dim, ], dim=-1, ) # Beta stays raw (bf16) here: the recurrent kernel sigmoids it in fp32 # at load (SIGMOID_BETA), and only the chunked prefill path needs the # pre-computed fp32 sigmoid — computed lazily in _forward. Pure decode # / spec-verify steps then skip the _cast_sigmoid kernel and its fp32 # intermediate entirely. beta = beta_raw.unsqueeze(0) g1 = self.f_b_proj(f_a)[0] g1 = g1.reshape(1, -1, self.local_num_heads, self.head_dim) g_proj_states = self.g_b_proj(g_a)[0] # Must stay 3D: rms_norm_gated reads H from g.shape[-2]. g2 = g_proj_states.reshape(-1, self.local_num_heads, self.head_dim) core_attn_out = torch.empty( (1, num_tokens, self.local_num_heads, self.head_dim), dtype=hidden_states.dtype, device=hidden_states.device, ) # Call the decorated eager break directly so host-side prefill branches # are not captured by PIECEWISE CUDA graphs. self._forward( qkv_proj_states=qkv, g1=g1, beta=beta, core_attn_out=core_attn_out, ) core_attn_out = self.o_norm(core_attn_out, g2) core_attn_out = core_attn_out.reshape(core_attn_out.size(1), -1) return self.o_proj(core_attn_out)[0] @eager_break_during_capture def _forward( self, qkv_proj_states: torch.Tensor, g1: torch.Tensor, beta: torch.Tensor, core_attn_out: torch.Tensor, ) -> None: forward_context = get_forward_context() attn_metadata_raw = forward_context.attn_metadata if attn_metadata_raw is None: return assert isinstance(attn_metadata_raw, dict) attn_metadata_narrowed = attn_metadata_raw.get(self.prefix) if attn_metadata_narrowed is None: # Profile/warmup dummy runs may omit mamba-family metadata. return assert isinstance(attn_metadata_narrowed, GDNAttentionMetadata) has_initial_state = attn_metadata_narrowed.has_initial_state non_spec_query_start_loc = attn_metadata_narrowed.non_spec_query_start_loc non_spec_state_indices_tensor = ( attn_metadata_narrowed.non_spec_state_indices_tensor ) # noqa: E501 num_actual_tokens = attn_metadata_narrowed.num_actual_tokens # Spec-decode metadata (all None when speculative decoding is disabled). spec_sequence_masks = attn_metadata_narrowed.spec_sequence_masks spec_query_start_loc = attn_metadata_narrowed.spec_query_start_loc spec_state_indices_tensor = attn_metadata_narrowed.spec_state_indices_tensor spec_token_indx = attn_metadata_narrowed.spec_token_indx non_spec_token_indx = attn_metadata_narrowed.non_spec_token_indx num_accepted_tokens = attn_metadata_narrowed.num_accepted_tokens num_spec_decodes = attn_metadata_narrowed.num_spec_decodes # RecoverSSM metadata leaves the mask None. num_spec_decodes marks spec rows. use_spec = num_spec_decodes > 0 and ( spec_sequence_masks is not None or self._recoverssm ) # Safe-gate checkpoints use the bounded sigmoid variant. safe_gate = self.kda_safe_gate lower_bound = self.kda_lower_bound constant_caches = self.kv_cache qkv_proj_states = qkv_proj_states[:num_actual_tokens] g1 = g1[:, :num_actual_tokens] beta = beta[:, :num_actual_tokens] (conv_state, recurrent_state) = constant_caches # conv_state must be (..., dim, width-1) for the conv kernels. # DS layout stores it that way directly; SD layout needs a transpose. # Layout is process-global and resolved once at init (see __init__). if not self._conv_state_dim_first: conv_state = conv_state.transpose(-1, -2) # One merged short-conv over q|k|v instead of three separate calls. The # 1D conv is independent per channel, so concatenating q/k/v along the # channel dim and running a single causal_conv1d is bit-identical to # three calls. The merged weight is q|k|v conv weights concatenated; # built once and cached (params are fixed after load). conv_state is # already stored as the merged q|k|v state, so it is used directly. if self._merged_conv_weight is None: def _w(m): return m.weight.view(m.weight.size(0), m.weight.size(2)) self._merged_conv_weight = torch.cat( [_w(self.q_conv1d), _w(self.k_conv1d), _w(self.v_conv1d)], dim=0, ).contiguous() conv_weights = self._merged_conv_weight conv_bias = self.q_conv1d.bias # Split projections / gating into spec (draft-verify) and non-spec token # groups when speculative decoding is active. Spec tokens carry # num_spec+1 recurrent-state columns each and are advanced with # num_accepted_tokens for rejection-sampling rollback; non-spec tokens # are one-per-request. Mirrors olmo_gdn_linear_attn.py. Projections are # [n, *] (token dim 0); g1/beta are [1, n, h, d] (token dim 1). if use_spec: # In a pure spec-verify step (no non-spec tokens) the metadata # builder sets spec_token_indx = arange(num_actual_tokens), making # the index_select calls below identity copies. Skip them on this # steady-state decode hot path. The outputs alias the inputs here; # the downstream conv/recurrent kernels read them without mutating # in place, so the aliasing is safe. if non_spec_token_indx is None or non_spec_token_indx.numel() == 0: qkv_spec = qkv_proj_states g1_spec = g1 beta_spec = beta else: qkv_spec = qkv_proj_states.index_select(0, spec_token_indx) g1_spec = g1.index_select(1, spec_token_indx) beta_spec = beta.index_select(1, spec_token_indx) if non_spec_token_indx is not None and non_spec_token_indx.numel() > 0: qkv_ns = qkv_proj_states.index_select(0, non_spec_token_indx) g1_ns = g1.index_select(1, non_spec_token_indx) beta_ns = beta.index_select(1, non_spec_token_indx) else: qkv_ns = g1_ns = beta_ns = None else: qkv_spec = g1_spec = beta_spec = None qkv_ns, g1_ns, beta_ns = qkv_proj_states, g1, beta # --- causal conv1d: spec (draft-verify) path --- if use_spec: assert spec_state_indices_tensor is not None assert num_accepted_tokens is not None conv_idx = spec_state_indices_tensor[:, 0][:num_spec_decodes] # RecoverSSM keeps one state slot, but the window is still k + 1. conv_mql = ( self.num_spec + 1 if self._recoverssm else spec_state_indices_tensor.size(-1) ) qkv_spec = causal_conv1d_update( qkv_spec, conv_state, conv_weights, conv_bias, activation="silu", conv_state_indices=conv_idx, num_accepted_tokens=num_accepted_tokens, query_start_loc=spec_query_start_loc, max_query_len=conv_mql, ) q_spec, k_spec, v_spec = qkv_spec.split(self.local_projection_size, dim=-1) # --- causal conv1d: non-spec path (prefill or plain decode) --- q_ns = k_ns = v_ns = None if attn_metadata_narrowed.num_prefills > 0: assert qkv_ns is not None # q, k and v come out as separate dense tensors, which FlashKDA # needs; split from one [T, 3 * P] output they had to be copied. q_ns, k_ns, v_ns = causal_conv1d_fn( qkv_ns.transpose(0, 1), conv_weights, conv_bias, activation="silu", conv_states=conv_state, has_initial_state=has_initial_state, cache_indices=non_spec_state_indices_tensor, query_start_loc=non_spec_query_start_loc, metadata=attn_metadata_narrowed, out_group=self.local_projection_size, ).unbind(0) elif attn_metadata_narrowed.num_decodes > 0: assert non_spec_state_indices_tensor is not None decode_conv_indices = non_spec_state_indices_tensor[ : attn_metadata_narrowed.num_decodes ] qkv_ns = causal_conv1d_update( qkv_ns, conv_state, conv_weights, conv_bias, activation="silu", conv_state_indices=decode_conv_indices, ) q_ns, k_ns, v_ns = qkv_ns.split(self.local_projection_size, dim=-1) def _rearr(x): return x.reshape(1, -1, self.local_num_heads, self.head_dim) # --- core attention: spec (draft-verify) path --- core_attn_out_spec = None # In a pure spec-verify step (no non-spec tokens) the recurrent kernel # can write straight into the layer output buffer, skipping the # fresh allocation + copy below. Mixed steps must scatter via # spec_token_indx, so they keep the kernel-managed output. spec_out = ( core_attn_out[0, :num_actual_tokens].unsqueeze(0) if non_spec_token_indx is None or non_spec_token_indx.numel() == 0 else None ) if use_spec and self._recoverssm: assert spec_state_indices_tensor is not None assert spec_query_start_loc is not None assert self.recoverssm_records is not None core_attn_out_spec = kda_recoverssm_verify( q=_rearr(q_spec), k=_rearr(k_spec), v=_rearr(v_spec), g=g1_spec, beta=beta_spec, a_log=self.A_log, g_bias=self.dt_bias, lower_bound=lower_bound, checkpoint_state=recurrent_state, records=self.recoverssm_records, query_start_loc=spec_query_start_loc[: num_spec_decodes + 1], state_indices=spec_state_indices_tensor[:num_spec_decodes, 0], out=spec_out, ) elif use_spec: assert spec_state_indices_tensor is not None assert num_accepted_tokens is not None assert spec_query_start_loc is not None # Gate computed inside the recurrent kernel (COMPUTE_GATE) from # raw g1 — replicates fused_kda_gate's arithmetic bit-for-bit and # skips its launch + fp32 [n, H, D] intermediate per layer. core_attn_out_spec, _ = fused_recurrent_kda( q=_rearr(q_spec), k=_rearr(k_spec), v=_rearr(v_spec), g=g1_spec, beta=beta_spec, initial_state=recurrent_state, use_qk_l2norm_in_kernel=True, cu_seqlens=spec_query_start_loc[: num_spec_decodes + 1], ssm_state_indices=spec_state_indices_tensor, num_accepted_tokens=num_accepted_tokens, out=spec_out, sigmoid_beta=True, a_log=self.A_log, g_bias=self.dt_bias, compute_gate=True, lower_bound=lower_bound, ) # --- core attention: non-spec path (prefill or plain decode) --- core_attn_out_non_spec = None # Only the plain-decode recurrent kernel can write straight into the # layer output buffer; the chunked prefill kernel cannot, so this # stays None there and the merge copy below runs as before. ns_out = None if attn_metadata_narrowed.num_prefills > 0: assert q_ns is not None assert non_spec_state_indices_tensor is not None assert has_initial_state is not None initial_state = gather_initial_states( recurrent_state, non_spec_state_indices_tensor, has_initial_state ) if self.kda_prefill_backend == "flashkda": # Non-spec step: write straight into the layer output buffer # (dense token order, no merge copy). Step with spec-decode # tokens: write to the workspace buffer and scatter below. ns_out = None if use_spec else core_attn_out[:, :num_actual_tokens] core_attn_out_non_spec, last_recurrent_state = self._flashkda_prefill( q=_rearr(q_ns), k=_rearr(k_ns), v=_rearr(v_ns), g=g1_ns, beta=beta_ns, initial_state=initial_state, cu_seqlens=non_spec_query_start_loc, out=ns_out, ) else: ( core_attn_out_non_spec, last_recurrent_state, ) = chunk_kda_with_fused_gate( q=_rearr(q_ns), k=_rearr(k_ns), v=_rearr(v_ns), raw_g=g1_ns, # Chunk path wants the pre-sigmoided fp32 beta (its # kernels don't sigmoid); beta_ns is raw bf16 from forward. beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0), A_log=self.A_log, g_bias=self.dt_bias, initial_state=initial_state, output_final_state=True, use_qk_l2norm_in_kernel=True, cu_seqlens=non_spec_query_start_loc, safe_gate=safe_gate, lower_bound=lower_bound, ) # Init cache scatter_states( recurrent_state, last_recurrent_state, non_spec_state_indices_tensor, ) elif attn_metadata_narrowed.num_decodes > 0: assert non_spec_query_start_loc is not None assert non_spec_state_indices_tensor is not None # Plain decode step (no spec tokens): token order is dense, so the # kernel can write straight into the layer output buffer. A mixed # step scatters non-spec output via non_spec_token_indx instead. # Gate computed in-kernel (COMPUTE_GATE), beta sigmoided in-kernel. if not use_spec: ns_out = spec_out core_attn_out_non_spec, _ = fused_recurrent_kda( q=_rearr(q_ns), k=_rearr(k_ns), v=_rearr(v_ns), g=g1_ns, beta=beta_ns, initial_state=recurrent_state, use_qk_l2norm_in_kernel=True, cu_seqlens=non_spec_query_start_loc[ : attn_metadata_narrowed.num_decodes + 1 ], ssm_state_indices=non_spec_state_indices_tensor, out=ns_out, sigmoid_beta=True, a_log=self.A_log, g_bias=self.dt_bias, compute_gate=True, lower_bound=lower_bound, ) # --- merge spec / non-spec outputs back into token order --- if use_spec and core_attn_out_non_spec is not None: assert core_attn_out_spec is not None core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec) core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) elif use_spec: assert core_attn_out_spec is not None if spec_out is None: core_attn_out[0, :num_actual_tokens] = core_attn_out_spec.squeeze(0) else: assert core_attn_out_non_spec is not None if ns_out is None: core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ 0, :num_actual_tokens ]