class KimiK3DeltaAttention(GatedDeltaNetAttention):
def get_attn_backend(self) -> type[AttentionBackend]:
return KimiK3KDAAttentionBackend
def get_state_dtype(
self,
) -> tuple[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,
self.cache_config.mamba_ssm_cache_dtype,
)
def get_state_shape(
self,
) -> tuple[tuple[int, ...], ...]:
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 __init__(
self,
config: KimiLinearConfig,
vllm_config: VllmConfig,
prefix: str = "",
aux_stream: torch.cuda.Stream | None = None,
run_gemm_rs_ar: bool = False,
) -> None:
super().__init__(config, vllm_config, prefix)
self.use_recoverssm = self.cache_config.use_kda_recoverssm
if self.cache_config.use_replayssm and not self.use_recoverssm:
raise ValueError(
"Kimi-K3 supports --use-replayssm only with speculative decoding"
)
self.spec_query_len = 1 + self.num_spec
kda_config = config.linear_attn_config # type: ignore[attr-defined]
assert kda_config is not None, "linear_attn_config must be set"
self.head_dim = kda_config["head_dim"]
self.num_heads = kda_config["num_heads"]
assert self.num_heads % self.tp_size == 0
self.local_num_heads = divide(self.num_heads, self.tp_size)
self.projection_size = self.head_dim * self.num_heads
self.local_projection_size = divide(self.projection_size, self.tp_size)
self.conv_size = kda_config["short_conv_kernel_size"]
self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None)
if self.gate_lower_bound is not None:
assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, (
"KDA gate lower bound must be in "
f"[{_KDA_GATE_LOGBOUND_MIN}, 0). "
f"Got {self.gate_lower_bound}."
)
assert kda_config.get("use_full_rank_gate", False), (
"KimiK3DeltaAttention requires a full-rank gate"
)
self._projection_aux_stream = aux_stream
self._projection_events = (
(torch.cuda.Event(), torch.cuda.Event()) if aux_stream is not None else None
)
self._projection_overlap_max_tokens = 0
# Keep f_a before the narrow beta shard, then align each TP-local row.
qkvg_output_sizes = [self.projection_size] * 4
in_proj_output_sizes = qkvg_output_sizes + [
self.head_dim,
self.num_heads,
]
local_output_size = (
4 * self.local_projection_size + self.head_dim + self.local_num_heads
)
in_proj_prefix = f"{prefix}.in_proj_qkvgfab"
alignment = (
128
if isinstance(self.quant_config, ModelOptMixedPrecisionConfig)
and self.quant_config._resolve_quant_algo(in_proj_prefix) == "FP8_PB_WO"
else 16
)
self.in_proj_padding = -local_output_size % alignment
if self.in_proj_padding:
in_proj_output_sizes.append(self.in_proj_padding * self.tp_size)
self.in_proj_qkvgfab = _KimiGDNMergedColumnParallelLinear(
self.hidden_size,
in_proj_output_sizes,
replicated_shard_id=4,
tp_size=self.tp_size,
padding_size=self.in_proj_padding,
bias=False,
quant_config=self.quant_config,
prefix=in_proj_prefix,
)
self.f_b_proj = ColumnParallelLinear(
self.head_dim,
self.projection_size,
bias=False,
quant_config=self.quant_config,
prefix=f"{prefix}.f_b_proj",
)
self.dt_bias = nn.Parameter(
torch.empty(self.local_projection_size, dtype=torch.float32)
)
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
# One packed parameter and cache let decode run a single conv update.
# Prefill slices them back into Q/K/V to obtain dense outputs cheaply.
self.conv1d = ColumnParallelLinear(
input_size=self.conv_size,
output_size=3 * self.projection_size,
bias=False,
params_dtype=torch.float32,
prefix=f"{prefix}.conv1d",
)
self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
# Keep a width-major copy for fused decode without changing the layout
# consumed by the prefill and fallback decode kernels.
conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()[:2]
additional_config = vllm_config.additional_config
decode_backend = (
additional_config.get("kda_decode_backend", "auto")
if isinstance(additional_config, dict)
else "auto"
)
self.kda_decode_backend = resolve_kda_decode_backend(
decode_backend,
self.local_num_heads,
self.head_dim,
self.conv_size,
self.num_spec,
vllm_config.model_config.dtype,
conv_state_dtype,
recurrent_state_dtype,
)
spec_decode_backend = (
additional_config.get("kda_spec_decode_backend", "auto")
if isinstance(additional_config, dict)
else "auto"
)
self.kda_spec_decode_backend = resolve_kda_spec_decode_backend(
spec_decode_backend,
self.local_num_heads,
self.head_dim,
self.conv_size,
self.num_spec,
vllm_config.model_config.dtype,
conv_state_dtype,
recurrent_state_dtype,
self.gate_lower_bound,
self.use_recoverssm,
)
# RecoverSSM: this layer's verify records of a speculative step, and its
# verify and commit kernels: "cute" (SM10x, the short convolution fused
# into the verify) where they can run, else "triton".
self.recoverssm_backend: Literal["cute", "triton"] | None = None
self.recoverssm_records: KDARecoverSSMRecords | None = None
self.recoverssm_verify: CuteRecoverSSMVerify | None = None
if self.use_recoverssm:
self.recoverssm_records = KDARecoverSSMRecords.allocate(
vllm_config.scheduler_config.max_num_seqs,
self.local_num_heads,
self.spec_query_len,
self.head_dim,
vllm_config.model_config.dtype,
self.conv1d.weight.device,
)
max_num_seqs = vllm_config.scheduler_config.max_num_seqs
records_bytes = len(kda_config["kda_layers"]) * sum(
t.numel() * t.element_size()
for t in (
self.recoverssm_records.correction,
self.recoverssm_records.key_gate,
)
)
logger.info_once(
"Allocated the KDA RecoverSSM verify records: %.2f GiB in total "
"for %d requests (%.2f MiB per request).",
records_bytes / 2**30,
max_num_seqs,
records_bytes / max_num_seqs / 2**20,
)
supported, reason = can_use_cute_recoverssm(
head_dim=self.head_dim,
window=self.spec_query_len,
conv_width=self.conv_size,
has_conv_bias=self.conv1d.bias is not None,
activation_dtype=vllm_config.model_config.dtype,
conv_state_dtype=conv_state_dtype,
state_dtype=recurrent_state_dtype,
conv_weight_dtype=self.conv1d.weight.dtype,
lower_bound=self.gate_lower_bound,
)
self.recoverssm_backend = "cute" if supported else "triton"
if self.recoverssm_backend == "triton":
logger.info_once(
"Using the Triton KDA RecoverSSM verify and commit: the CuTe "
"kernels cannot run (%s).",
reason,
)
else:
logger.info_once("Using the CuTe KDA RecoverSSM verify and commit.")
self.recoverssm_verify = CuteRecoverSSMVerify(
num_heads=self.local_num_heads,
window=self.spec_query_len,
conv_dim_first=is_conv_state_dim_first(),
conv_weight_dtype=self.conv1d.weight.dtype,
device=self.conv1d.weight.device,
)
self.recoverssm_verify.compile()
compile_state_commit(
self.local_num_heads,
self.spec_query_len,
self.conv1d.weight.device,
align=self.cache_config.mamba_cache_mode == "align",
)
decode_conv1d_weight = None
if (
self.kda_decode_backend != "triton"
or self.kda_spec_decode_backend == "flashinfer"
):
decode_conv1d_weight = torch.empty(
3,
self.conv_size,
self.local_projection_size,
dtype=self.conv1d.weight.dtype,
device=self.conv1d.weight.device,
)
self.register_buffer(
"decode_conv1d_weight", decode_conv1d_weight, persistent=False
)
delattr(self.conv1d.weight, "weight_loader")
set_weight_attrs(
self.conv1d.weight,
{
"weight_loader": _make_decode_conv1d_weight_loader(
[self.projection_size] * 3,
self.tp_size,
self.tp_rank,
decode_conv1d_weight,
)
},
)
self.A_log = nn.Parameter(
torch.empty(self.local_num_heads, dtype=torch.float32)
)
set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)})
backend = (
additional_config.get("kda_prefill_backend", "auto")
if isinstance(additional_config, dict)
else "auto"
)
self.kda_prefill_backend = resolve_kda_prefill_backend(
backend,
self.head_dim,
vllm_config.model_config.dtype,
recurrent_state_dtype,
self.gate_lower_bound,
)
self._flashkda_buffer_specs: (
tuple[tuple[tuple[int, ...], torch.dtype], ...] | None
) = None
self._flashinfer_kda_output_spec: tuple[tuple[int, ...], torch.dtype] | None = (
None
)
self._checkpoint_exporter: FlashKDAPrefillCheckpointExporter | None = None
if self.kda_prefill_backend == "flashkda":
T = vllm_config.scheduler_config.max_num_batched_tokens
N = vllm_config.scheduler_config.max_num_seqs
H, D = self.local_num_heads, self.head_dim
import vllm._flashkda_C # noqa: F401
workspace_size = torch.ops._flashkda_C.get_workspace_size(T, H, N)
self._flashkda_buffer_specs = (
((1, T, H, D), self.model_config.dtype),
((N, H, D, D), self.get_state_dtype()[1]),
((N, H, D, D), self.get_state_dtype()[1]),
((workspace_size,), torch.uint8),
)
self._checkpoint_exporter = FlashKDAPrefillCheckpointExporter()
elif self.kda_prefill_backend == "flashinfer":
T = vllm_config.scheduler_config.max_num_batched_tokens
H, D = self.local_num_heads, self.head_dim
self._flashinfer_kda_output_spec = (
(1, T, H, D),
self.model_config.dtype,
)
# Profiling skips attention, so reserve its runtime layout before the
# KV cache is sized. Layers share the manager's grow-only allocation.
if is_workspace_manager_initialized():
if self._flashkda_buffer_specs is not None:
current_workspace_manager().get_simultaneous(
*self._flashkda_buffer_specs
)
elif self._flashinfer_kda_output_spec is not None:
current_workspace_manager().get_simultaneous(
self._flashinfer_kda_output_spec
)
self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
decode_norm_weight = None
if decode_conv1d_weight is not None:
decode_norm_weight = torch.empty(
self.head_dim,
dtype=torch.float32,
device=self.o_norm.weight.device,
)
self.register_buffer("decode_norm_weight", decode_norm_weight, persistent=False)
if decode_norm_weight is not None:
# Upcast once while loading; direct BF16 norm weights slow the
# fully fused decode kernel.
if hasattr(self.o_norm.weight, "weight_loader"):
delattr(self.o_norm.weight, "weight_loader")
set_weight_attrs(
self.o_norm.weight,
{"weight_loader": _make_decode_norm_weight_loader(decode_norm_weight)},
)
self.o_proj = RowParallelLinear(
self.projection_size,
self.hidden_size,
bias=False,
quant_config=self.quant_config,
prefix=f"{prefix}.o_proj",
)
self.gemm_rs_ar = None
if run_gemm_rs_ar:
from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import (
get_gemm_rs_ar,
)
gemm_rs_ar = get_gemm_rs_ar()
if gemm_rs_ar.can_run(self.o_proj):
self.gemm_rs_ar = gemm_rs_ar
else:
gemm_rs_ar.warn_incompatible_projection()
compilation_config = 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
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec:
spec = super().get_kv_cache_spec(vllm_config)
assert isinstance(spec, MambaSpec)
alignment = kda_prefill_checkpoint_alignment(self.kda_prefill_backend)
return replace(
spec,
num_prefill_checkpoint_blocks=int(alignment is not None),
prefill_checkpoint_alignment=alignment,
)
def forward(
self,
hidden_states: torch.Tensor,
positions: torch.Tensor,
) -> torch.Tensor:
num_tokens = hidden_states.size(0)
projection_events = self._projection_events
projection_aux_stream = self._projection_aux_stream
if (
0 < num_tokens <= self._projection_overlap_max_tokens
and hidden_states.stride() == (self.hidden_size, 1)
and projection_events is not None
and projection_aux_stream is not None
and torch.cuda.is_current_stream_capturing()
):
from vllm.models.kimi_k3.nvidia.low_latency_gemm import (
run_kda_projection_overlap,
)
projected_qkvg, g1, beta = run_kda_projection_overlap(
hidden_states,
self.in_proj_qkvgfab.weight,
self.f_b_proj.weight,
projection_aux_stream,
projection_events,
)
mixed_qkv, g_proj_states = projected_qkvg.split(
[3 * self.local_projection_size, self.local_projection_size],
dim=-1,
)
else:
projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]
split_sizes = [
3 * self.local_projection_size,
self.local_projection_size,
self.head_dim,
self.local_num_heads,
]
if self.in_proj_padding:
split_sizes.append(self.in_proj_padding)
projected = projected_qkvgfab.split(split_sizes, dim=-1)
mixed_qkv, g_proj_states, f_a, beta = projected[:4]
g1 = self.f_b_proj(f_a)[0]
beta = beta.unsqueeze(0)
g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)
g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=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,
)
self._forward(
mixed_qkv=mixed_qkv,
g1=g1,
g2=g2,
beta=beta,
core_attn_out=core_attn_out,
)
core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)")
if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(core_attn_out):
return self.gemm_rs_ar.apply(core_attn_out, self.o_proj)
return self.o_proj(core_attn_out)[0]
def _recoverssm_verify(
self,
m: KimiK3KDAMetadata,
mixed_qkv: torch.Tensor,
raw_g: torch.Tensor,
raw_beta: torch.Tensor,
conv_weight: torch.Tensor,
query_start_loc: torch.Tensor,
out: torch.Tensor | None,
) -> torch.Tensor:
"""RecoverSSM verify of the speculative requests: the short convolution
and the KDA recurrence from each request's checkpoint, which is left
unchanged. The verify records it writes let the commit after sampling
fold the accepted tokens in (KimiK3KDAMetadata.commit_recoverssm_state).
"""
num_tokens = mixed_qkv.shape[0]
if out is None:
out = mixed_qkv.new_empty(
(1, num_tokens, self.local_num_heads, self.head_dim)
)
assert m.spec_state_indices_tensor is not None
assert m.num_accepted_tokens is not None
state_indices = m.spec_state_indices_tensor[: m.num_spec_decodes, 0]
records = self.recoverssm_records
assert records is not None
conv_state, state = self.kv_cache[:2]
if self.recoverssm_backend == "cute":
assert self.recoverssm_verify is not None
assert self.gate_lower_bound is not None
if self.recoverssm_verify(
mixed_qkv=mixed_qkv,
raw_gate=raw_g[0].flatten(-2),
raw_beta=raw_beta[0],
A_log=self.A_log,
dt_bias=self.dt_bias,
lower_bound=self.gate_lower_bound,
state=state,
records=records,
query_start_loc=query_start_loc,
state_indices=state_indices,
num_accepted_tokens=m.num_accepted_tokens,
conv_state=conv_state,
conv_weight=conv_weight,
out=out[0].flatten(-2),
):
return out
# The CuTe verify declined misaligned activations: a mixed batch's
# speculative tokens can start at an odd row of a projection whose
# rows are not 16-byte aligned (e.g. the beta of the projection
# overlap path under graph capture). The Triton verify runs instead,
# leaving the same conv state, records and SSM state.
if not is_conv_state_dim_first():
conv_state = conv_state.transpose(-1, -2)
mixed_qkv = causal_conv1d_update(
mixed_qkv,
conv_state,
conv_weight,
self.conv1d.bias,
activation="silu",
conv_state_indices=state_indices,
num_accepted_tokens=m.num_accepted_tokens,
query_start_loc=query_start_loc,
max_query_len=self.spec_query_len,
validate_data=False,
out=torch.empty_like(mixed_qkv),
)
q, k, v = (
rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
for x in mixed_qkv.split(self.local_projection_size, dim=-1)
)
return kda_recoverssm_verify(
q=q,
k=k,
v=v,
raw_g=raw_g,
raw_beta=raw_beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
lower_bound=self.gate_lower_bound,
checkpoint_state=state,
records=records,
query_start_loc=query_start_loc,
state_indices=state_indices,
spec_query_len=self.spec_query_len,
out=out,
)
@eager_break_during_capture
def _forward(
self,
mixed_qkv: torch.Tensor,
g1: torch.Tensor,
g2: 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
from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
chunk_kda_with_fused_gate,
fused_recurrent_kda,
fused_recurrent_kda_packed_decode,
)
assert isinstance(attn_metadata_raw, dict)
attn_metadata_narrowed = attn_metadata_raw.get(self.prefix)
if attn_metadata_narrowed is None:
return
assert isinstance(attn_metadata_narrowed, KimiK3KDAMetadata)
m = attn_metadata_narrowed
has_initial_state = m.has_initial_state
non_spec_query_start_loc = m.non_spec_query_start_loc
non_spec_state_indices_tensor = m.non_spec_state_indices_tensor
spec_token_indx = m.spec_token_indx
non_spec_token_indx = m.non_spec_token_indx
spec_token_start = m.spec_token_start
non_spec_token_start = m.non_spec_token_start
spec_state_indices_tensor = m.spec_state_indices_tensor
spec_query_start_loc = m.spec_query_start_loc
num_accepted_tokens = m.num_accepted_tokens
num_actual_tokens = m.num_actual_tokens
checkpoint = m.checkpoint
has_spec_decode = m.num_spec_decodes > 0
mixed_qkv = mixed_qkv[:num_actual_tokens]
g1 = g1[:, :num_actual_tokens]
beta = beta[:, :num_actual_tokens]
conv_state, recurrent_state = self.kv_cache[:2]
# The convolution kernels consume (..., dim, width - 1).
if not is_conv_state_dim_first():
conv_state = conv_state.transpose(-1, -2)
if (
self.kda_spec_decode_backend == "flashinfer"
and has_spec_decode
and m.num_prefills == 0
and m.num_decodes == 0
):
assert self.decode_conv1d_weight is not None
assert self.decode_norm_weight is not None
assert spec_state_indices_tensor is not None
assert spec_query_start_loc is not None
assert num_accepted_tokens is not None
assert self.gate_lower_bound is not None
# "Packed" is FlashInfer's name for the varlen T>1 speculative path.
flashinfer_packed_fused_kda_decode(
x=mixed_qkv,
weight=self.decode_conv1d_weight,
conv_state=conv_state,
raw_gate=g1,
raw_beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
state_indices=spec_state_indices_tensor,
query_start_loc=spec_query_start_loc,
num_accepted_tokens=num_accepted_tokens,
state=recurrent_state,
output_gate=g2[:num_actual_tokens],
norm_weight=self.decode_norm_weight,
lower_bound=self.gate_lower_bound,
norm_eps=self.o_norm.eps,
output=core_attn_out[:, :num_actual_tokens],
)
return
if (
self.kda_decode_backend != "triton"
and self.decode_conv1d_weight is not None
and self.decode_norm_weight is not None
and not has_spec_decode
and m.num_prefills == 0
and m.num_decodes > 0
):
assert non_spec_state_indices_tensor is not None
state_indices = non_spec_state_indices_tensor[:num_actual_tokens]
if self.kda_decode_backend == "flashinfer":
flashinfer_fused_kda_decode(
x=mixed_qkv,
weight=self.decode_conv1d_weight,
conv_state=conv_state,
raw_gate=g1,
raw_beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
state_indices=state_indices,
state=recurrent_state,
output_gate=g2[:num_actual_tokens],
norm_weight=self.decode_norm_weight,
lower_bound=self.gate_lower_bound,
norm_eps=self.o_norm.eps,
output=core_attn_out[:, :num_actual_tokens],
)
else:
ops.fused_kda_decode(
x=mixed_qkv,
weight=self.decode_conv1d_weight,
bias=self.conv1d.bias,
conv_state=conv_state,
raw_g=g1,
raw_beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
state_indices=state_indices,
state=recurrent_state,
out=core_attn_out[:, :num_actual_tokens],
lower_bound=self.gate_lower_bound,
output_gate=g2[:num_actual_tokens],
norm_weight=self.decode_norm_weight,
norm_eps=self.o_norm.eps,
)
return
conv_weights = self.conv1d.weight.view(
self.conv1d.weight.size(0), self.conv1d.weight.size(2)
)
q_conv_weight, k_conv_weight, v_conv_weight = conv_weights.split(
self.local_projection_size, dim=0
)
q_conv_state, k_conv_state, v_conv_state = conv_state.split(
self.local_projection_size, dim=-2
)
# Separate multi-query speculative tokens from prefill/plain decode.
if has_spec_decode:
if m.num_prefills == 0 and m.num_decodes == 0:
mixed_qkv_spec = mixed_qkv
g1_spec, beta_spec = g1, beta
mixed_qkv_ns = g1_ns = beta_ns = None
elif spec_token_start is not None:
assert non_spec_token_start is not None
spec_end = spec_token_start + m.num_spec_decode_tokens
non_spec_end = (
non_spec_token_start + m.num_prefill_tokens + m.num_decode_tokens
)
spec_slice = slice(spec_token_start, spec_end)
non_spec_slice = slice(non_spec_token_start, non_spec_end)
mixed_qkv_spec = mixed_qkv[spec_slice]
g1_spec, beta_spec = g1[:, spec_slice], beta[:, spec_slice]
mixed_qkv_ns = mixed_qkv[non_spec_slice]
g1_ns, beta_ns = g1[:, non_spec_slice], beta[:, non_spec_slice]
else:
assert spec_token_indx is not None
assert non_spec_token_indx is not None
mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
g1_spec = g1.index_select(1, spec_token_indx)
beta_spec = beta.index_select(1, spec_token_indx)
mixed_qkv_ns = mixed_qkv.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:
mixed_qkv_spec = g1_spec = beta_spec = None
mixed_qkv_ns, g1_ns, beta_ns = mixed_qkv, g1, beta
# Spec-decode multi-query path.
core_attn_out_spec = None
if has_spec_decode:
assert spec_state_indices_tensor is not None
assert spec_query_start_loc is not None
spec_cu_seqlens = spec_query_start_loc[: m.num_spec_decodes + 1]
spec_out = (
core_attn_out[:, : mixed_qkv_spec.shape[0]]
if m.num_prefills == 0 and m.num_decodes == 0
else None
)
if spec_token_start is not None:
spec_out = core_attn_out[:, spec_slice]
if self.use_recoverssm:
core_attn_out_spec = self._recoverssm_verify(
m,
mixed_qkv_spec,
g1_spec,
beta_spec,
conv_weights,
spec_cu_seqlens,
spec_out,
)
else:
mixed_qkv_spec = causal_conv1d_update(
mixed_qkv_spec,
conv_state,
conv_weights,
self.conv1d.bias,
activation="silu",
conv_state_indices=spec_state_indices_tensor[:, 0][
: m.num_spec_decodes
],
num_accepted_tokens=num_accepted_tokens,
query_start_loc=spec_query_start_loc,
max_query_len=spec_state_indices_tensor.size(-1),
validate_data=False,
out=torch.empty_like(mixed_qkv_spec),
)
q_spec, k_spec, v_spec = (
rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
for x in mixed_qkv_spec.split(self.local_projection_size, dim=-1)
)
core_attn_out_spec, _ = fused_recurrent_kda(
q=q_spec,
k=k_spec,
v=v_spec,
raw_g=g1_spec,
raw_beta=beta_spec,
A_log=self.A_log,
dt_bias=self.dt_bias,
lower_bound=self.gate_lower_bound,
initial_state=recurrent_state,
cu_seqlens=spec_cu_seqlens,
ssm_state_indices=spec_state_indices_tensor,
num_accepted_tokens=num_accepted_tokens,
out=spec_out,
)
# Prefill or plain-decode path.
core_attn_out_non_spec = None
if mixed_qkv_ns is not None:
assert g1_ns is not None and beta_ns is not None
non_spec_out = None
if non_spec_token_start is not None:
non_spec_out = core_attn_out[:, non_spec_slice]
if m.num_prefills > 0:
q_ns, k_ns, v_ns = mixed_qkv_ns.split(
self.local_projection_size, dim=-1
)
# Separate convolution calls accept row-strided packed inputs
# and produce dense Q/K/V without an additional V copy.
def _prefill_conv(
x: torch.Tensor,
state: torch.Tensor,
weight: torch.Tensor,
) -> torch.Tensor:
return causal_conv1d_fn(
x.transpose(0, 1),
weight,
None,
activation="silu",
conv_states=state,
has_initial_state=has_initial_state,
cache_indices=non_spec_state_indices_tensor,
query_start_loc=non_spec_query_start_loc,
metadata=m,
).transpose(0, 1)
q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight)
k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight)
v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight)
q_ns, k_ns, v_ns = (
rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim)
for x in (q_ns, k_ns, v_ns)
)
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":
assert self.gate_lower_bound is not None
assert self._flashkda_buffer_specs is not None
workspace_out, final_state, checkpoint_state, workspace = (
current_workspace_manager().get_simultaneous(
*self._flashkda_buffer_specs
)
)
flashkda_out = (
workspace_out if has_spec_decode else core_attn_out
)[:, : q_ns.shape[1]]
if non_spec_out is not None:
flashkda_out = non_spec_out
if checkpoint is not None:
assert non_spec_query_start_loc is not None
num_sequences = initial_state.shape[0]
assert checkpoint.checkpoint_offsets.shape == (num_sequences,)
final_state = final_state[:num_sequences]
checkpoint_state = checkpoint_state[:num_sequences]
checkpoint_offsets = checkpoint.checkpoint_offsets
_flashkda_prefill(
q=q_ns,
k=k_ns,
v=v_ns,
g=g1_ns,
beta=beta_ns,
A_log=self.A_log,
dt_bias=self.dt_bias,
lower_bound=self.gate_lower_bound,
initial_state=initial_state,
cu_seqlens=non_spec_query_start_loc,
out=flashkda_out,
final_state=final_state,
workspace=workspace,
checkpoint_state=checkpoint_state,
checkpoint_offsets=checkpoint_offsets,
)
core_attn_out_non_spec = flashkda_out
last_recurrent_state = final_state
assert self._checkpoint_exporter is not None
self._checkpoint_exporter.export(
checkpoint,
raw_qkv=mixed_qkv_ns,
conv_state=conv_state,
recurrent_checkpoint=checkpoint_state,
recurrent_state=recurrent_state,
cu_seqlens=non_spec_query_start_loc,
)
else:
(
core_attn_out_non_spec,
last_recurrent_state,
) = _flashkda_prefill(
q=q_ns,
k=k_ns,
v=v_ns,
g=g1_ns,
beta=beta_ns,
A_log=self.A_log,
dt_bias=self.dt_bias,
lower_bound=self.gate_lower_bound,
initial_state=initial_state,
cu_seqlens=non_spec_query_start_loc,
out=flashkda_out,
final_state=final_state[: initial_state.shape[0]],
workspace=workspace,
)
elif self.kda_prefill_backend == "flashinfer":
assert self.gate_lower_bound is not None
assert m.flashinfer_prefill_query_start_loc is not None
if q_ns.shape[1] > initial_state.shape[0]:
assert m.flashinfer_prefill_seq_order is not None
flashinfer_out = core_attn_out[:, : q_ns.shape[1]]
if non_spec_out is not None:
flashinfer_out = non_spec_out
elif has_spec_decode:
assert self._flashinfer_kda_output_spec is not None
(workspace_out,) = current_workspace_manager().get_simultaneous(
self._flashinfer_kda_output_spec
)
flashinfer_out = workspace_out[:, : q_ns.shape[1]]
(
core_attn_out_non_spec,
last_recurrent_state,
) = _flashinfer_kda_prefill(
q=q_ns,
k=k_ns,
v=v_ns,
raw_g=g1_ns,
raw_beta=beta_ns,
A_log=self.A_log,
dt_bias=self.dt_bias,
lower_bound=self.gate_lower_bound,
initial_state=initial_state,
cu_seqlens=m.flashinfer_prefill_query_start_loc,
out=flashinfer_out,
seq_order=m.flashinfer_prefill_seq_order,
)
else:
(
core_attn_out_non_spec,
last_recurrent_state,
) = chunk_kda_with_fused_gate(
q=q_ns,
k=k_ns,
v=v_ns,
raw_g=g1_ns,
raw_beta=beta_ns,
A_log=self.A_log,
g_bias=self.dt_bias,
lower_bound=self.gate_lower_bound,
initial_state=initial_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
cu_seqlens=non_spec_query_start_loc,
out=non_spec_out,
)
recurrent_state[non_spec_state_indices_tensor] = (
last_recurrent_state.to(recurrent_state.dtype)
)
else:
# Pure non-speculative decode.
assert non_spec_state_indices_tensor is not None
decode_conv_indices = non_spec_state_indices_tensor[
: mixed_qkv_ns.size(0)
]
packed_conv_out = torch.empty_like(mixed_qkv_ns)
mixed_qkv_ns = causal_conv1d_update(
mixed_qkv_ns,
conv_state,
conv_weights,
self.conv1d.bias,
activation="silu",
conv_state_indices=decode_conv_indices,
validate_data=True,
out=packed_conv_out,
)
(
core_attn_out_non_spec,
_,
) = fused_recurrent_kda_packed_decode(
mixed_qkv=mixed_qkv_ns,
raw_g=g1_ns,
raw_beta=beta_ns,
A_log=self.A_log,
dt_bias=self.dt_bias,
lower_bound=self.gate_lower_bound,
initial_state=recurrent_state,
state_indices=decode_conv_indices,
out=non_spec_out,
)
# Restore the scheduler's original token order for mixed batches.
if core_attn_out_spec is not None and core_attn_out_non_spec is not None:
if spec_token_start is None:
assert spec_token_indx is not None
assert non_spec_token_indx 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 core_attn_out_non_spec is not None:
if (
self.kda_prefill_backend not in ("flashkda", "flashinfer")
or m.num_prefills == 0
):
# TODO: decode kernels write directly to core_attn_out
core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[
0, :num_actual_tokens
]
else:
assert core_attn_out_spec is not None
# Triton normalizes in place, so this is a self-copy with no device
# work. Keep it for the out-of-place native implementation.
core_attn_out.copy_(self.o_norm(core_attn_out, g2))