diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index d056e318ce87..eedc69166df2 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -109,12 +109,13 @@ from ...mapping import Mapping from ...models.modeling_utils import QuantAlgo, QuantConfig from ..attention_backend import AttentionMetadata -from ..distributed import AllReduce, AllReduceParams, AllReduceStrategy +from ..distributed import AllReduce, AllReduceParams from ..model_config import ModelConfig from ..modules.fused_moe import ConfigurableMoE, create_moe from ..modules.fused_moe.interface import _compute_ep_partition from ..modules.fused_moe.routing import DeepSeekV3MoeRoutingMethod from ..modules.gated_mlp import GatedMLP +from ..modules.kimi_kda import KimiKDALinearAttention from ..modules.linear import Linear as TrtllmLinear from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm @@ -129,14 +130,6 @@ os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" ) -_KDA_INDEXED_STATE_POOL_ENABLED = os.environ.get("TLLM_KDA_ENABLE_INDEXED_STATE_POOL", "1") == "1" -# Heuristic ported from SGLang's Blackwell cutoff: -# https://github.com/sgl-project/sglang/blob/e84bbf68efb683c9e2eef4168c5198042544599d/python/sglang/srt/models/kimi_k3.py#L946-L954 -# It has not been tuned for TensorRT-LLM; benchmark and retune it for TRT-LLM's -# projection kernels. Verify intentionally counts B * num_steps because those -# flattened token rows form the projection GEMMs' M dimension. -_KDA_BFA_MULTISTREAM_MAX_ROWS = 128 - # Routed-expert MoE TP/EP split overrides (read per model init, not import). # Highest precedence; either one may be set alone, the other is derived from # tp_size. Without them, an explicit moe_tensor_parallel_size / @@ -193,9 +186,10 @@ # the SM100 gate still apply. _KIMI_K3_FP8_WEIGHT_READ_MLA_ENV = "KIMI_K3_FP8_WEIGHT_READ_MLA" -# Expert override (prototype): set to "0" to drop the KimiKDARuntime decode +# Expert override (prototype): set to "0" to drop the +# KimiKDALinearAttention decode # fast path — fused qkvg and [f_a | b] projections, persistent conv staging, -# and precomputed kernel-layout constants (``_forward_decode``) — when the +# and precomputed kernel-layout constants (``forward_decode``) — when the # KDA projections are read at FP8 block-scale. With the fast path kept (the # default on an enabled master), decode issues the loader's fused FP8 # ``qkvg_proj`` GEMM for q/k/v/g plus one small BF16 GEMV for [f_a | b] @@ -821,7 +815,7 @@ def _convert_kda_projections_to_fp8_weight_read(model: nn.Module) -> int: for layer in model.layers: if not getattr(layer, "is_kda", False) or not _has_weights(layer): continue - mixer = getattr(getattr(layer, "self_attn", None), "mixer", None) + mixer = getattr(layer, "linear_attn", None) if mixer is None: continue @@ -1442,965 +1436,10 @@ def _routed_output(): # --------------------------------------------------------------------------- -# KDA runtime (pool-backed prefill / decode via the FLA kernels). +# MLA runtime. # --------------------------------------------------------------------------- -def _kda_split_conv_sections( - cs: torch.Tensor, d: int -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Split a gathered ``[N, 3D, W]`` conv-cache into contiguous q/k/v.""" - return (cs[:, :d].contiguous(), cs[:, d : 2 * d].contiguous(), cs[:, 2 * d :].contiguous()) - - -class KimiKDARuntime(nn.Module): - """Wraps the parity-tested ``KimiKDALinearAttention`` parameters with a - cache-pool-aware forward for the executor flow. - - Parameter names mirror the HF checkpoint 1:1 (the wrapped mixer is - registered under the layer as ``self_attn``, so e.g. - ``model.layers.N.self_attn.q_proj.weight`` maps identically). - """ - - def __init__( - self, - cfg, - layer_idx: int, - mapping=None, - allreduce_strategy=AllReduceStrategy.AUTO, - aux_stream: Optional[torch.cuda.Stream] = None, - ): - super().__init__() - # Lazy import: pulls in fla/einops. - from ..modules.kimi_kda.kimi_kda_mixer import KimiKDALinearAttention - - lin = cfg.linear_attn_config - self.layer_idx = layer_idx - self._use_indexed_ssm_pool = _KDA_INDEXED_STATE_POOL_ENABLED - # Attention-family TP semantics (Qwen3-Next GatedDeltaNet pattern, - # gdn_mixer.py): replicated under attention-DP — each rank runs - # its own batch with the full head set — and head-sharded across - # mapping.tp_size otherwise: every rank holds the same batch, runs - # its 1/tp head slice, and the row-sharded o_proj partials are - # all-reduced at the end of forward(). - if mapping is not None and mapping.tp_size > 1 and not mapping.enable_attention_dp: - self._kda_tp_size = mapping.tp_size - else: - self._kda_tp_size = 1 - self._kda_tp_rank = mapping.tp_rank if self._kda_tp_size > 1 else 0 - self._o_allreduce = ( - AllReduce(mapping=mapping, strategy=allreduce_strategy, dtype=torch.bfloat16) - if self._kda_tp_size > 1 - else None - ) - num_heads = lin["num_heads"] - assert num_heads % self._kda_tp_size == 0, ( - f"KDA num_heads {num_heads} not divisible by tp_size {self._kda_tp_size}" - ) - self.mixer = KimiKDALinearAttention( - hidden_size=cfg.hidden_size, - num_heads=num_heads // self._kda_tp_size, - head_dim=lin["head_dim"], - conv_kernel_size=lin["short_conv_kernel_size"], - use_full_rank_gate=lin.get("use_full_rank_gate", True), - gate_lower_bound=lin.get("gate_lower_bound", None), - rms_norm_eps=cfg.rms_norm_eps, - dtype=torch.bfloat16, - layer_idx=layer_idx, - # Use TLLM_KDA_ENABLE_OPT_PREFILL=0 to opt out of the optimized - # prefill kernel. - use_optimized_prefill=os.getenv("TLLM_KDA_ENABLE_OPT_PREFILL", "1") == "1", - use_optimized_decode=True, - ) - self.proj_size = (num_heads // self._kda_tp_size) * lin["head_dim"] - # Fused prefill/decode/verify projection weights, built after checkpoint - # load. BF16 uses separate fused [q | k | v | g] and [f_a | b] - # GEMMs; FP8 supplies qkvg through the mixer's fused projection and - # reuses the BF16 [f_a | b] weight. - self._qkvg_proj_weight: Optional[torch.Tensor] = None - self._bfa_proj_weight: Optional[torch.Tensor] = None - self._w_q_t = self._w_k_t = self._w_v_t = None - self._A_log_f32 = self._dt_bias_f32 = self._onorm_w_f32 = None - # Fork/join state for overlapping the small [f_a | b] -> f_b chain - # with the wide qkvg projection during CUDA-graph execution. - self._projection_aux_stream = aux_stream - self._projection_fork_event = torch.cuda.Event() - self._projection_join_event = torch.cuda.Event() - # Persistent batch-row-dense staging for the fused decode kernel's - # per-section conv windows. Sized once, on the first decode call, - # to the conv pool's slot count and never reallocated (see - # ``_forward_decode``). - self._cs_dense: Optional[torch.Tensor] = None - # fp32 [dim, W] conv weights for the fused verify kernel, prebuilt - # by ``_build_mtp_conv_weights()`` at weight-load finalize time. - self._mtp_conv_weights: Optional[Tuple[torch.Tensor, ...]] = None - - def finalize_decode_weights(self) -> None: - """Build fused projection weights and decode constants after weight load. - - 1. Separate fused ``[q | k | v | g]`` and ``[f_a | b]`` projections. - Keeping the wide qkvg output aligned avoids degrading its GEMM - kernel selection with the small f_a and b tails. Source parameters - are repointed to row views of the fused buffers, so prefill and - verify paths keep using them without duplicate weight storage. - 2. Kernel-layout constants that ``_decode_via_optimized`` used to - rebuild with ~6 device kernels per layer per decode step: - transposed conv weights (bf16 ``[W, D]``) and fp32 copies of - ``A_log`` / ``dt_bias`` / ``o_norm.weight``. - """ - mixer = self.mixer - if mixer._dispatch.decode_kernel_path != "optimized" or not mixer.use_full_rank_gate: - return - if mixer.q_proj.weight.device.type != "cuda": - return - with torch.no_grad(): - qkvg_modules = ( - mixer.q_proj, - mixer.k_proj, - mixer.v_proj, - mixer.g_proj, - ) - qkvg_weight = self._merge_projection_weights(qkvg_modules) - # Eight BF16 outputs occupy 16 bytes, so padding keeps each output row - # aligned for vectorized f_b consumption; it is not a kernel requirement. - bfa_weight = self._merge_projection_weights( - (mixer.f_a_proj, mixer.b_proj), pad_rows_to=8 - ) - self._build_decode_kernel_constants() - self._bfa_proj_weight = bfa_weight - # Publish last: both weights are required by the BF16 fast path. - self._qkvg_proj_weight = qkvg_weight - - @staticmethod - def _merge_projection_weights( - modules: tuple[nn.Linear, ...], pad_rows_to: int = 1 - ) -> torch.Tensor: - """Concatenate linear weights and repoint the modules to row views.""" - weights = [module.weight.data for module in modules] - padding = (-sum(weight.shape[0] for weight in weights)) % pad_rows_to - if padding: - weights.append(weights[0].new_zeros((padding, weights[0].shape[1]))) - fused = torch.cat(weights, dim=0).contiguous() - offset = 0 - for module in modules: - rows = module.weight.shape[0] - module.weight.data = fused[offset : offset + rows] - offset += rows - return fused - - def _build_decode_kernel_constants(self) -> None: - """Kernel-layout constants shared by both finalize variants.""" - mixer = self.mixer - self._w_q_t = ( - mixer.q_conv1d.weight.detach() - .squeeze(1) - .transpose(0, 1) - .to(torch.bfloat16) - .contiguous() - ) - self._w_k_t = ( - mixer.k_conv1d.weight.detach() - .squeeze(1) - .transpose(0, 1) - .to(torch.bfloat16) - .contiguous() - ) - self._w_v_t = ( - mixer.v_conv1d.weight.detach() - .squeeze(1) - .transpose(0, 1) - .to(torch.bfloat16) - .contiguous() - ) - self._A_log_f32 = mixer.A_log.detach().float().contiguous() - self._dt_bias_f32 = mixer.dt_bias.detach().float().contiguous() - self._onorm_w_f32 = mixer.o_norm.weight.detach().float().contiguous() - # Build the fused-verify conv constants eagerly too, so the first - # verify call never allocates (a capture-unsafe lazy allocation). - self._build_mtp_conv_weights() - - def finalize_decode_weights_fp8(self) -> None: - """FP8 counterpart of ``finalize_decode_weights()``. - - Runs AFTER ``_convert_kda_projections_to_fp8_weight_read``, so - q/k/v/g already live in the mixer's fused FP8 ``qkvg_proj`` GEMM. - Only the two small BF16 projections reading the same hidden — - ``f_a_proj`` and ``b_proj`` (kept BF16 by the FP8 conversion: outputs - are not 128-multiples and feed the accuracy-sensitive recurrent - decay) — are fused here into one ``[f_a | b]`` weight, with the source - parameters repointed to row views. Prefill, decode, and verification - then share both fused projections; the kernel-layout constants are - decode-only. - """ - mixer = self.mixer - if mixer._dispatch.decode_kernel_path != "optimized" or not mixer.use_full_rank_gate: - return - fused_qkvg = getattr(mixer, "qkvg_proj", None) - split_sizes = getattr(mixer, "qkvg_split_sizes", None) - if fused_qkvg is None or split_sizes is None or len(split_sizes) != 4: - return - if mixer.f_a_proj.weight.device.type != "cuda": - return - with torch.no_grad(): - bfa_weight = self._merge_projection_weights( - (mixer.f_a_proj, mixer.b_proj), pad_rows_to=8 - ) - self._build_decode_kernel_constants() - # Publish last: enables fused [f_a | b] in prefill/decode/verify. - self._bfa_proj_weight = bfa_weight - - def forward( - self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata - ) -> torch.Tensor: - """``hidden_states``: flattened ``[num_tokens, hidden]`` (ctx tokens - first, then one token per generation request).""" - mamba_metadata = attn_metadata.mamba_metadata - num_prefills = attn_metadata.num_contexts - num_ctx_tokens = attn_metadata.num_ctx_tokens - batch_size = attn_metadata.seq_lens.shape[0] - # index_copy_/index_select need int64 indices; the int64 mirror is - # prepared once per step by Mamba2Metadata.prepare() so KDA layers - # do not each replay an int32->int64 cast inside the decode graph. - state_indices = getattr(mamba_metadata, "state_indices_long", None) - if state_indices is None or state_indices.shape[0] != batch_size: - state_indices = mamba_metadata.state_indices[:batch_size].long() - cu_seqlens = mamba_metadata.query_start_loc_long[: num_prefills + 1] - num_decodes = batch_size - num_prefills - - layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) - conv_pool = layer_cache.conv # [slots, 3D, W] bf16 - ssm_pool = layer_cache.temporal # [slots, H, V, K] fp32 - - outputs: List[torch.Tensor] = [] - if num_prefills > 0: - outputs.append( - self._forward_prefill( - hidden_states[:num_ctx_tokens], - cu_seqlens, - mamba_metadata, - num_prefills, - conv_pool, - ssm_pool, - state_indices[:num_prefills], - layer_cache, - ) - ) - if num_decodes > 0: - decode_rows = hidden_states.shape[0] - num_ctx_tokens - if decode_rows == num_decodes: - outputs.append( - self._forward_decode( - hidden_states[num_ctx_tokens:], - conv_pool, - ssm_pool, - state_indices[num_prefills:], - mamba_metadata, - layer_cache, - ssm_state_indices=( - mamba_metadata.state_indices[num_prefills:batch_size] - if self._use_indexed_ssm_pool - else None - ), - ) - ) - else: - # Speculative verification: each generation request carries - # 1 + draft_len tokens (drafts are padded to the static max, - # so T is uniform). Per-step states go to the manager's - # SpeculativeState scratch buffers — never the live pools — - # and kv_cache_manager.update_mamba_states() promotes the - # accepted step after sampling. - assert decode_rows % num_decodes == 0, ( - f"ragged generation batch: {decode_rows} tokens for {num_decodes} requests" - ) - outputs.append( - self._forward_verify( - hidden_states[num_ctx_tokens:], - decode_rows // num_decodes, - layer_cache, - conv_pool, - ssm_pool, - state_indices[num_prefills:], - ) - ) - out = outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0) - if self._o_allreduce is not None: - # Head-sharded TP: every rank ran its head shard on the same - # local batch; sum the row-sharded o_proj partials. - out = self._o_allreduce(out) - return out - - def _has_kda_replay_caches(self, layer_cache) -> bool: - """True when the manager allocated the fused-verify replay caches.""" - return getattr(layer_cache, "kda_qkg_cache", None) is not None - - def _sync_kda_replay_conv_window( - self, layer_cache, slot_indices, conv_q, conv_k, conv_v - ) -> None: - """Seed the replay conv caches' committed window from FLA windows. - - The fused verify kernel keeps its own extended fp32 dim-contiguous - conv caches; their committed window (columns ``[0, W-1)``) must hold - the last ``W-1`` raw conv inputs whenever another path (prefill, - plain decode) advances the base conv pool. The FLA window's oldest - column drops out of every future convolution, so columns ``[1, W)`` - of the FLA cache map 1:1 onto the committed window. - """ - if not self._has_kda_replay_caches(layer_cache): - return - w = self.mixer.conv_size - for cache, window in ( - (layer_cache.kda_conv_q, conv_q), - (layer_cache.kda_conv_k, conv_k), - (layer_cache.kda_conv_v, conv_v), - ): - cache[:, :, : w - 1].index_copy_(0, slot_indices, window[:, :, 1:].to(cache.dtype)) - - def _forward_prefill( - self, - x2d, - cu_seqlens, - mamba_metadata, - num_prefills, - conv_pool, - ssm_pool, - slot_indices, - layer_cache=None, - ) -> torch.Tensor: - from einops import rearrange - - mixer = self.mixer - d = self.proj_size - x = x2d.unsqueeze(0) # [1, T, hidden] - - onorm_g = None - if self._qkvg_proj_weight is not None: - qkvg = torch.nn.functional.linear(x, self._qkvg_proj_weight) - q_proj_states, k_proj_states, v_proj_states = qkvg[..., : 3 * d].split(d, dim=-1) - onorm_g = qkvg[..., 3 * d : 4 * d] - else: - fused_qkvg = getattr(mixer, "qkvg_proj", None) - if fused_qkvg is not None: - qkvg = fused_qkvg(x) - q_proj_states, k_proj_states, v_proj_states = qkvg[..., : 3 * d].split(d, dim=-1) - qkvg_split_sizes = getattr(mixer, "qkvg_split_sizes", None) - if ( - mixer.use_full_rank_gate - and qkvg_split_sizes is not None - and len(qkvg_split_sizes) == 4 - ): - onorm_g = qkvg[..., 3 * d : 4 * d] - else: - q_proj_states = mixer.q_proj(x) - k_proj_states = mixer.k_proj(x) - v_proj_states = mixer.v_proj(x) - - # Initial states: present for continuation chunks (chunked prefill) - # and for prefix-cache hits (block reuse), where the previous - # conv/recurrent state was onboarded into this request's slot. - conv_q_in = conv_k_in = conv_v_in = None - recurrent_in = None - if mamba_metadata.use_initial_states: - has_init = mamba_metadata.has_initial_states[:num_prefills] - cs = conv_pool.index_select(0, slot_indices) - cs[~has_init] = 0 - conv_q_in, conv_k_in, conv_v_in = _kda_split_conv_sections(cs, d) - recurrent_in = ssm_pool.index_select(0, slot_indices) - recurrent_in[~has_init] = 0 - - q, conv_q = mixer.q_conv1d( - q_proj_states, cache=conv_q_in, output_final_state=True, cu_seqlens=cu_seqlens - ) - k, conv_k = mixer.k_conv1d( - k_proj_states, cache=conv_k_in, output_final_state=True, cu_seqlens=cu_seqlens - ) - v, conv_v = mixer.v_conv1d( - v_proj_states, cache=conv_v_in, output_final_state=True, cu_seqlens=cu_seqlens - ) - - if self._bfa_proj_weight is not None: - bfa = torch.nn.functional.linear(x, self._bfa_proj_weight) - f_a = bfa[..., : mixer.head_dim] - beta = bfa[..., mixer.head_dim : mixer.head_dim + mixer.num_heads].float() - g = mixer.f_b_proj(f_a) - else: - g = mixer.f_b_proj(mixer.f_a_proj(x)) - beta = mixer.b_proj(x).float() - g = rearrange(g, "... (h d) -> ... h d", d=mixer.head_dim) - - q = rearrange(q, "... (h d) -> ... h d", d=mixer.head_k_dim) - k = rearrange(k, "... (h d) -> ... h d", d=mixer.head_k_dim) - v = rearrange(v, "... (h d) -> ... h d", d=mixer.head_dim) - - # Kernel dispatch (in-tree trtllm::kda_prefill or FLA chunk_kda). - # Both paths exchange states in the pool's V-first [N, H, V, K] - # layout, so recurrent_in / final_state map to ssm_pool 1:1. - lower_bound = mixer.gate_lower_bound - o, final_state = mixer.prefill_chunk_kda( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A_log=mixer.A_log, - dt_bias=mixer.dt_bias, - scale=mixer.head_k_dim**-0.5, - initial_state=recurrent_in, - safe_gate=lower_bound is not None, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - ) - - # Persist per-request states into the pools. - conv_pool.index_copy_( - 0, slot_indices, torch.cat([conv_q, conv_k, conv_v], dim=1).to(conv_pool.dtype) - ) - ssm_pool.index_copy_(0, slot_indices, final_state.to(ssm_pool.dtype)) - # Fused-verify replay caches: seed the committed conv window so the - # first verify round convolves the correct history (pending drafts - # are zero for a fresh request, so the tail columns are unused). - self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_q, conv_k, conv_v) - - return self._output_gate_and_proj(x, o, onorm_g) - - def _forward_decode( - self, - x2d, - conv_pool, - ssm_pool, - slot_indices, - mamba_metadata=None, - layer_cache=None, - ssm_state_indices=None, - ) -> torch.Tensor: - """Plain T=1 decode, fast path. - - Calls ``trtllm::kda_decode`` directly with kernel-native layouts - (nsys 07-24: the reference path spent ~70 us/layer on glue around - the 5 us kernel — 6 separate in-projection GEMV pairs, per-step - re-transposition of constant weights, conv-window slice/roll - copies, per-call torch.arange defaults, and redundant dtype - casts): - - * one wide fused qkvg GEMV on the main stream, overlapped with the - fused [f_a | b] GEMV and f_b GEMV on the auxiliary stream for - CUDA-graph batches up to 128 tokens; - * conv windows staged with one gather + one repack copy into a - persistent dense per-section buffer; - * conv-pool write-back with one cat + one index_copy_; - * constant tensors (transposed conv weights, fp32 A_log/dt_bias/ - o_norm weight) reused instead of rebuilt per step. - - The conv windows remain gathered batch-row-dense. When stable - int32 slot indices are supplied, the recurrent-state pool is passed - directly and the CUDA wrapper selects its indexed-state launch; - otherwise the state uses the batch-row-dense static layout. - """ - mixer = self.mixer - if mixer.decode_kernel_path != "optimized" or mixer.wrong_state_layout: - ssm_state_indices = None - if ssm_state_indices is not None: - logger.info_once( - "Kimi K3 KDA indexed recurrent-state pool path is active", - key="kimi_k3_kda_indexed_state_pool", - ) - else: - logger.info_once( - "Kimi K3 KDA static recurrent-state path is active", key="kimi_k3_kda_static_state" - ) - has_qkvg_projection = ( - self._qkvg_proj_weight is not None or getattr(mixer, "qkvg_proj", None) is not None - ) - if ( - not has_qkvg_projection - or self._bfa_proj_weight is None - or mamba_metadata is None - or ssm_pool.dtype != torch.float32 - ): - return self._forward_decode_ref( - x2d, conv_pool, ssm_pool, slot_indices, layer_cache, ssm_state_indices - ) - - d = self.proj_size - hd = mixer.head_dim - H = mixer.num_heads - B = x2d.shape[0] - W = mixer.conv_size - - # Allocated ONCE at the pool slot count (== per-rank max batch on - # the Mixed manager; ``slot_indices`` are distinct pool rows and - # this is the plain one-token-per-request path, so B never exceeds - # it) and never reallocated: captured CUDA graphs hold this - # pointer, so a realloc would leave earlier graphs writing into - # freed memory. Footprint: slots x ~9(H=6)..222(H=96) KB per layer. - buf = self._cs_dense - if buf is None: - if torch.cuda.is_current_stream_capturing(): - # Never allocate inside CUDA graph capture; the reference - # path is capture-safe (just slower). - return self._forward_decode_ref( - x2d, conv_pool, ssm_pool, slot_indices, layer_cache, ssm_state_indices - ) - buf = torch.empty( - 3, max(conv_pool.shape[0], B), d, W - 1, dtype=torch.bfloat16, device=x2d.device - ) - self._cs_dense = buf - else: - # Fail loudly if the sizing invariant ever breaks: silently - # reallocating here would hand previously captured CUDA graphs - # a dangling pointer. - assert buf.shape[1] >= B, ( - f"KDA decode staging buffer holds {buf.shape[1]} rows but the " - f"decode batch is {B}; reallocating would corrupt previously " - f"captured CUDA graphs" - ) - - def _project_qkvg() -> torch.Tensor: - if self._qkvg_proj_weight is not None: - return torch.nn.functional.linear(x2d, self._qkvg_proj_weight) - # FP8 weight read (KIMI_K3_KDA_GLUE_FP8=1) uses the loader's - # fused FP8 [q | k | v | g] GEMM. - return mixer.qkvg_proj(x2d) - - def _project_bfa_and_fb() -> tuple[torch.Tensor, torch.Tensor]: - bfa = torch.nn.functional.linear(x2d, self._bfa_proj_weight) - f_a = bfa[:, :hd] - beta = bfa[:, hd : hd + H] - return beta, mixer.f_b_proj(f_a) - - projection_aux_stream = ( - self._projection_aux_stream if B <= _KDA_BFA_MULTISTREAM_MAX_ROWS else None - ) - qkvg, (beta, g) = maybe_execute_in_parallel( - _project_qkvg, - _project_bfa_and_fb, - self._projection_fork_event, - self._projection_join_event, - projection_aux_stream, - disable_on_compile=True, - ) - x_qkv = qkvg[:, : 3 * d] - onorm_g = qkvg[:, 3 * d : 4 * d] - - # Gather the HF-layout conv windows once, then repack the - # historical W-1 columns into the kernel's dense per-section - # [B, d, W-1] layout (single strided copy kernel). - cs = conv_pool.index_select(0, slot_indices) # [B, 3d, W] - cs_dense = buf[:, :B] - cs_dense.copy_(cs.view(B, 3, d, W)[:, :, :, 1:].permute(1, 0, 2, 3)) - - state = ( - ssm_pool if ssm_state_indices is not None else ssm_pool.index_select(0, slot_indices) - ) - - o = mixer._dispatch.decode_kda( - x_q=x_qkv[:, :d].unflatten(-1, (H, hd)).unsqueeze(0), - x_k=x_qkv[:, d : 2 * d].unflatten(-1, (H, hd)).unsqueeze(0), - x_v=x_qkv[:, 2 * d :].unflatten(-1, (H, hd)).unsqueeze(0), - w_q_t=self._w_q_t, - w_k_t=self._w_k_t, - w_v_t=self._w_v_t, - bias_q=None, - bias_k=None, - bias_v=None, - cs_q=cs_dense[0], - cs_k=cs_dense[1], - cs_v=cs_dense[2], - A_log=self._A_log_f32, - g=g.unflatten(-1, (H, hd)).unsqueeze(0), - dt_bias=self._dt_bias_f32, - beta=beta.unsqueeze(0), - state=state, - onorm_g=onorm_g.unflatten(-1, (H, hd)).unsqueeze(0), - onorm_weight=self._onorm_w_f32, - out=None, - ssm_state_indices=ssm_state_indices, - cu_seqlens=mamba_metadata._arange_buffer[: B + 1], - scale=hd**-0.5, - onorm_eps=mixer.o_norm.eps, - lower_bound=mixer.gate_lower_bound, - use_beta_sigmoid_in_kernel=True, - verbose=False, - update_conv_cache=False, - ) - if ssm_state_indices is None: - ssm_pool.index_copy_(0, slot_indices, state) - - # Roll the HF-layout conv pool by one token: new window = - # [old columns 1..W-1, x_new]. One cat + one scatter. - new_win = torch.cat([cs[:, :, 1:], x_qkv.unsqueeze(-1)], dim=-1) - if new_win.dtype != conv_pool.dtype: - new_win = new_win.to(conv_pool.dtype) - conv_pool.index_copy_(0, slot_indices, new_win) - # Fused-verify replay caches (spec decoding only): keep the - # committed conv window in sync with the plain-decode advance. - self._sync_kda_replay_conv_window( - layer_cache, slot_indices, new_win[:, :d], new_win[:, d : 2 * d], new_win[:, 2 * d :] - ) - - return mixer.o_proj(o.view(B, d)) - - def _forward_decode_ref( - self, x2d, conv_pool, ssm_pool, slot_indices, layer_cache=None, ssm_state_indices=None - ) -> torch.Tensor: - from ..modules.kimi_kda.kimi_kda_mixer import KimiKDACachedState - - mixer = self.mixer - d = self.proj_size - x = x2d.unsqueeze(1) # [B, 1, hidden] - - cs = conv_pool.index_select(0, slot_indices) - conv_q, conv_k, conv_v = _kda_split_conv_sections(cs, d) - cache = KimiKDACachedState( - conv_state_q=conv_q, - conv_state_k=conv_k, - conv_state_v=conv_v, - recurrent_state=( - ssm_pool - if ssm_state_indices is not None - else ssm_pool.index_select(0, slot_indices) - ), - ) - out, new_cache = mixer.forward_decode( - x, - cache, - ssm_state_indices=ssm_state_indices, - ) - - conv_pool.index_copy_( - 0, - slot_indices, - torch.cat( - [ - new_cache.conv_state_q, - new_cache.conv_state_k, - new_cache.conv_state_v, - ], - dim=1, - ).to(conv_pool.dtype), - ) - if ssm_state_indices is None: - ssm_pool.index_copy_(0, slot_indices, new_cache.recurrent_state.to(ssm_pool.dtype)) - # Fused-verify replay caches: keep the committed conv window in - # sync with the plain-decode advance. NOTE: this path is only - # correct for requests with no pending accepted drafts - # (prev_num_accepted_tokens == 0); with drafts pending, the live - # pools lag by the pending prefix and only the fused verify kernel - # can advance them. The spec workers pad drafts to the static max, - # so drafted batches always take the verify path. - self._sync_kda_replay_conv_window( - layer_cache, - slot_indices, - new_cache.conv_state_q, - new_cache.conv_state_k, - new_cache.conv_state_v, - ) - - return out.squeeze(1) - - def _forward_verify( - self, x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices - ) -> torch.Tensor: - """Speculative verification: advance each request ``num_steps`` - tokens (1 golden + ``num_steps - 1`` padded drafts). - - Two paths: - - * Fused (``trtllm::kda_mtp_decode``, when the manager allocated the - KDA replay caches): one kernel launch replays the previous - round's accepted drafts from the per-slot replay caches, then - processes the new tokens, committing the recurrent state and conv - windows **in place** after the golden token and caching the new - drafts. ``update_mamba_states()`` afterwards only records the - accepted count for the next round's replay. - * Legacy (sequential per-step FLA): per-step states go to the - manager's batch-row-indexed intermediate scratch buffers and - ``update_mamba_states()`` promotes the accepted step's state - after sampling. - """ - if self._has_kda_replay_caches(layer_cache): - assert self.mixer.verify_kernel_path == "optimized", ( - "KDA replay caches are allocated but the fused verify " - "kernel is unavailable; the legacy intermediate buffers " - "were not allocated so there is no fallback" - ) - return self._forward_verify_fused(x2d, num_steps, layer_cache, ssm_pool, slot_indices) - return self._forward_verify_sequential( - x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices - ) - - def _project_verify_inputs( - self, x: torch.Tensor, num_rows: int - ) -> Optional[ - tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - Optional[torch.Tensor], - ] - ]: - """Project fused QKVG and [f_a | b] inputs for target verification.""" - mixer = self.mixer - qkvg_weight = self._qkvg_proj_weight - fused_qkvg = getattr(mixer, "qkvg_proj", None) - if qkvg_weight is None and fused_qkvg is None: - return None - - def _project_qkvg() -> torch.Tensor: - if qkvg_weight is not None: - return torch.nn.functional.linear(x, qkvg_weight) - return fused_qkvg(x) - - bfa_weight = self._bfa_proj_weight - if bfa_weight is not None: - - def _project_bfa_and_fb() -> tuple[torch.Tensor, torch.Tensor]: - bfa = torch.nn.functional.linear(x, bfa_weight) - f_a = bfa[..., : mixer.head_dim] - beta = bfa[..., mixer.head_dim : mixer.head_dim + mixer.num_heads] - return beta, mixer.f_b_proj(f_a) - - projection_aux_stream = ( - self._projection_aux_stream - if 0 < num_rows <= _KDA_BFA_MULTISTREAM_MAX_ROWS - else None - ) - qkvg, (beta, forget_gate) = maybe_execute_in_parallel( - _project_qkvg, - _project_bfa_and_fb, - self._projection_fork_event, - self._projection_join_event, - projection_aux_stream, - disable_on_compile=True, - ) - else: - qkvg = _project_qkvg() - beta = mixer.b_proj(x) - forget_gate = mixer.f_b_proj(mixer.f_a_proj(x)) - - d = self.proj_size - q_proj, k_proj, v_proj = (part.contiguous() for part in qkvg[..., : 3 * d].split(d, dim=-1)) - qkvg_split_sizes = getattr(mixer, "qkvg_split_sizes", None) - has_onorm_gate = qkvg_weight is not None or ( - mixer.use_full_rank_gate and qkvg_split_sizes is not None and len(qkvg_split_sizes) == 4 - ) - onorm_g = qkvg[..., 3 * d : 4 * d].contiguous() if has_onorm_gate else None - return q_proj, k_proj, v_proj, forget_gate, beta, onorm_g - - def _forward_verify_fused( - self, x2d, num_steps, layer_cache, ssm_pool, slot_indices - ) -> torch.Tensor: - """Fused multi-token verify via ``trtllm::kda_mtp_decode``. - - Token layout: the kernel indexes each request's new tokens at - ``cu_seqlens[n] + num_accepted[n] + i``. The runtime packs the - ``num_steps`` new tokens per request contiguously, so we pass - ``cu_seqlens[n] = n * num_steps - num_accepted[n]`` — the shift - lands the kernel's reads/writes exactly on the packed rows. A - negative entry for request 0 is fine: ``bos`` is only ever used - additively with a token offset ``>= num_accepted``. - """ - mixer = self.mixer - num_decodes = x2d.shape[0] // num_steps - num_spec = num_steps - 1 - H = mixer.num_heads - K = mixer.head_k_dim - x = x2d.view(num_decodes, num_steps, -1) # [B, T, hidden] - T_total = num_decodes * num_steps - - projections = self._project_verify_inputs(x, T_total) - if projections is None: - q_proj = mixer.q_proj(x) - k_proj = mixer.k_proj(x) - v_proj = mixer.v_proj(x) - forget_gate = mixer.f_b_proj(mixer.f_a_proj(x)) - beta_proj = mixer.b_proj(x) - onorm_g = None - else: - q_proj, k_proj, v_proj, forget_gate, beta_proj, onorm_g = projections - x_q = q_proj.view(1, T_total, H, K) - x_k = k_proj.view(1, T_total, H, K) - x_v = v_proj.view(1, T_total, H, mixer.head_dim) - # Raw gate / beta: the kernel applies dt_bias, A_log, the - # lower-bound sigmoid gate, and the beta sigmoid itself. - g = forget_gate.view(1, T_total, H, K) - beta = beta_proj.contiguous().view(1, T_total, H) - - w_q, w_k, w_v = self._get_mtp_conv_weights() - lower_bound = ( - mixer.gate_lower_bound_override - if mixer.gate_lower_bound_override is not None - else mixer.gate_lower_bound - ) - - pending = layer_cache.prev_num_accepted_tokens[slot_indices].to( - torch.int32 - ) # accepted drafts of the previous round, per req - cu_seqlens = torch.arange( - 0, (num_decodes + 1) * num_steps, num_steps, dtype=torch.int32, device=x2d.device - ) - cu_seqlens[:num_decodes].sub_(pending) - - out = mixer._dispatch.mtp_verify( - x_q=x_q, - x_k=x_k, - x_v=x_v, - w_q=w_q, - w_k=w_k, - w_v=w_v, - cs_q=layer_cache.kda_conv_q, - cs_k=layer_cache.kda_conv_k, - cs_v=layer_cache.kda_conv_v, - g=g, - beta=beta, - # .detach(): the CuTe DSL DLPack bridge rejects grad-tracking - # tensors. - A_log=mixer.A_log.detach(), - dt_bias=mixer.dt_bias.detach(), - recurrent_state=ssm_pool, - qkg_cache=layer_cache.kda_qkg_cache, - v_cache=layer_cache.kda_v_cache, - beta_cache=layer_cache.kda_beta_cache, - ssm_state_indices=slot_indices.to(torch.int32), - cu_seqlens=cu_seqlens, - num_spec=num_spec, - num_accepted_tokens=pending, - lower_bound=lower_bound, - scale=mixer.head_k_dim**-0.5, - ) - o = out.view(num_decodes, num_steps, H, mixer.head_dim) - return self._output_gate_and_proj(x, o, onorm_g) - - def _build_mtp_conv_weights(self) -> None: - """Prebuild the fp32 ``[dim, W]`` conv weights for the fused verify - kernel (once, at weight-load finalize time). Building them lazily - at first use would allocate at runtime; under CUDA graph capture - that bakes capture-pool pointers into the cached tuple.""" - mixer = self.mixer - self._mtp_conv_weights = tuple( - conv.weight.detach().squeeze(1).float().contiguous() - for conv in (mixer.q_conv1d, mixer.k_conv1d, mixer.v_conv1d) - ) - - def _get_mtp_conv_weights(self) -> Tuple[torch.Tensor, ...]: - """fp32 ``[dim, W]`` conv weights for the fused verify kernel, - prebuilt by ``_build_mtp_conv_weights()``.""" - cached = self._mtp_conv_weights - if cached is None: - raise RuntimeError( - "Kimi K3 fused-verify conv weights were not prebuilt; call " - "_build_mtp_conv_weights() (done by load_weights() and by " - "finalize_decode_weights() / finalize_decode_weights_fp8()) " - "after weight load and before the first verify step." - ) - return cached - - def _forward_verify_sequential( - self, x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices - ) -> torch.Tensor: - """Sequential per-step FLA verification (legacy intermediate-buffer - path). Live pools are read-only here; ``update_mamba_states()`` - commits the accepted step's state after sampling. - """ - from einops import rearrange - from fla.ops.kda import fused_recurrent_kda - - intermediate_conv = layer_cache.intermediate_conv_window - intermediate_ssm = layer_cache.intermediate_ssm - assert intermediate_conv is not None and intermediate_ssm is not None, ( - "speculative verification requires the cache manager's " - "SpeculativeState (legacy intermediate-buffer path)" - ) - - mixer = self.mixer - d = self.proj_size - num_decodes = x2d.shape[0] // num_steps - x = x2d.view(num_decodes, num_steps, -1) # [B, T, hidden] - - projections = self._project_verify_inputs(x, x2d.shape[0]) - if projections is None: - q_proj_states = mixer.q_proj(x) - k_proj_states = mixer.k_proj(x) - v_proj_states = mixer.v_proj(x) - g = mixer.f_b_proj(mixer.f_a_proj(x)) - beta = mixer.b_proj(x).float() - onorm_g = None - else: - q_proj_states, k_proj_states, v_proj_states, g, beta, onorm_g = projections - beta = beta.float() - g = rearrange(g, "... (h d) -> ... h d", d=mixer.head_dim) - - # Gathered copies — mutated across steps, never written back to the - # live pools. - cs = conv_pool.index_select(0, slot_indices) - conv_q, conv_k, conv_v = _kda_split_conv_sections(cs, d) - state = ssm_pool.index_select(0, slot_indices) - - step_outputs: List[torch.Tensor] = [] - for t in range(num_steps): - # ShortConvolution.step updates the (gathered) caches in place. - q_t, conv_q = mixer.q_conv1d( - q_proj_states[:, t : t + 1], cache=conv_q, output_final_state=True - ) - k_t, conv_k = mixer.k_conv1d( - k_proj_states[:, t : t + 1], cache=conv_k, output_final_state=True - ) - v_t, conv_v = mixer.v_conv1d( - v_proj_states[:, t : t + 1], cache=conv_v, output_final_state=True - ) - - q_t = rearrange(q_t, "... (h d) -> ... h d", d=mixer.head_k_dim) - k_t = rearrange(k_t, "... (h d) -> ... h d", d=mixer.head_k_dim) - v_t = rearrange(v_t, "... (h d) -> ... h d", d=mixer.head_dim) - - o_t, state = fused_recurrent_kda( - q=q_t, - k=k_t, - v=v_t, - g=g[:, t : t + 1], - beta=beta[:, t : t + 1], - A_log=mixer.A_log, - dt_bias=mixer.dt_bias, - initial_state=state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=True, - lower_bound=mixer.gate_lower_bound, - state_v_first=True, - ) - step_outputs.append(o_t) - - # Batch-row indexed ([:num_decodes] prefix), matching - # update_mamba_states()'s intermediate_state_indices. - intermediate_conv[:num_decodes, t] = torch.cat([conv_q, conv_k, conv_v], dim=1).to( - intermediate_conv.dtype - ) - intermediate_ssm[:num_decodes, t] = state.to(intermediate_ssm.dtype) - - o = torch.cat(step_outputs, dim=1) # [B, T, H, V] - return self._output_gate_and_proj(x, o, onorm_g) - - def _output_gate_and_proj( - self, x: torch.Tensor, o: torch.Tensor, onorm_g: Optional[torch.Tensor] = None - ) -> torch.Tensor: - from einops import rearrange - - mixer = self.mixer - if onorm_g is not None: - g_out = onorm_g - elif mixer.use_full_rank_gate: - g_out = mixer.g_proj(x) - else: - g_out = mixer.g_b_proj(mixer.g_a_proj(x)) - g_out = rearrange(g_out, "... (h d) -> ... h d", d=mixer.head_dim) - o = mixer.o_norm(o, g_out) - o = rearrange(o, "b t h d -> (b t) (h d)") - return mixer.o_proj(o) - - class KimiMLARuntime(nn.Module): """Wraps K3 MLA and applies its external TP output reduction.""" @@ -2486,7 +1525,7 @@ def __init__( raise ValueError(f"Kimi K3 layer {layer_idx} must be exactly one of KDA/MLA") if self.is_kda: - self.self_attn = KimiKDARuntime( + self.linear_attn = KimiKDALinearAttention( cfg, layer_idx, mapping=model_config.mapping, @@ -2611,7 +1650,10 @@ def forward( prefix_sum = None hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn(hidden_states, attn_metadata) + if self.is_kda: + hidden_states = self.linear_attn(hidden_states, attn_metadata) + else: + hidden_states = self.self_attn(hidden_states, attn_metadata) if prefix_sum is not None: prefix_sum = prefix_sum + hidden_states @@ -2954,9 +1996,12 @@ def checkpoint_name_plan( if name == "lm_head.weight": ckpt_key = prefix + "lm_head.weight" else: - # Runtime wrapper modules hold the parity-tested mixers as a - # "mixer" submodule; the checkpoint names have no such scope. - ckpt_key = prefix + name.replace(".self_attn.mixer.", ".self_attn.") + if ".linear_attn." in name: + ckpt_key = prefix + name.replace(".linear_attn.", ".self_attn.") + else: + # MLA retains its runtime/mixer hierarchy; the checkpoint + # has no intermediate ``mixer`` scope. + ckpt_key = prefix + name.replace(".self_attn.mixer.", ".self_attn.") name_map[name] = ckpt_key if name.endswith(_GATE_UP_FUSED_SUFFIX): # Fused [gate | up] MLP layout (dense mlp / shared_experts): @@ -3065,8 +2110,8 @@ def _load_trunk_params( kda_tp_size, kda_tp_rank = 1, 0 for layer in self.model.layers: if getattr(layer, "is_kda", False): - kda_tp_size = layer.self_attn._kda_tp_size - kda_tp_rank = layer.self_attn._kda_tp_rank + kda_tp_size = layer.linear_attn._kda_tp_size + kda_tp_rank = layer.linear_attn._kda_tp_rank break def load_param(name: str, param: torch.nn.Parameter): @@ -3168,7 +2213,7 @@ def load_param(name: str, param: torch.nn.Parameter): # weights on dim 0 (rows), o_proj on dim 1 (columns). # MLA head-sharded projections were handled by parameter # identity above, so shape ratios identify the KDA slices. - if kda_tp_size > 1 and ".self_attn." in name: + if kda_tp_size > 1 and ".linear_attn." in name: if ( src.shape[0] == param.shape[0] * kda_tp_size and src.shape[1:] == param.shape[1:] @@ -3483,19 +2528,19 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: for layer in self.model.layers: if getattr(layer, "is_kda", False) and _has_weights(layer): if not kda_fp8: - layer.self_attn.finalize_decode_weights() + layer.linear_attn.finalize_decode_weights() num_kda_fused += int( - layer.self_attn._qkvg_proj_weight is not None - and layer.self_attn._bfa_proj_weight is not None + layer.linear_attn._qkvg_proj_weight is not None + and layer.linear_attn._bfa_proj_weight is not None ) # The fused-verify conv constants are needed on every - # configuration that can reach _forward_verify_fused, + # configuration that can reach forward_verify_fused, # including ones where neither finalize variant runs (e.g. # FP8 KDA weight read with the fused decode glue disabled), # and are never computed lazily (a first verify under CUDA # graph capture must not allocate). Build them # unconditionally; three small fp32 tensors per layer. - layer.self_attn._build_mtp_conv_weights() + layer.linear_attn._build_mtp_conv_weights() logger.info( f"Kimi K3: loaded {num_params} parameters and the expert " f"slices of {num_moe_layers} MoE layers; fused prefill/decode/verify " @@ -3537,8 +2582,8 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: n_glue = 0 for layer in self.model.layers: if getattr(layer, "is_kda", False) and _has_weights(layer): - layer.self_attn.finalize_decode_weights_fp8() - n_glue += int(layer.self_attn._bfa_proj_weight is not None) + layer.linear_attn.finalize_decode_weights_fp8() + n_glue += int(layer.linear_attn._bfa_proj_weight is not None) logger.info( f"Kimi K3: FP8 fused prefill/decode/verify projections on " f"{n_glue} KDA layers" diff --git a/tensorrt_llm/_torch/modules/kimi_kda/__init__.py b/tensorrt_llm/_torch/modules/kimi_kda/__init__.py index c457a986744c..c3a05eb31da8 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/__init__.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/__init__.py @@ -9,9 +9,6 @@ interface. """ -from .kimi_kda_mixer import KimiKDAKernelPath, KimiKDALinearAttention +from .kimi_kda_mixer import KimiKDALinearAttention -__all__ = [ - "KimiKDAKernelPath", - "KimiKDALinearAttention", -] +__all__ = ["KimiKDALinearAttention"] diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py index 283c721b1c1b..f0caa5eb64cc 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -101,10 +101,6 @@ def is_intree_prefill_available() -> bool: return False -def _load_fla_chunk_kda() -> ModuleType: - return importlib.import_module("fla.ops.kda") - - # --------------------------------------------------------------------------- # In-tree KDA multi-token verify op (CuTe DSL, trtllm::kda_mtp_decode). # --------------------------------------------------------------------------- @@ -209,16 +205,6 @@ def __init__( f"verify={self.verify_kernel_path}" ) - def get_prefill_source(self) -> str: - if self.prefill_kernel_path == "optimized": - return _load_prefill_module().__file__ or "" - return _load_fla_chunk_kda().__file__ or "" - - def get_decode_source(self) -> str: - if self.decode_kernel_path == "optimized": - return _kda_decode.__file__ or "" - return _load_fla_chunk_kda().__file__ or "" - def mtp_verify(self, **kwargs) -> torch.Tensor: """Run the fused KDA multi-token verify kernel. @@ -354,7 +340,7 @@ def prefill_chunk_kda( from fla.ops.kda import chunk_kda - o, final_state = chunk_kda( + return chunk_kda( q=q, k=k, v=v, @@ -373,7 +359,6 @@ def prefill_chunk_kda( state_v_first=True, cu_seqlens=cu_seqlens, ) - return o, final_state def decode_kda(self, **kwargs) -> torch.Tensor: """Run the fused KDA single-token decode kernel. diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py index 3f123c707daa..1caad9da1e67 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -1,83 +1,69 @@ # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""KimiKDALinearAttention — Kimi K3 linear-attention module for the PyTorch backend. +"""Kimi K3 KDA production frontend for the TensorRT-LLM PyTorch executor. -Structural mirror of the HF reference ``KimiDeltaAttention`` in -``modeling_kimi.py``. Same parameter names, same layer shapes, same short -convolution + FLA gating + FusedRMSNormGated output-gate stack. The -delta-rule inner loop is routed through :mod:`_kda_kernels`, which selects -the optimized sm_100 CuTe/Triton chunked prefill and fused CUDA decode -kernels on Blackwell and falls back to the FLA references elsewhere. +The checkpoint-visible projections, short convolutions, forget/beta gates, +and gated output norm mirror the HF ``KimiDeltaAttention`` structure. The +delta-rule inner loop dispatches to the optimized Blackwell prefill, +decode, and verification kernels, with explicit FLA fallbacks. Cache ownership --------------- -KDA carries three short-convolution states (``conv_state_{q,k,v}``, HF -layout ``[B, D, W]`` bf16) and one delta-rule recurrent state -(``recurrent_state``, layout ``[B, HV, V, K]`` fp32, matching the optimized -kernel's transposed convention). These match the hybrid-cache ownership -pattern used by the mamba modules; the runtime cache-manager plumbing -(``AttentionMetadata`` split, cache indices, spec/verify path) is deferred -to the model-assembly wiring goal. This module exposes parity entry points -that consume and return the state tensors directly so module-level tests -can prove state roundtrip without the runtime plumbing. - -Kernel mutations for negative controls --------------------------------------- -Two invariants have their own construction switches so parity tests can -prove they are actually being enforced: - -* ``gate_lower_bound_override`` — replace the ``linear_attn_config`` - gate lower bound at forward time. A value that disagrees with the HF - reference must fail parity. -* ``wrong_state_layout`` — permute the recurrent state's V/K axes before - and after the decode kernel call so read/write hit mislabeled slots. - Because K == V the shape check still passes but the numerics break. +The module reads and updates three short-convolution states in the combined +``[slots, 3D, W]`` bf16 pool and one delta-rule recurrent state in the +``[slots, H, V, K]`` fp32 pool. It owns the ``AttentionMetadata`` split, +cache-slot indexing, and speculative-verification state handling directly. """ from __future__ import annotations import math -from dataclasses import dataclass -from typing import Optional, Tuple +import os +from typing import List, Optional, Tuple import torch -from einops import rearrange from fla.modules import FusedRMSNormGated, ShortConvolution -from fla.ops.kda import fused_recurrent_kda from torch import nn -from ._kda_kernels import KDAKernelDispatch, is_kda_optimized_supported +from ...attention_backend import AttentionMetadata +from ...distributed import AllReduce, AllReduceStrategy +from ...modules.multi_stream_utils import maybe_execute_in_parallel +from ._kda_kernels import KDAKernelDispatch +_KDA_INDEXED_STATE_POOL_ENABLED = os.environ.get("TLLM_KDA_ENABLE_INDEXED_STATE_POOL", "1") == "1" +# Heuristic ported from SGLang's Blackwell cutoff: +# https://github.com/sgl-project/sglang/blob/e84bbf68efb683c9e2eef4168c5198042544599d/python/sglang/srt/models/kimi_k3.py#L946-L954 +# It has not been tuned for TensorRT-LLM; benchmark and retune it for TRT-LLM's +# projection kernels. Verify intentionally counts B * num_steps because those +# flattened token rows form the projection GEMMs' M dimension. +_KDA_BFA_MULTISTREAM_MAX_ROWS = 128 -def _meta_safe_cast_dtype(module, dtype): - """``module.to(dtype=dtype)`` that also works under ``MetaInitMode``. - ``Module.to`` dispatches ``aten._to_copy``, which MetaInitMode rejects - (it would silently fall back to full CPU construction of the model — - ~70 GB of host RAM per rank for Kimi K3). Under meta init the values - are garbage anyway, so a dtype-only re-allocation via ``empty_like`` - (an allowed init op) is equivalent; off meta this matches ``.to``. +def _meta_safe_cast_dtype(module: nn.Module, dtype: torch.dtype) -> None: + """Cast floating parameters while preserving meta-device construction. + + ``Module.to`` dispatches ``aten._to_copy``, which ``MetaInitMode`` rejects + and would force the full model to fall back to eager CPU construction. + Meta values are uninitialized, so dtype-only reallocation via + ``empty_like`` is equivalent there; materialized tensors use ``to``. """ - import torch as _torch - def _cast(t): - if not t.is_floating_point(): - return t - if t.is_meta: - return _torch.empty_like(t, dtype=dtype) - return t.to(dtype=dtype) + def _cast(tensor: torch.Tensor) -> torch.Tensor: + if not tensor.is_floating_point(): + return tensor + if tensor.is_meta: + return torch.empty_like(tensor, dtype=dtype) + return tensor.to(dtype=dtype) module._apply(_cast) class _MetaSafeFusedRMSNormGated(FusedRMSNormGated): - """FusedRMSNormGated whose init survives the model loader's MetaInitMode. + """FLA gated RMSNorm whose initialization supports ``MetaInitMode``. ``FusedRMSNormGated.reset_parameters`` uses ``nn.init.ones_`` (a plain - ``fill_``), which MetaInitMode rejects on meta tensors and which would - force the whole model construction to fall back to eager CPU init. - ``uniform_(1, 1)`` produces identical values and is on MetaInitMode's - random-init allowlist. + ``fill_``), which ``MetaInitMode`` rejects. ``uniform_(1, 1)`` produces + the same values and is on the allowed random-initialization path. """ def reset_parameters(self) -> None: @@ -86,162 +72,82 @@ def reset_parameters(self) -> None: self.weight.uniform_(1.0, 1.0) -@dataclass -class KimiKDACachedState: - """Per-layer KDA cache tensors in HF layout. - - ``conv_state_*`` — shape ``[B, D, W]`` bf16 where ``W`` is - ``short_conv_kernel_size`` and the newest processed token sits at - position ``W-1``. ``D`` is ``num_heads * head_dim`` for q/k and - ``num_heads * head_dim`` for v (K3 uses HV == H). - - ``recurrent_state`` — shape ``[B, HV, V, K]`` fp32. This is the - transposed layout the optimized KDA kernels expect, and it is the same - layout HF stores when running with ``transpose_state_layout=True``. - ``None`` fields are treated as zero. - """ - - conv_state_q: Optional[torch.Tensor] - conv_state_k: Optional[torch.Tensor] - conv_state_v: Optional[torch.Tensor] - recurrent_state: Optional[torch.Tensor] - - -class KimiKDAKernelPath: - """Enum-like string tags for the selected KDA kernel path.""" - - OPTIMIZED = "optimized" - FLA = "fla" - - -def _hf_conv_to_kernel_conv( - hf_cache: Optional[torch.Tensor], - b: int, - d: int, - w: int, - device: torch.device, - dtype: torch.dtype, -) -> torch.Tensor: - """HF ``[B, D, W]`` conv cache -> optimized kernel's ``[B, D, W-1]``. - - HF stores W positions with the newest processed token last. The - optimized decode kernel's ``cs_*`` argument stores the ``W-1`` - historical positions before the incoming token; drop the oldest column. - """ - if hf_cache is None: - return torch.zeros(b, d, w - 1, device=device, dtype=dtype) - return hf_cache[:, :, 1:].contiguous() - - -def _roll_hf_conv( - prev_hf: Optional[torch.Tensor], - x_new_col: torch.Tensor, - b: int, - d: int, - w: int, - device: torch.device, - dtype: torch.dtype, -) -> torch.Tensor: - """Roll an HF-layout conv cache by one token. - - HF ``ShortConvolution.step`` does - ``cache.copy_(cache.roll(shifts=-1, dims=-1)); cache[:, :, -1] = x``. - We implement the same semantics via ``torch.cat`` so the update is - independent of the kernel's internal cs handling. - """ - if prev_hf is None: - prev = torch.zeros(b, d, w, device=device, dtype=dtype) - else: - prev = prev_hf.to(dtype=dtype) - return torch.cat([prev[:, :, 1:], x_new_col.to(dtype)], dim=-1).contiguous() +def _kda_split_conv_sections( + conv_state: torch.Tensor, dim: int +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Split a gathered ``[N, 3D, W]`` convolution cache into Q/K/V.""" + return ( + conv_state[:, :dim].contiguous(), + conv_state[:, dim : 2 * dim].contiguous(), + conv_state[:, 2 * dim :].contiguous(), + ) class KimiKDALinearAttention(nn.Module): - """Kimi K3 linear-attention module — in-tree production version. - - Parameters - ---------- - hidden_size : int - num_heads : int - head_dim : int - conv_kernel_size : int - use_full_rank_gate : bool - gate_lower_bound : Optional[float] - rms_norm_eps : float - dtype : Optional[torch.dtype] - layer_idx : int - use_optimized_prefill : bool - Enable the optimized prefill path when supported. - use_optimized_decode : bool - Enable the optimized decode path when supported. - gate_lower_bound_override : Optional[float] - Override the ``linear_attn_config`` gate lower bound. Test knob for - the "wrong gate lower bound" mutation control. - wrong_state_layout : bool - Swap the V and K axes of the recurrent state around the decode - kernel call. Test knob for the "wrong state layout" mutation - control on the decode path. - """ + """Production Kimi K3 KDA module with direct cache-pool ownership.""" def __init__( self, - *, - hidden_size: int, - num_heads: int, - head_dim: int, - conv_kernel_size: int, - use_full_rank_gate: bool, - gate_lower_bound: Optional[float], - rms_norm_eps: float = 1e-5, - dtype: Optional[torch.dtype] = None, - layer_idx: int = 0, - use_optimized_prefill: bool = True, - use_optimized_decode: bool = True, - gate_lower_bound_override: Optional[float] = None, - wrong_state_layout: bool = False, + cfg, + layer_idx: int, + mapping=None, + allreduce_strategy=AllReduceStrategy.AUTO, + aux_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__() - self.hidden_size = hidden_size - self.num_heads = num_heads - self.head_dim = head_dim - self.head_k_dim = head_dim - self.num_k_heads = num_heads - self.conv_size = conv_kernel_size - self.use_full_rank_gate = use_full_rank_gate - self.gate_lower_bound = gate_lower_bound - self.rms_norm_eps = rms_norm_eps + lin = cfg.linear_attn_config + self.hidden_size = cfg.hidden_size self.layer_idx = layer_idx - self.gate_lower_bound_override = gate_lower_bound_override - self.wrong_state_layout = wrong_state_layout - - projection_k_size = self.head_k_dim * self.num_k_heads - projection_size = self.head_dim * self.num_heads - - self.q_proj = nn.Linear(hidden_size, projection_k_size, bias=False) - self.k_proj = nn.Linear(hidden_size, projection_k_size, bias=False) - self.v_proj = nn.Linear(hidden_size, projection_size, bias=False) + self.head_dim = lin["head_dim"] + self.head_k_dim = self.head_dim + self.conv_size = lin["short_conv_kernel_size"] + self.use_full_rank_gate = lin.get("use_full_rank_gate", True) + self.gate_lower_bound = lin.get("gate_lower_bound", None) + self.rms_norm_eps = cfg.rms_norm_eps + self._use_indexed_ssm_pool = _KDA_INDEXED_STATE_POOL_ENABLED + + if mapping is not None and mapping.tp_size > 1 and not mapping.enable_attention_dp: + self._kda_tp_size = mapping.tp_size + else: + self._kda_tp_size = 1 + self._kda_tp_rank = mapping.tp_rank if self._kda_tp_size > 1 else 0 + self._o_allreduce = ( + AllReduce(mapping=mapping, strategy=allreduce_strategy, dtype=torch.bfloat16) + if self._kda_tp_size > 1 + else None + ) + num_heads = lin["num_heads"] + if num_heads % self._kda_tp_size != 0: + raise ValueError( + f"KDA num_heads {num_heads} not divisible by tp_size {self._kda_tp_size}" + ) + self.num_heads = num_heads // self._kda_tp_size + self.num_k_heads = self.num_heads + projection_size = self.num_heads * self.head_dim + self.proj_size = projection_size + + # Keep the logical projections separate for checkpoint-compatible + # parameter names and portable fallbacks. After weight loading, the + # Blackwell path aliases them into one fused QKVG weight buffer. + self.q_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + self.k_proj = nn.Linear(self.hidden_size, projection_size, bias=False) + self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False) self.q_conv1d = ShortConvolution( - hidden_size=projection_k_size, - kernel_size=conv_kernel_size, - activation="silu", + hidden_size=projection_size, kernel_size=self.conv_size, activation="silu" ) self.k_conv1d = ShortConvolution( - hidden_size=projection_k_size, - kernel_size=conv_kernel_size, - activation="silu", + hidden_size=projection_size, kernel_size=self.conv_size, activation="silu" ) self.v_conv1d = ShortConvolution( - hidden_size=projection_size, - kernel_size=conv_kernel_size, - activation="silu", + hidden_size=projection_size, kernel_size=self.conv_size, activation="silu" ) self.A_log = nn.Parameter( - torch.log(torch.empty(num_heads, dtype=torch.float32).uniform_(1, 16)) + torch.log(torch.empty(self.num_heads, dtype=torch.float32).uniform_(1, 16)) ) - self.f_a_proj = nn.Linear(hidden_size, head_dim, bias=False) - self.f_b_proj = nn.Linear(head_dim, projection_size, bias=False) + self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) # dt_bias must be initialized: torch.empty heap garbage can contain # NaN bit patterns, which poison both the optimized and FLA gates in # randomly-constructed modules (parity tests) — TRTLLM-15204. This @@ -256,138 +162,331 @@ def __init__( math.log(1e-3), math.log(1e-1) ) ) - self.b_proj = nn.Linear(hidden_size, num_heads, bias=False) + self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False) - if use_full_rank_gate: - self.g_proj = nn.Linear(hidden_size, projection_size, bias=False) + if self.use_full_rank_gate: + self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False) else: - self.g_a_proj = nn.Linear(hidden_size, head_dim, bias=False) - self.g_b_proj = nn.Linear(head_dim, projection_size, bias=False) - - self.o_norm = _MetaSafeFusedRMSNormGated(head_dim, eps=rms_norm_eps, activation="sigmoid") - self.o_proj = nn.Linear(projection_size, hidden_size, bias=False) - - # Installed together by the FP8 weight loader (fused [q | k | v | g] - # decode GEMM). Declared here so the decode path never sees a - # half-installed pair. + self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False) + self.o_norm = _MetaSafeFusedRMSNormGated( + self.head_dim, eps=self.rms_norm_eps, activation="sigmoid" + ) + self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False) + # Installed together by the FP8 weight loader as the fused + # [q | k | v | g] projection and its output-section metadata. self.qkvg_proj: Optional[nn.Module] = None - self.qkvg_split_sizes: Optional[list[int]] = None - - if dtype is not None: - _meta_safe_cast_dtype(self, dtype) + self.qkvg_split_sizes: Optional[List[int]] = None + _meta_safe_cast_dtype(self, torch.bfloat16) - # The optimized decode/verify kernels are specialized for the Kimi - # K3 shape (K == V == 128). Reduced-dim test configurations must - # fall back to FLA instead of hard-failing inside the kernels. + # The optimized decode/verify kernels specialize the Kimi K3 + # K == V == 128 shape. Other shapes must use the portable fallback. kernel_shape_ok = self.head_k_dim == 128 and self.head_dim == 128 self._dispatch = KDAKernelDispatch( - use_optimized_prefill=use_optimized_prefill, - use_optimized_decode=use_optimized_decode and kernel_shape_ok, + use_optimized_decode=kernel_shape_ok, use_optimized_verify=kernel_shape_ok, ) - # ------------------------------------------------------------------ - # Introspection helpers used by the test smoke. - # ------------------------------------------------------------------ - - @property - def prefill_kernel_path(self) -> str: - return self._dispatch.prefill_kernel_path - - @property - def decode_kernel_path(self) -> str: - return self._dispatch.decode_kernel_path - - @property - def verify_kernel_path(self) -> str: - return self._dispatch.verify_kernel_path - - @property - def sm_100_optimized_supported(self) -> bool: - return is_kda_optimized_supported() - - def prefill_kernel_source(self) -> str: - return self._dispatch.get_prefill_source() - - def decode_kernel_source(self) -> str: - return self._dispatch.get_decode_source() - - def prefill_chunk_kda(self, **kwargs): - """Kernel-level chunked prefill via the dispatch. - - Used by the executor runtime (``KimiKDARuntime``), which owns the - projections, convs, and cache pools itself and only needs the - delta-rule inner loop. States are exchanged in the V-first - ``[N, H, V, K]`` pool layout on both dispatch paths — see - ``KDAKernelDispatch.prefill_chunk_kda``. + # Fused prefill/decode/verify projection weights, built after checkpoint + # load. BF16 uses separate fused [q | k | v | g] and [f_a | b] + # GEMMs; FP8 supplies qkvg through the fused projection and reuses the + # BF16 [f_a | b] weight. + self._qkvg_proj_weight: Optional[torch.Tensor] = None + self._bfa_proj_weight: Optional[torch.Tensor] = None + self._w_q_t = self._w_k_t = self._w_v_t = None + self._A_log_f32 = self._dt_bias_f32 = self._onorm_w_f32 = None + # Fork/join state for overlapping the small [f_a | b] -> f_b chain + # with the wide QKVG projection during CUDA-graph execution. + self._projection_aux_stream = aux_stream + self._projection_fork_event = torch.cuda.Event() + self._projection_join_event = torch.cuda.Event() + # Persistent batch-row-dense staging for the fused decode kernel's + # convolution windows. It is allocated once and never reallocated. + self._cs_dense: Optional[torch.Tensor] = None + # FP32 [dim, W] convolution weights for fused verification, built + # after weight loading so the first captured call does not allocate. + self._mtp_conv_weights: Optional[Tuple[torch.Tensor, ...]] = None + + def finalize_decode_weights(self) -> None: + """Build fused projection weights and decode constants after weight load. + + 1. Separate fused ``[q | k | v | g]`` and ``[f_a | b]`` projections. + Keeping the wide qkvg output aligned avoids degrading its GEMM + kernel selection with the small f_a and b tails. Source parameters + are repointed to row views of the fused buffers, so prefill and + verify paths keep using them without duplicate weight storage. + 2. Kernel-layout constants that ``_decode_via_optimized`` used to + rebuild with ~6 device kernels per layer per decode step: + transposed conv weights (bf16 ``[W, D]``) and fp32 copies of + ``A_log`` / ``dt_bias`` / ``o_norm.weight``. """ - return self._dispatch.prefill_chunk_kda(**kwargs) - - # ------------------------------------------------------------------ - # Prefill entry (Goal 2.1 pass path). - # ------------------------------------------------------------------ + if self._dispatch.decode_kernel_path != "optimized" or not self.use_full_rank_gate: + return + if self.q_proj.weight.device.type != "cuda": + return + with torch.no_grad(): + qkvg_modules = ( + self.q_proj, + self.k_proj, + self.v_proj, + self.g_proj, + ) + qkvg_weight = self._merge_projection_weights(qkvg_modules) + # Eight BF16 outputs occupy 16 bytes, so padding keeps each output row + # aligned for vectorized f_b consumption; it is not a kernel requirement. + bfa_weight = self._merge_projection_weights((self.f_a_proj, self.b_proj), pad_rows_to=8) + self._build_decode_kernel_constants() + self._bfa_proj_weight = bfa_weight + # Publish last: both weights are required by the BF16 fast path. + self._qkvg_proj_weight = qkvg_weight + + @staticmethod + def _merge_projection_weights( + modules: tuple[nn.Linear, ...], pad_rows_to: int = 1 + ) -> torch.Tensor: + """Concatenate linear weights and repoint the modules to row views.""" + weights = [module.weight.data for module in modules] + padding = (-sum(weight.shape[0] for weight in weights)) % pad_rows_to + if padding: + weights.append(weights[0].new_zeros((padding, weights[0].shape[1]))) + fused = torch.cat(weights, dim=0).contiguous() + offset = 0 + for module in modules: + rows = module.weight.shape[0] + module.weight.data = fused[offset : offset + rows] + offset += rows + return fused + + def _build_decode_kernel_constants(self) -> None: + """Kernel-layout constants shared by both finalize variants.""" + self._w_q_t = ( + self.q_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + ) + self._w_k_t = ( + self.k_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + ) + self._w_v_t = ( + self.v_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + ) + self._A_log_f32 = self.A_log.detach().float().contiguous() + self._dt_bias_f32 = self.dt_bias.detach().float().contiguous() + self._onorm_w_f32 = self.o_norm.weight.detach().float().contiguous() + # Build the fused-verify conv constants eagerly too, so the first + # verify call never allocates (a capture-unsafe lazy allocation). + self._build_mtp_conv_weights() + + def finalize_decode_weights_fp8(self) -> None: + """FP8 counterpart of ``finalize_decode_weights()``. + + Runs AFTER ``_convert_kda_projections_to_fp8_weight_read``, so + q/k/v/g already live in the fused FP8 ``qkvg_proj`` GEMM. Only the + two small BF16 projections reading the same hidden — ``f_a_proj`` and + ``b_proj`` — are fused here into one ``[f_a | b]`` weight, with the + source parameters repointed to row views. Prefill, decode, and + verification then share both fused projections; the kernel-layout + constants are decode-only. + """ + if self._dispatch.decode_kernel_path != "optimized" or not self.use_full_rank_gate: + return + fused_qkvg = self.qkvg_proj + split_sizes = self.qkvg_split_sizes + if fused_qkvg is None or split_sizes is None or len(split_sizes) != 4: + return + if self.f_a_proj.weight.device.type != "cuda": + return + with torch.no_grad(): + bfa_weight = self._merge_projection_weights((self.f_a_proj, self.b_proj), pad_rows_to=8) + self._build_decode_kernel_constants() + # Publish last: enables fused [f_a | b] in prefill/decode/verify. + self._bfa_proj_weight = bfa_weight + + def forward( + self, hidden_states: torch.Tensor, attn_metadata: AttentionMetadata + ) -> torch.Tensor: + """``hidden_states``: flattened ``[num_tokens, hidden]`` (ctx tokens + first, then one token per generation request).""" + mamba_metadata = attn_metadata.mamba_metadata + num_prefills = attn_metadata.num_contexts + num_ctx_tokens = attn_metadata.num_ctx_tokens + batch_size = attn_metadata.seq_lens.shape[0] + # index_copy_/index_select need int64 indices; the int64 mirror is + # prepared once per step by Mamba2Metadata.prepare() so KDA layers + # do not each replay an int32->int64 cast inside the decode graph. + state_indices = getattr(mamba_metadata, "state_indices_long", None) + if state_indices is None or state_indices.shape[0] != batch_size: + state_indices = mamba_metadata.state_indices[:batch_size].long() + cu_seqlens = mamba_metadata.query_start_loc_long[: num_prefills + 1] + num_decodes = batch_size - num_prefills + + layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) + conv_pool = layer_cache.conv # [slots, 3D, W] bf16 + ssm_pool = layer_cache.temporal # [slots, H, V, K] fp32 + + outputs: List[torch.Tensor] = [] + if num_prefills > 0: + outputs.append( + self.forward_prefill( + hidden_states[:num_ctx_tokens], + cu_seqlens, + mamba_metadata, + num_prefills, + conv_pool, + ssm_pool, + state_indices[:num_prefills], + layer_cache, + ) + ) + if num_decodes > 0: + decode_rows = hidden_states.shape[0] - num_ctx_tokens + if decode_rows == num_decodes: + outputs.append( + self.forward_decode( + hidden_states[num_ctx_tokens:], + conv_pool, + ssm_pool, + state_indices[num_prefills:], + mamba_metadata, + layer_cache, + ssm_state_indices=( + mamba_metadata.state_indices[num_prefills:batch_size] + if self._use_indexed_ssm_pool + else None + ), + ) + ) + else: + # Speculative verification: each generation request carries + # 1 + draft_len tokens (drafts are padded to the static max, + # so T is uniform). Per-step states go to the manager's + # SpeculativeState scratch buffers — never the live pools — + # and kv_cache_manager.update_mamba_states() promotes the + # accepted step after sampling. + assert decode_rows % num_decodes == 0, ( + f"ragged generation batch: {decode_rows} tokens for {num_decodes} requests" + ) + outputs.append( + self.forward_verify( + hidden_states[num_ctx_tokens:], + decode_rows // num_decodes, + layer_cache, + conv_pool, + ssm_pool, + state_indices[num_prefills:], + ) + ) + out = outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0) + if self._o_allreduce is not None: + # Head-sharded TP: every rank ran its head shard on the same + # local batch; sum the row-sharded o_proj partials. + out = self._o_allreduce(out) + return out + + def _has_kda_replay_caches(self, layer_cache) -> bool: + """True when the manager allocated the fused-verify replay caches.""" + return getattr(layer_cache, "kda_qkg_cache", None) is not None + + def _sync_kda_replay_conv_window( + self, layer_cache, slot_indices, conv_q, conv_k, conv_v + ) -> None: + """Seed the replay conv caches' committed window from FLA windows. + + The fused verify kernel keeps its own extended fp32 dim-contiguous + conv caches; their committed window (columns ``[0, W-1)``) must hold + the last ``W-1`` raw conv inputs whenever another path (prefill, + plain decode) advances the base conv pool. The FLA window's oldest + column drops out of every future convolution, so columns ``[1, W)`` + of the FLA cache map 1:1 onto the committed window. + """ + if not self._has_kda_replay_caches(layer_cache): + return + w = self.conv_size + for cache, window in ( + (layer_cache.kda_conv_q, conv_q), + (layer_cache.kda_conv_k, conv_k), + (layer_cache.kda_conv_v, conv_v), + ): + cache[:, :, : w - 1].index_copy_(0, slot_indices, window[:, :, 1:].to(cache.dtype)) def forward_prefill( self, - hidden_states: torch.Tensor, - cu_seqlens: Optional[torch.Tensor] = None, + x2d, + cu_seqlens, + mamba_metadata, + num_prefills, + conv_pool, + ssm_pool, + slot_indices, + layer_cache=None, ) -> torch.Tensor: - """Prefill forward matching HF ``KimiDeltaAttention.forward`` in chunk mode. - - Parameters - ---------- - hidden_states : ``(B, T, hidden_size)`` for equal-length prefill or - ``(1, sum(seq_lens), hidden_size)`` when ``cu_seqlens`` is given. - cu_seqlens : optional cumulative sequence lengths for varlen inputs. - - Returns - ------- - ``(B, T, hidden_size)`` output tensor (equal-length case) or - ``(1, sum(seq_lens), hidden_size)`` (varlen case). - """ - if cu_seqlens is not None: - cu_seqlens = cu_seqlens.to(device=hidden_states.device, dtype=torch.long) + from einops import rearrange - q_proj_states = self.q_proj(hidden_states) - k_proj_states = self.k_proj(hidden_states) - v_proj_states = self.v_proj(hidden_states) + d = self.proj_size + x = x2d.unsqueeze(0) # [1, T, hidden] - q, _ = self.q_conv1d( - x=q_proj_states, - cache=None, - output_final_state=False, - cu_seqlens=cu_seqlens, + onorm_g = None + if self._qkvg_proj_weight is not None: + qkvg = torch.nn.functional.linear(x, self._qkvg_proj_weight) + q_proj_states, k_proj_states, v_proj_states = qkvg[..., : 3 * d].split(d, dim=-1) + onorm_g = qkvg[..., 3 * d : 4 * d] + else: + fused_qkvg = self.qkvg_proj + if fused_qkvg is not None: + qkvg = fused_qkvg(x) + q_proj_states, k_proj_states, v_proj_states = qkvg[..., : 3 * d].split(d, dim=-1) + qkvg_split_sizes = self.qkvg_split_sizes + if ( + self.use_full_rank_gate + and qkvg_split_sizes is not None + and len(qkvg_split_sizes) == 4 + ): + onorm_g = qkvg[..., 3 * d : 4 * d] + else: + q_proj_states = self.q_proj(x) + k_proj_states = self.k_proj(x) + v_proj_states = self.v_proj(x) + + # Initial states: present for continuation chunks (chunked prefill) + # and for prefix-cache hits (block reuse), where the previous + # conv/recurrent state was onboarded into this request's slot. + conv_q_in = conv_k_in = conv_v_in = None + recurrent_in = None + if mamba_metadata.use_initial_states: + has_init = mamba_metadata.has_initial_states[:num_prefills] + cs = conv_pool.index_select(0, slot_indices) + cs[~has_init] = 0 + conv_q_in, conv_k_in, conv_v_in = _kda_split_conv_sections(cs, d) + recurrent_in = ssm_pool.index_select(0, slot_indices) + recurrent_in[~has_init] = 0 + + q, conv_q = self.q_conv1d( + q_proj_states, cache=conv_q_in, output_final_state=True, cu_seqlens=cu_seqlens ) - k, _ = self.k_conv1d( - x=k_proj_states, - cache=None, - output_final_state=False, - cu_seqlens=cu_seqlens, + k, conv_k = self.k_conv1d( + k_proj_states, cache=conv_k_in, output_final_state=True, cu_seqlens=cu_seqlens ) - v, _ = self.v_conv1d( - x=v_proj_states, - cache=None, - output_final_state=False, - cu_seqlens=cu_seqlens, + v, conv_v = self.v_conv1d( + v_proj_states, cache=conv_v_in, output_final_state=True, cu_seqlens=cu_seqlens ) - g = self.f_b_proj(self.f_a_proj(hidden_states)) + if self._bfa_proj_weight is not None: + bfa = torch.nn.functional.linear(x, self._bfa_proj_weight) + f_a = bfa[..., : self.head_dim] + beta = bfa[..., self.head_dim : self.head_dim + self.num_heads].float() + g = self.f_b_proj(f_a) + else: + g = self.f_b_proj(self.f_a_proj(x)) + beta = self.b_proj(x).float() g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim) - beta = self.b_proj(hidden_states).float() q = rearrange(q, "... (h d) -> ... h d", d=self.head_k_dim) k = rearrange(k, "... (h d) -> ... h d", d=self.head_k_dim) v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) - lower_bound = ( - self.gate_lower_bound_override - if self.gate_lower_bound_override is not None - else self.gate_lower_bound - ) - safe_gate = lower_bound is not None - scale = self.head_k_dim**-0.5 - - o, _final_state = self._dispatch.prefill_chunk_kda( + # Kernel dispatch (in-tree trtllm::kda_prefill or FLA chunk_kda). + # Both paths exchange states in the pool's V-first [N, H, V, K] + # layout, so recurrent_in / final_state map to ssm_pool 1:1. + lower_bound = self.gate_lower_bound + o, final_state = self._dispatch.prefill_chunk_kda( q=q, k=k, v=v, @@ -395,343 +494,611 @@ def forward_prefill( beta=beta, A_log=self.A_log, dt_bias=self.dt_bias, - scale=scale, - initial_state=None, - safe_gate=safe_gate, + scale=self.head_k_dim**-0.5, + initial_state=recurrent_in, + safe_gate=lower_bound is not None, lower_bound=lower_bound, cu_seqlens=cu_seqlens, - chunk_size=64, ) - if self.use_full_rank_gate: - g_out = self.g_proj(hidden_states) - else: - g_out = self.g_b_proj(self.g_a_proj(hidden_states)) - g_out = rearrange(g_out, "... (h d) -> ... h d", d=self.head_dim) - o = self.o_norm(o, g_out) - - o = rearrange(o, "b t h d -> b t (h d)") - o = self.o_proj(o) - return o + # Persist per-request states into the pools. + conv_pool.index_copy_( + 0, slot_indices, torch.cat([conv_q, conv_k, conv_v], dim=1).to(conv_pool.dtype) + ) + assert final_state is not None + ssm_pool.index_copy_(0, slot_indices, final_state.to(ssm_pool.dtype)) + # Fused-verify replay caches: seed the committed conv window so the + # first verify round convolves the correct history (pending drafts + # are zero for a fresh request, so the tail columns are unused). + self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_q, conv_k, conv_v) - # ------------------------------------------------------------------ - # Decode entry (Goal 2.1 pass path). - # ------------------------------------------------------------------ + return self._output_gate_and_proj(x, o, onorm_g) def forward_decode( self, - hidden_states: torch.Tensor, - cache: Optional[KimiKDACachedState] = None, - ssm_state_indices: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, KimiKDACachedState]: - """T=1 cached-decode forward. Returns ``(o, new_cache)``. - - ``hidden_states`` shape ``(B, 1, hidden_size)``. Cache is - ``KimiKDACachedState`` in HF layout; ``None`` fields become zero - tensors. + x2d, + conv_pool, + ssm_pool, + slot_indices, + mamba_metadata=None, + layer_cache=None, + ssm_state_indices=None, + ) -> torch.Tensor: + """Plain T=1 decode, fast path. + + Calls ``trtllm::kda_decode`` directly with kernel-native layouts + (nsys 07-24: the reference path spent ~70 us/layer on glue around + the 5 us kernel — 6 separate in-projection GEMV pairs, per-step + re-transposition of constant weights, conv-window slice/roll + copies, per-call torch.arange defaults, and redundant dtype + casts): + + * one wide fused qkvg GEMV on the main stream, overlapped with the + fused [f_a | b] GEMV and f_b GEMV on the auxiliary stream for + CUDA-graph batches up to 128 tokens; + * conv windows staged with one gather + one repack copy into a + persistent dense per-section buffer; + * conv-pool write-back with one cat + one index_copy_; + * constant tensors (transposed conv weights, fp32 A_log/dt_bias/ + o_norm weight) reused instead of rebuilt per step. + + The conv windows remain gathered batch-row-dense. When stable + int32 slot indices are supplied, the recurrent-state pool is passed + directly and the CUDA wrapper selects its indexed-state launch; + otherwise the state uses the batch-row-dense static layout. """ - b, q_len, _ = hidden_states.shape - assert q_len == 1, f"KimiKDALinearAttention.forward_decode expects T=1, got T={q_len}" - - if self._dispatch.decode_kernel_path == KimiKDAKernelPath.OPTIMIZED: - return self._decode_via_optimized(hidden_states, cache, b, ssm_state_indices) - if ssm_state_indices is not None: - raise ValueError("ssm_state_indices requires the optimized KDA decode kernel") - return self._decode_via_fla(hidden_states, cache, b) - - # ------------------------------------------------------------------ - # Internals — optimized decode dispatch. - # ------------------------------------------------------------------ - - def _decode_via_optimized( - self, - hidden_states: torch.Tensor, - cache: Optional[KimiKDACachedState], - b: int, - ssm_state_indices: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, KimiKDACachedState]: - dev = hidden_states.device + if self._dispatch.decode_kernel_path != "optimized": + ssm_state_indices = None + has_qkvg_projection = self._qkvg_proj_weight is not None or self.qkvg_proj is not None + if ( + self._dispatch.decode_kernel_path != "optimized" + or not has_qkvg_projection + or self._bfa_proj_weight is None + or mamba_metadata is None + or ssm_pool.dtype != torch.float32 + ): + return self.forward_decode_fallback( + x2d, + conv_pool, + ssm_pool, + slot_indices, + layer_cache, + ssm_state_indices, + ) + d = self.proj_size + hd = self.head_dim H = self.num_heads - HV = self.num_heads - K_dim = self.head_dim - V_dim = self.head_dim + B = x2d.shape[0] W = self.conv_size - projection_size = H * K_dim - projection_v_size = HV * V_dim - - # q/k/v and the full-rank output gate all read this same normed hidden. - # When their weights are read at FP8 block-scale (Blackwell decode), the - # loader fuses them into one ``qkvg_proj`` GEMM: one activation quant and - # one GEMM launch replace four, which is what the launch-bound - # generation step needs. The split is output-identical to the per - # projection GEMMs (same activation, same weight slices). The forget - # gate (f_a/f_b), beta and low-rank output gate stay BF16, so they keep - # their own calls. - fused_qkvg = self.qkvg_proj if self.qkvg_split_sizes is not None else None - if fused_qkvg is not None: - parts = fused_qkvg(hidden_states).split(self.qkvg_split_sizes, dim=-1) - q_proj_states, k_proj_states, v_proj_states = parts[0], parts[1], parts[2] - onorm_g_hidden = ( - parts[3] if self.use_full_rank_gate else self.g_b_proj(self.g_a_proj(hidden_states)) + # Allocated ONCE at the pool slot count (== per-rank max batch on + # the Mixed manager; ``slot_indices`` are distinct pool rows and + # this is the plain one-token-per-request path, so B never exceeds + # it) and never reallocated: captured CUDA graphs hold this + # pointer, so a realloc would leave earlier graphs writing into + # freed memory. Footprint: slots x ~9(H=6)..222(H=96) KB per layer. + buf = self._cs_dense + if buf is None: + if torch.cuda.is_current_stream_capturing(): + # Never allocate inside CUDA graph capture; the reference + # path is capture-safe (just slower). + return self.forward_decode_fallback( + x2d, conv_pool, ssm_pool, slot_indices, layer_cache, ssm_state_indices + ) + buf = torch.empty( + 3, max(conv_pool.shape[0], B), d, W - 1, dtype=torch.bfloat16, device=x2d.device ) + self._cs_dense = buf else: - q_proj_states = self.q_proj(hidden_states) - k_proj_states = self.k_proj(hidden_states) - v_proj_states = self.v_proj(hidden_states) - onorm_g_hidden = ( - self.g_proj(hidden_states) - if self.use_full_rank_gate - else self.g_b_proj(self.g_a_proj(hidden_states)) + # Fail loudly if the sizing invariant ever breaks: silently + # reallocating here would hand previously captured CUDA graphs + # a dangling pointer. + assert buf.shape[1] >= B, ( + f"KDA decode staging buffer holds {buf.shape[1]} rows but the " + f"decode batch is {B}; reallocating would corrupt previously " + f"captured CUDA graphs" ) - g_hidden = self.f_b_proj(self.f_a_proj(hidden_states)) + def _project_qkvg() -> torch.Tensor: + if self._qkvg_proj_weight is not None: + return torch.nn.functional.linear(x2d, self._qkvg_proj_weight) + return self.qkvg_proj(x2d) - beta_hidden = self.b_proj(hidden_states).float() + def _project_bfa_and_fb() -> tuple[torch.Tensor, torch.Tensor]: + bfa = torch.nn.functional.linear(x2d, self._bfa_proj_weight) + f_a = bfa[:, :hd] + beta = bfa[:, hd : hd + H] + return beta, self.f_b_proj(f_a) - def _kernel_input(proj: torch.Tensor, h: int, d: int) -> torch.Tensor: - x = rearrange(proj, "b t (h d) -> t b h d", h=h, d=d) - return x.to(dtype=torch.bfloat16).contiguous() - - x_q_full = _kernel_input(q_proj_states, H, K_dim) - x_k_full = _kernel_input(k_proj_states, H, K_dim) - x_v_full = _kernel_input(v_proj_states, HV, V_dim) - g_full = _kernel_input(g_hidden, H, K_dim) - onorm_g_full = _kernel_input(onorm_g_hidden, HV, V_dim) - beta_full = rearrange(beta_hidden, "b t h -> t b h").to(torch.bfloat16).contiguous() - - w_q_t_full = ( - self.q_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() - ) - w_k_t_full = ( - self.k_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + projection_aux_stream = ( + self._projection_aux_stream if B <= _KDA_BFA_MULTISTREAM_MAX_ROWS else None ) - w_v_t_full = ( - self.v_conv1d.weight.detach().squeeze(1).transpose(0, 1).to(torch.bfloat16).contiguous() + qkvg, (beta, g) = maybe_execute_in_parallel( + _project_qkvg, + _project_bfa_and_fb, + self._projection_fork_event, + self._projection_join_event, + projection_aux_stream, + disable_on_compile=True, ) - - if cache is not None and cache.conv_state_q is not None: - hf_cs_q_pre = cache.conv_state_q.to(torch.bfloat16) - else: - hf_cs_q_pre = torch.zeros(b, projection_size, W, device=dev, dtype=torch.bfloat16) - if cache is not None and cache.conv_state_k is not None: - hf_cs_k_pre = cache.conv_state_k.to(torch.bfloat16) - else: - hf_cs_k_pre = torch.zeros(b, projection_size, W, device=dev, dtype=torch.bfloat16) - if cache is not None and cache.conv_state_v is not None: - hf_cs_v_pre = cache.conv_state_v.to(torch.bfloat16) - else: - hf_cs_v_pre = torch.zeros(b, projection_v_size, W, device=dev, dtype=torch.bfloat16) - - cs_q_full = _hf_conv_to_kernel_conv(hf_cs_q_pre, b, projection_size, W, dev, torch.bfloat16) - cs_k_full = _hf_conv_to_kernel_conv(hf_cs_k_pre, b, projection_size, W, dev, torch.bfloat16) - cs_v_full = _hf_conv_to_kernel_conv( - hf_cs_v_pre, b, projection_v_size, W, dev, torch.bfloat16 - ) - - x_q_col = q_proj_states.transpose(1, 2).to(torch.bfloat16) - x_k_col = k_proj_states.transpose(1, 2).to(torch.bfloat16) - x_v_col = v_proj_states.transpose(1, 2).to(torch.bfloat16) - new_hf_cs_q = _roll_hf_conv( - hf_cs_q_pre, x_q_col, b, projection_size, W, dev, torch.bfloat16 - ) - new_hf_cs_k = _roll_hf_conv( - hf_cs_k_pre, x_k_col, b, projection_size, W, dev, torch.bfloat16 - ) - new_hf_cs_v = _roll_hf_conv( - hf_cs_v_pre, x_v_col, b, projection_v_size, W, dev, torch.bfloat16 - ) - - if cache is not None and cache.recurrent_state is not None: - if ssm_state_indices is not None: - if self.wrong_state_layout: - raise ValueError("ssm_state_indices is incompatible with wrong_state_layout") - state_full = cache.recurrent_state - else: - state_full = cache.recurrent_state.to(dtype=torch.float32).contiguous() - else: - if ssm_state_indices is not None: - raise ValueError("ssm_state_indices requires a recurrent state pool") - state_full = torch.zeros(b, HV, V_dim, K_dim, device=dev, dtype=torch.float32) - - # The decode op requires fp32 A_log/dt_bias even in a bf16-cast module. - A_log_full = self.A_log.detach().float().contiguous() - dt_bias_full = self.dt_bias.detach().float().contiguous() - onorm_weight_full = self.o_norm.weight.detach().to(torch.float32).contiguous() - lower_bound = ( - self.gate_lower_bound_override - if self.gate_lower_bound_override is not None - else self.gate_lower_bound + x_qkv = qkvg[:, : 3 * d] + onorm_g = qkvg[:, 3 * d : 4 * d] + + # Gather the HF-layout conv windows once, then repack the + # historical W-1 columns into the kernel's dense per-section + # [B, d, W-1] layout (single strided copy kernel). + cs = conv_pool.index_select(0, slot_indices) # [B, 3d, W] + cs_dense = buf[:, :B] + cs_dense.copy_(cs.view(B, 3, d, W)[:, :, :, 1:].permute(1, 0, 2, 3)) + + state = ( + ssm_pool if ssm_state_indices is not None else ssm_pool.index_select(0, slot_indices) ) - kernel_state = ( - state_full.transpose(-1, -2).contiguous() if self.wrong_state_layout else state_full - ) - o_bfhvk = self._dispatch.decode_kda( - x_q=x_q_full, - x_k=x_k_full, - x_v=x_v_full, - w_q_t=w_q_t_full, - w_k_t=w_k_t_full, - w_v_t=w_v_t_full, + o = self._dispatch.decode_kda( + x_q=x_qkv[:, :d].unflatten(-1, (H, hd)).unsqueeze(0), + x_k=x_qkv[:, d : 2 * d].unflatten(-1, (H, hd)).unsqueeze(0), + x_v=x_qkv[:, 2 * d :].unflatten(-1, (H, hd)).unsqueeze(0), + w_q_t=self._w_q_t, + w_k_t=self._w_k_t, + w_v_t=self._w_v_t, bias_q=None, bias_k=None, bias_v=None, - cs_q=cs_q_full, - cs_k=cs_k_full, - cs_v=cs_v_full, - A_log=A_log_full, - g=g_full, - dt_bias=dt_bias_full, - beta=beta_full, - state=kernel_state, - onorm_g=onorm_g_full, - onorm_weight=onorm_weight_full, + cs_q=cs_dense[0], + cs_k=cs_dense[1], + cs_v=cs_dense[2], + A_log=self._A_log_f32, + g=g.unflatten(-1, (H, hd)).unsqueeze(0), + dt_bias=self._dt_bias_f32, + beta=beta.unsqueeze(0), + state=state, + onorm_g=onorm_g.unflatten(-1, (H, hd)).unsqueeze(0), + onorm_weight=self._onorm_w_f32, out=None, ssm_state_indices=ssm_state_indices, - cu_seqlens=None, - scale=K_dim**-0.5, + cu_seqlens=mamba_metadata._arange_buffer[: B + 1], + scale=hd**-0.5, onorm_eps=self.o_norm.eps, - lower_bound=lower_bound, + lower_bound=self.gate_lower_bound, use_beta_sigmoid_in_kernel=True, verbose=False, update_conv_cache=False, ) - state_full = ( - kernel_state.transpose(-1, -2).contiguous() if self.wrong_state_layout else kernel_state + if ssm_state_indices is None: + ssm_pool.index_copy_(0, slot_indices, state) + + # Roll the HF-layout conv pool by one token: new window = + # [old columns 1..W-1, x_new]. One cat + one scatter. + new_win = torch.cat([cs[:, :, 1:], x_qkv.unsqueeze(-1)], dim=-1) + if new_win.dtype != conv_pool.dtype: + new_win = new_win.to(conv_pool.dtype) + conv_pool.index_copy_(0, slot_indices, new_win) + # Fused-verify replay caches (spec decoding only): keep the + # committed conv window in sync with the plain-decode advance. + self._sync_kda_replay_conv_window( + layer_cache, slot_indices, new_win[:, :d], new_win[:, d : 2 * d], new_win[:, 2 * d :] ) - o_flat = rearrange(o_bfhvk, "b t h d -> b t (h d)") - o = self.o_proj(o_flat) + return self.o_proj(o.view(B, d)) - new_cache = KimiKDACachedState( - conv_state_q=new_hf_cs_q, - conv_state_k=new_hf_cs_k, - conv_state_v=new_hf_cs_v, - recurrent_state=state_full, + def forward_decode_fallback( + self, x2d, conv_pool, ssm_pool, slot_indices, layer_cache=None, ssm_state_indices=None + ) -> torch.Tensor: + """Portable or unfused decode with the production pool contract.""" + from einops import rearrange + + d = self.proj_size + x = x2d.unsqueeze(1) # [B, 1, hidden] + cs = conv_pool.index_select(0, slot_indices) + conv_q, conv_k, conv_v = _kda_split_conv_sections(cs, d) + state = ( + ssm_pool if ssm_state_indices is not None else ssm_pool.index_select(0, slot_indices) ) - return o, new_cache - # ------------------------------------------------------------------ - # Internals — FLA fallback decode (non-sm_100 path). - # ------------------------------------------------------------------ + q_proj = self.q_proj(x) + k_proj = self.k_proj(x) + v_proj = self.v_proj(x) + g_hidden = self.f_b_proj(self.f_a_proj(x)) + beta = self.b_proj(x).float() - def _decode_via_fla( - self, - hidden_states: torch.Tensor, - cache: Optional[KimiKDACachedState], - b: int, - ) -> Tuple[torch.Tensor, KimiKDACachedState]: - """FLA ``fused_recurrent_kda`` decode path — used when sm_100 is unavailable. - - Matches HF ``KimiDeltaAttention`` in ``fused_recurrent`` mode: uses - the ``ShortConvolution.step`` semantics and dispatches the delta - update to ``fla.ops.kda.fused_recurrent_kda``. - """ - q_proj_states = self.q_proj(hidden_states) - k_proj_states = self.k_proj(hidden_states) - v_proj_states = self.v_proj(hidden_states) + if self._dispatch.decode_kernel_path == "optimized": + onorm_g = self.g_proj(x) if self.use_full_rank_gate else self.g_b_proj(self.g_a_proj(x)) - conv_q_in = cache.conv_state_q if cache is not None else None - conv_k_in = cache.conv_state_k if cache is not None else None - conv_v_in = cache.conv_state_v if cache is not None else None - recurrent_in = cache.recurrent_state if cache is not None else None - - q, new_conv_q = self.q_conv1d(x=q_proj_states, cache=conv_q_in, output_final_state=True) - k, new_conv_k = self.k_conv1d(x=k_proj_states, cache=conv_k_in, output_final_state=True) - v, new_conv_v = self.v_conv1d(x=v_proj_states, cache=conv_v_in, output_final_state=True) + def _kernel_input(value: torch.Tensor) -> torch.Tensor: + return ( + rearrange(value, "b t (h d) -> t b h d", h=self.num_heads, d=self.head_dim) + .to(dtype=torch.bfloat16) + .contiguous() + ) - g_hidden = self.f_b_proj(self.f_a_proj(hidden_states)) - if self.use_full_rank_gate: - onorm_g_hidden = self.g_proj(hidden_states) + out = self._dispatch.decode_kda( + x_q=_kernel_input(q_proj), + x_k=_kernel_input(k_proj), + x_v=_kernel_input(v_proj), + w_q_t=self.q_conv1d.weight.detach() + .squeeze(1) + .transpose(0, 1) + .to(torch.bfloat16) + .contiguous(), + w_k_t=self.k_conv1d.weight.detach() + .squeeze(1) + .transpose(0, 1) + .to(torch.bfloat16) + .contiguous(), + w_v_t=self.v_conv1d.weight.detach() + .squeeze(1) + .transpose(0, 1) + .to(torch.bfloat16) + .contiguous(), + bias_q=None, + bias_k=None, + bias_v=None, + cs_q=conv_q[:, :, 1:].contiguous(), + cs_k=conv_k[:, :, 1:].contiguous(), + cs_v=conv_v[:, :, 1:].contiguous(), + A_log=self.A_log.detach().float().contiguous(), + g=_kernel_input(g_hidden), + dt_bias=self.dt_bias.detach().float().contiguous(), + beta=rearrange(beta, "b t h -> t b h").to(torch.bfloat16).contiguous(), + state=state, + onorm_g=_kernel_input(onorm_g), + onorm_weight=self.o_norm.weight.detach().float().contiguous(), + out=None, + ssm_state_indices=ssm_state_indices, + cu_seqlens=None, + scale=self.head_k_dim**-0.5, + onorm_eps=self.o_norm.eps, + lower_bound=self.gate_lower_bound, + use_beta_sigmoid_in_kernel=True, + verbose=False, + update_conv_cache=False, + ) + new_conv_q = torch.cat([conv_q[:, :, 1:], q_proj.transpose(1, 2)], dim=-1) + new_conv_k = torch.cat([conv_k[:, :, 1:], k_proj.transpose(1, 2)], dim=-1) + new_conv_v = torch.cat([conv_v[:, :, 1:], v_proj.transpose(1, 2)], dim=-1) + out = self.o_proj(out.flatten(2)) else: - onorm_g_hidden = self.g_b_proj(self.g_a_proj(hidden_states)) - beta = self.b_proj(hidden_states).float() + from fla.ops.kda import fused_recurrent_kda + + q, new_conv_q = self.q_conv1d(q_proj, cache=conv_q, output_final_state=True) + k, new_conv_k = self.k_conv1d(k_proj, cache=conv_k, output_final_state=True) + v, new_conv_v = self.v_conv1d(v_proj, cache=conv_v, output_final_state=True) + q = rearrange(q, "... (h d) -> ... h d", d=self.head_k_dim) + k = rearrange(k, "... (h d) -> ... h d", d=self.head_k_dim) + v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) + g = rearrange(g_hidden, "... (h d) -> ... h d", d=self.head_dim) + out, state = fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=self.gate_lower_bound, + state_v_first=True, + ) + out = self._output_gate_and_proj(x, out).unsqueeze(1) - g = rearrange(g_hidden, "... (h d) -> ... h d", d=self.head_dim) - q = rearrange(q, "... (h d) -> ... h d", d=self.head_k_dim) - k = rearrange(k, "... (h d) -> ... h d", d=self.head_k_dim) - v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) + conv_pool.index_copy_( + 0, + slot_indices, + torch.cat([new_conv_q, new_conv_k, new_conv_v], dim=1).to(conv_pool.dtype), + ) + if ssm_state_indices is None: + ssm_pool.index_copy_(0, slot_indices, state.to(ssm_pool.dtype)) + # Fused-verify replay caches: keep the committed conv window in + # sync with the plain-decode advance. NOTE: this path is only + # correct for requests with no pending accepted drafts + # (prev_num_accepted_tokens == 0); with drafts pending, the live + # pools lag by the pending prefix and only the fused verify kernel + # can advance them. The spec workers pad drafts to the static max, + # so drafted batches always take the verify path. + self._sync_kda_replay_conv_window( + layer_cache, + slot_indices, + new_conv_q, + new_conv_k, + new_conv_v, + ) + + return out.squeeze(1) + + def forward_verify( + self, x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices + ) -> torch.Tensor: + """Speculative verification: advance each request ``num_steps`` + tokens (1 golden + ``num_steps - 1`` padded drafts). + + Two paths: + + * Fused (``trtllm::kda_mtp_decode``, when the manager allocated the + KDA replay caches): one kernel launch replays the previous + round's accepted drafts from the per-slot replay caches, then + processes the new tokens, committing the recurrent state and conv + windows **in place** after the golden token and caching the new + drafts. ``update_mamba_states()`` afterwards only records the + accepted count for the next round's replay. + * Legacy (sequential per-step FLA): per-step states go to the + manager's batch-row-indexed intermediate scratch buffers and + ``update_mamba_states()`` promotes the accepted step's state + after sampling. + """ + if self._has_kda_replay_caches(layer_cache): + assert self._dispatch.verify_kernel_path == "optimized", ( + "KDA replay caches are allocated but the fused verify " + "kernel is unavailable; the legacy intermediate buffers " + "were not allocated so there is no fallback" + ) + return self.forward_verify_fused(x2d, num_steps, layer_cache, ssm_pool, slot_indices) + return self.forward_verify_sequential( + x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices + ) - lower_bound = ( - self.gate_lower_bound_override - if self.gate_lower_bound_override is not None - else self.gate_lower_bound + def _project_verify_inputs( + self, x: torch.Tensor, num_rows: int + ) -> Optional[ + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + Optional[torch.Tensor], + ] + ]: + """Project fused QKVG and [f_a | b] inputs for target verification.""" + qkvg_weight = self._qkvg_proj_weight + fused_qkvg = self.qkvg_proj + if qkvg_weight is None and fused_qkvg is None: + return None + + def _project_qkvg() -> torch.Tensor: + if qkvg_weight is not None: + return torch.nn.functional.linear(x, qkvg_weight) + return fused_qkvg(x) + + bfa_weight = self._bfa_proj_weight + if bfa_weight is not None: + + def _project_bfa_and_fb() -> tuple[torch.Tensor, torch.Tensor]: + bfa = torch.nn.functional.linear(x, bfa_weight) + f_a = bfa[..., : self.head_dim] + beta = bfa[..., self.head_dim : self.head_dim + self.num_heads] + return beta, self.f_b_proj(f_a) + + projection_aux_stream = ( + self._projection_aux_stream + if 0 < num_rows <= _KDA_BFA_MULTISTREAM_MAX_ROWS + else None + ) + qkvg, (beta, forget_gate) = maybe_execute_in_parallel( + _project_qkvg, + _project_bfa_and_fb, + self._projection_fork_event, + self._projection_join_event, + projection_aux_stream, + disable_on_compile=True, + ) + else: + qkvg = _project_qkvg() + beta = self.b_proj(x) + forget_gate = self.f_b_proj(self.f_a_proj(x)) + + d = self.proj_size + q_proj, k_proj, v_proj = (part.contiguous() for part in qkvg[..., : 3 * d].split(d, dim=-1)) + qkvg_split_sizes = self.qkvg_split_sizes + has_onorm_gate = qkvg_weight is not None or ( + self.use_full_rank_gate and qkvg_split_sizes is not None and len(qkvg_split_sizes) == 4 ) + onorm_g = qkvg[..., 3 * d : 4 * d].contiguous() if has_onorm_gate else None + return q_proj, k_proj, v_proj, forget_gate, beta, onorm_g - o, new_recurrent = fused_recurrent_kda( - q=q, - k=k, - v=v, + def forward_verify_fused( + self, x2d, num_steps, layer_cache, ssm_pool, slot_indices + ) -> torch.Tensor: + """Fused multi-token verify via ``trtllm::kda_mtp_decode``. + + Token layout: the kernel indexes each request's new tokens at + ``cu_seqlens[n] + num_accepted[n] + i``. The runtime packs the + ``num_steps`` new tokens per request contiguously, so we pass + ``cu_seqlens[n] = n * num_steps - num_accepted[n]`` — the shift + lands the kernel's reads/writes exactly on the packed rows. A + negative entry for request 0 is fine: ``bos`` is only ever used + additively with a token offset ``>= num_accepted``. + """ + num_decodes = x2d.shape[0] // num_steps + num_spec = num_steps - 1 + H = self.num_heads + K = self.head_k_dim + x = x2d.view(num_decodes, num_steps, -1) # [B, T, hidden] + T_total = num_decodes * num_steps + + projections = self._project_verify_inputs(x, T_total) + if projections is None: + q_proj = self.q_proj(x) + k_proj = self.k_proj(x) + v_proj = self.v_proj(x) + forget_gate = self.f_b_proj(self.f_a_proj(x)) + beta_proj = self.b_proj(x) + onorm_g = None + else: + q_proj, k_proj, v_proj, forget_gate, beta_proj, onorm_g = projections + x_q = q_proj.view(1, T_total, H, K) + x_k = k_proj.view(1, T_total, H, K) + x_v = v_proj.view(1, T_total, H, self.head_dim) + # Raw gate / beta: the kernel applies dt_bias, A_log, the + # lower-bound sigmoid gate, and the beta sigmoid itself. + g = forget_gate.view(1, T_total, H, K) + beta = beta_proj.contiguous().view(1, T_total, H) + + w_q, w_k, w_v = self._get_mtp_conv_weights() + lower_bound = self.gate_lower_bound + + pending = layer_cache.prev_num_accepted_tokens[slot_indices].to( + torch.int32 + ) # accepted drafts of the previous round, per req + cu_seqlens = torch.arange( + 0, (num_decodes + 1) * num_steps, num_steps, dtype=torch.int32, device=x2d.device + ) + cu_seqlens[:num_decodes].sub_(pending) + + out = self._dispatch.mtp_verify( + x_q=x_q, + x_k=x_k, + x_v=x_v, + w_q=w_q, + w_k=w_k, + w_v=w_v, + cs_q=layer_cache.kda_conv_q, + cs_k=layer_cache.kda_conv_k, + cs_v=layer_cache.kda_conv_v, g=g, beta=beta, - A_log=self.A_log, - dt_bias=self.dt_bias, - initial_state=recurrent_in, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - use_beta_sigmoid_in_kernel=True, + # .detach(): the CuTe DSL DLPack bridge rejects grad-tracking + # tensors. + A_log=self.A_log.detach(), + dt_bias=self.dt_bias.detach(), + recurrent_state=ssm_pool, + qkg_cache=layer_cache.kda_qkg_cache, + v_cache=layer_cache.kda_v_cache, + beta_cache=layer_cache.kda_beta_cache, + ssm_state_indices=slot_indices.to(torch.int32), + cu_seqlens=cu_seqlens, + num_spec=num_spec, + num_accepted_tokens=pending, lower_bound=lower_bound, - state_v_first=True, + scale=self.head_k_dim**-0.5, + ) + o = out.view(num_decodes, num_steps, H, self.head_dim) + return self._output_gate_and_proj(x, o, onorm_g) + + def _build_mtp_conv_weights(self) -> None: + """Prebuild the fp32 ``[dim, W]`` conv weights for the fused verify + kernel (once, at weight-load finalize time). Building them lazily + at first use would allocate at runtime; under CUDA graph capture + that bakes capture-pool pointers into the cached tuple.""" + self._mtp_conv_weights = tuple( + conv.weight.detach().squeeze(1).float().contiguous() + for conv in (self.q_conv1d, self.k_conv1d, self.v_conv1d) ) - onorm_g = rearrange(onorm_g_hidden, "... (h d) -> ... h d", d=self.head_dim) - o = self.o_norm(o, onorm_g) - o = rearrange(o, "b t h d -> b t (h d)") - o = self.o_proj(o) + def _get_mtp_conv_weights(self) -> Tuple[torch.Tensor, ...]: + """fp32 ``[dim, W]`` conv weights for the fused verify kernel, + prebuilt by ``_build_mtp_conv_weights()``.""" + cached = self._mtp_conv_weights + if cached is None: + raise RuntimeError( + "Kimi K3 fused-verify conv weights were not prebuilt; call " + "_build_mtp_conv_weights() (done by load_weights() and by " + "finalize_decode_weights() / finalize_decode_weights_fp8()) " + "after weight load and before the first verify step." + ) + return cached - new_cache = KimiKDACachedState( - conv_state_q=new_conv_q, - conv_state_k=new_conv_k, - conv_state_v=new_conv_v, - recurrent_state=new_recurrent, + def forward_verify_sequential( + self, x2d, num_steps, layer_cache, conv_pool, ssm_pool, slot_indices + ) -> torch.Tensor: + """Sequential per-step FLA verification (legacy intermediate-buffer + path). Live pools are read-only here; ``update_mamba_states()`` + commits the accepted step's state after sampling. + """ + from einops import rearrange + from fla.ops.kda import fused_recurrent_kda + + intermediate_conv = layer_cache.intermediate_conv_window + intermediate_ssm = layer_cache.intermediate_ssm + assert intermediate_conv is not None and intermediate_ssm is not None, ( + "speculative verification requires the cache manager's " + "SpeculativeState (legacy intermediate-buffer path)" ) - return o, new_cache - # ------------------------------------------------------------------ - # Weight helper for random-weight parity tests. - # ------------------------------------------------------------------ + d = self.proj_size + num_decodes = x2d.shape[0] // num_steps + x = x2d.view(num_decodes, num_steps, -1) # [B, T, hidden] + + projections = self._project_verify_inputs(x, x2d.shape[0]) + if projections is None: + q_proj_states = self.q_proj(x) + k_proj_states = self.k_proj(x) + v_proj_states = self.v_proj(x) + g = self.f_b_proj(self.f_a_proj(x)) + beta = self.b_proj(x).float() + onorm_g = None + else: + q_proj_states, k_proj_states, v_proj_states, g, beta, onorm_g = projections + beta = beta.float() + g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim) - def copy_weights_from(self, source: nn.Module) -> "dict[str, Tuple[Tuple[int, ...], str]]": - """Copy every named parameter/buffer from ``source`` into ``self``. + # Gathered copies — mutated across steps, never written back to the + # live pools. + cs = conv_pool.index_select(0, slot_indices) + conv_q, conv_k, conv_v = _kda_split_conv_sections(cs, d) + state = ssm_pool.index_select(0, slot_indices) + + step_outputs: List[torch.Tensor] = [] + for t in range(num_steps): + # ShortConvolution.step updates the (gathered) caches in place. + q_t, conv_q = self.q_conv1d( + q_proj_states[:, t : t + 1], cache=conv_q, output_final_state=True + ) + k_t, conv_k = self.k_conv1d( + k_proj_states[:, t : t + 1], cache=conv_k, output_final_state=True + ) + v_t, conv_v = self.v_conv1d( + v_proj_states[:, t : t + 1], cache=conv_v, output_final_state=True + ) - Because ``KimiKDALinearAttention`` mirrors the HF reference's - parameter names 1:1, the mapping is identity: every source name is - assigned to the identically named target. Shape mismatches raise - loudly. Returns a ``{name: (shape, dtype)}`` provenance dict. - """ - src: dict[str, torch.Tensor] = {} - for name, p in source.named_parameters(recurse=True): - src[name] = p.data - for name, buf in source.named_buffers(recurse=True): - src[name] = buf - - dst: dict[str, torch.Tensor] = {} - for name, p in self.named_parameters(recurse=True): - dst[name] = p.data - for name, buf in self.named_buffers(recurse=True): - dst[name] = buf - - missing_on_dst = sorted(set(src) - set(dst)) - missing_on_src = sorted(set(dst) - set(src)) - if missing_on_dst: - raise KeyError( - f"copy_weights_from: source params missing on target: {missing_on_dst[:5]}" + q_t = rearrange(q_t, "... (h d) -> ... h d", d=self.head_k_dim) + k_t = rearrange(k_t, "... (h d) -> ... h d", d=self.head_k_dim) + v_t = rearrange(v_t, "... (h d) -> ... h d", d=self.head_dim) + + o_t, state = fused_recurrent_kda( + q=q_t, + k=k_t, + v=v_t, + g=g[:, t : t + 1], + beta=beta[:, t : t + 1], + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=self.gate_lower_bound, + state_v_first=True, ) - if missing_on_src: - raise KeyError( - f"copy_weights_from: target params missing on source: {missing_on_src[:5]}" + step_outputs.append(o_t) + + # Batch-row indexed ([:num_decodes] prefix), matching + # update_mamba_states()'s intermediate_state_indices. + intermediate_conv[:num_decodes, t] = torch.cat([conv_q, conv_k, conv_v], dim=1).to( + intermediate_conv.dtype ) + intermediate_ssm[:num_decodes, t] = state.to(intermediate_ssm.dtype) - provenance: "dict[str, Tuple[Tuple[int, ...], str]]" = {} - for name, srct in src.items(): - dstt = dst[name] - if srct.shape != dstt.shape: - raise ValueError( - f"shape mismatch for {name}: source {tuple(srct.shape)} " - f"vs target {tuple(dstt.shape)}" - ) - dstt.copy_(srct.to(dtype=dstt.dtype, device=dstt.device)) - provenance[name] = (tuple(srct.shape), str(srct.dtype)) - return provenance + o = torch.cat(step_outputs, dim=1) # [B, T, H, V] + return self._output_gate_and_proj(x, o, onorm_g) + + def _output_gate_and_proj( + self, x: torch.Tensor, o: torch.Tensor, onorm_g: Optional[torch.Tensor] = None + ) -> torch.Tensor: + from einops import rearrange + + if onorm_g is not None: + g_out = onorm_g + elif self.use_full_rank_gate: + g_out = self.g_proj(x) + else: + g_out = self.g_b_proj(self.g_a_proj(x)) + g_out = rearrange(g_out, "... (h d) -> ... h d", d=self.head_dim) + o = self.o_norm(o, g_out) + o = rearrange(o, "b t h d -> (b t) (h d)") + return self.o_proj(o) diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py b/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py index 53237af8992a..17f91cab2f72 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import nvtx from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate, Deepseekv3MoE @@ -5,7 +8,6 @@ from tensorrt_llm._torch.models.modeling_kimi_linear import ( KimiK3MoEGate, KimiK3MoERuntime, - KimiKDARuntime, KimiMLARuntime, ) from tensorrt_llm._torch.models.modeling_nemotron_h import MLPLayer, NemotronHMOE @@ -16,6 +18,7 @@ from tensorrt_llm._torch.modules.attention import Attention from tensorrt_llm._torch.modules.fused_moe.interface import MoE from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention from tensorrt_llm._torch.modules.mamba.mamba2_mixer import Mamba2Mixer from tensorrt_llm._torch.modules.mhc.hyper_connection import mHC from tensorrt_llm._torch.modules.mla import MLA @@ -32,10 +35,13 @@ def mark_ranges(): Qwen3NextSparseMoeBlock.forward = nvtx.annotate("Qwen3NextSparseMoeBlock")( Qwen3NextSparseMoeBlock.forward ) - # Kimi K3. `KimiK3MLAAttention` overrides `MLA.forward`, so the range below - # is on its `KimiMLARuntime` wrapper. The gate is entered through - # `compute_logits`, not `forward`. Its MLPs are the shared `GatedMLP`. - KimiKDARuntime.forward = nvtx.annotate("KimiKDARuntime")(KimiKDARuntime.forward) + # Kimi K3. KDA runs directly through `KimiKDALinearAttention`. + # `KimiK3MLAAttention` overrides `MLA.forward`, so its range is on the + # `KimiMLARuntime` wrapper. The gate is entered through `compute_logits`, + # not `forward`. Its MLPs are the shared `GatedMLP`. + KimiKDALinearAttention.forward = nvtx.annotate("KimiKDALinearAttention")( + KimiKDALinearAttention.forward + ) KimiMLARuntime.forward = nvtx.annotate("KimiMLARuntime")(KimiMLARuntime.forward) KimiK3MoERuntime.forward = nvtx.annotate("KimiK3MoERuntime")(KimiK3MoERuntime.forward) KimiK3MoEGate.compute_logits = nvtx.annotate("KimiK3MoEGate")(KimiK3MoEGate.compute_logits) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index f1b02c22217e..8d9baa61d2cf 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -95,8 +95,8 @@ l0_b200: - unittest/_torch/modules/test_mhc.py - unittest/_torch/modules/test_engram.py # ------------- Kimi K3 (KimiLinear) unit tests --------------- - - unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py - - unittest/_torch/modeling/test_kimi_kda_verify_parity.py + - unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py + - unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py - unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py - unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py - unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index e065a77cd043..d98a3f71c8cc 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -38,12 +38,12 @@ l0_gb300_multi_gpus: # Kimi K3 speculative decoding: exact kernel-level verify parity (random # weights, no checkpoint) + truncated-model e2e SA logits parity # (auto-skips when the Kimi K3 checkpoint is not staged). - - unittest/_torch/modeling -k "kimi_kda_verify" + - unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py # Fused KDA multi-token verify (trtllm::kda_mtp_decode): kernel-level # parity vs CPU golden + FLA sequential, and runtime-level fused-vs- # sequential two-round replay parity. - - unittest/_torch/modeling -k "kda_mtp_decode_cute_parity" - - unittest/_torch/modeling -k "kimi_kda_fused_verify" + - unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py + - unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py # Kimi K3 op parity tests (random weights, no checkpoint; SM100/SM103). - unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py - unittest/_torch/modules/kimi_kda/test_kda_decode_op.py diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py b/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py index 4bd8dd4efced..806afdec4ebb 100644 --- a/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py +++ b/tests/unittest/_torch/modeling/test_kimi_kda_fp8_packed_prefill.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Runtime parity for the Kimi K3 FP8 packed q/k/v prefill projection.""" +"""Parity for the Kimi K3 FP8 packed q/k/v prefill projection.""" from collections.abc import Callable from types import SimpleNamespace @@ -12,9 +12,9 @@ pytest.importorskip("fla") from tensorrt_llm._torch.models.modeling_kimi_linear import ( - KimiKDARuntime, _convert_kda_projections_to_fp8_weight_read, ) +from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention class _Cfg: @@ -30,16 +30,16 @@ class _Cfg: class _Layer(nn.Module): - def __init__(self, runtime: KimiKDARuntime) -> None: + def __init__(self, attention: KimiKDALinearAttention) -> None: super().__init__() self.is_kda = True - self.self_attn = runtime + self.linear_attn = attention class _Model(nn.Module): - def __init__(self, runtime: KimiKDARuntime) -> None: + def __init__(self, attention: KimiKDALinearAttention) -> None: super().__init__() - self.layers = nn.ModuleList([_Layer(runtime)]) + self.layers = nn.ModuleList([_Layer(attention)]) def _has_supported_gpu() -> bool: @@ -52,10 +52,10 @@ def _has_supported_gpu() -> bool: ) -def _make_runtime() -> KimiKDARuntime: - runtime = KimiKDARuntime(_Cfg(), layer_idx=0).to("cuda") - assert _convert_kda_projections_to_fp8_weight_read(_Model(runtime)) == 5 - return runtime +def _make_attention() -> KimiKDALinearAttention: + attention = KimiKDALinearAttention(_Cfg(), layer_idx=0).to("cuda") + assert _convert_kda_projections_to_fp8_weight_read(_Model(attention)) == 5 + return attention def _assert_numerically_close(actual: torch.Tensor, expected: torch.Tensor) -> None: @@ -70,13 +70,14 @@ def _assert_numerically_close(actual: torch.Tensor, expected: torch.Tensor) -> N @torch.no_grad() def test_fp8_packed_qkv_projection_matches_separate_views() -> None: torch.manual_seed(0) - runtime = _make_runtime() - mixer = runtime.mixer + attention = _make_attention() hidden = torch.randn(1, 193, _Cfg.hidden_size, device="cuda", dtype=torch.bfloat16) * 0.05 - packed = mixer.qkvg_proj(hidden)[..., : 3 * runtime.proj_size] - actual = packed.split(runtime.proj_size, dim=-1) - expected = (mixer.q_proj(hidden), mixer.k_proj(hidden), mixer.v_proj(hidden)) + qkvg_proj = attention.qkvg_proj + assert qkvg_proj is not None + packed = qkvg_proj(hidden)[..., : 3 * attention.proj_size] + actual = packed.split(attention.proj_size, dim=-1) + expected = (attention.q_proj(hidden), attention.k_proj(hidden), attention.v_proj(hidden)) for packed_part, separate_part in zip(actual, expected): _assert_numerically_close(packed_part, separate_part) @@ -94,12 +95,11 @@ def test_fp8_packed_qkv_prefill_matches_separate_path_and_updates_state( sequence_lengths: list[int], use_initial_states: bool, has_initial_states: list[bool] ) -> None: torch.manual_seed(1) - runtime = _make_runtime() - mixer = runtime.mixer + attention = _make_attention() num_prefills = len(sequence_lengths) num_tokens = sum(sequence_lengths) slots = num_prefills + 3 - d = runtime.proj_size + d = attention.proj_size h = _Cfg.linear_attn_config["num_heads"] head_dim = _Cfg.linear_attn_config["head_dim"] conv_size = _Cfg.linear_attn_config["short_conv_kernel_size"] @@ -126,18 +126,19 @@ def _hook(_module: nn.Module, _inputs: tuple, _output: object) -> None: return _hook + fused_qkvg = attention.qkvg_proj + assert fused_qkvg is not None handles = [ - mixer.qkvg_proj.register_forward_hook(_count("qkvg")), - mixer.q_proj.register_forward_hook(_count("q")), - mixer.k_proj.register_forward_hook(_count("k")), - mixer.v_proj.register_forward_hook(_count("v")), + fused_qkvg.register_forward_hook(_count("qkvg")), + attention.q_proj.register_forward_hook(_count("q")), + attention.k_proj.register_forward_hook(_count("k")), + attention.v_proj.register_forward_hook(_count("v")), ] - fused_qkvg = mixer.qkvg_proj try: - mixer.qkvg_proj = None + attention.qkvg_proj = None ref_conv = conv_seed.clone() ref_state = state_seed.clone() - expected = runtime._forward_prefill( + expected = attention.forward_prefill( hidden, cu_seqlens, metadata, @@ -149,10 +150,10 @@ def _hook(_module: nn.Module, _inputs: tuple, _output: object) -> None: assert calls == {"qkvg": 0, "q": 1, "k": 1, "v": 1} calls.update(qkvg=0, q=0, k=0, v=0) - mixer.qkvg_proj = fused_qkvg + attention.qkvg_proj = fused_qkvg actual_conv = conv_seed.clone() actual_state = state_seed.clone() - actual = runtime._forward_prefill( + actual = attention.forward_prefill( hidden, cu_seqlens, metadata, @@ -163,7 +164,7 @@ def _hook(_module: nn.Module, _inputs: tuple, _output: object) -> None: ) assert calls == {"qkvg": 1, "q": 0, "k": 0, "v": 0} finally: - mixer.qkvg_proj = fused_qkvg + attention.qkvg_proj = fused_qkvg for handle in handles: handle.remove() @@ -189,7 +190,7 @@ def _hook(_module: nn.Module, _inputs: tuple, _output: object) -> None: repeat_conv = conv_seed.clone() repeat_state = state_seed.clone() - repeated = runtime._forward_prefill( + repeated = attention.forward_prefill( hidden, cu_seqlens, metadata, diff --git a/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py b/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py new file mode 100644 index 000000000000..ae87fd0ae131 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_kimi_linear_checkpoint.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Checkpoint-name tests for the Kimi Linear model.""" + +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("fla") + +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearForCausalLM # noqa: E402 + + +def test_checkpoint_plan_preserves_external_attention_names(): + class _PlanHarness: + checkpoint_name_plan = KimiLinearForCausalLM.checkpoint_name_plan + model = SimpleNamespace(layers=[]) + + def _trunk_parameters(self): + return { + "model.layers.0.linear_attn.q_proj.weight": torch.empty(0), + "model.layers.1.self_attn.mixer.q_a_proj.weight": torch.empty(0), + "lm_head.weight": torch.empty(0), + } + + name_map, expected_keys, expert_jobs = _PlanHarness().checkpoint_name_plan("language_model.") + + assert name_map == { + "model.layers.0.linear_attn.q_proj.weight": ( + "language_model.model.layers.0.self_attn.q_proj.weight" + ), + "model.layers.1.self_attn.mixer.q_a_proj.weight": ( + "language_model.model.layers.1.self_attn.q_a_proj.weight" + ), + "lm_head.weight": "language_model.lm_head.weight", + } + assert expected_keys == set(name_map.values()) + assert expert_jobs == [] diff --git a/tests/unittest/_torch/modules/kimi_kda/kimi_kda_test_utils.py b/tests/unittest/_torch/modules/kimi_kda/kimi_kda_test_utils.py new file mode 100644 index 000000000000..5c7666acc9ca --- /dev/null +++ b/tests/unittest/_torch/modules/kimi_kda/kimi_kda_test_utils.py @@ -0,0 +1,351 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FLA-only Kimi KDA parity reference. + +This is a structural mirror of the HF reference ``KimiDeltaAttention`` in +``modeling_kimi.py``. Same parameter names, same layer shapes, same short +convolution + FLA gating + FusedRMSNormGated output-gate stack. The +delta-rule inner loop calls FLA directly so optimized production kernels +are always compared against an independent implementation. Production code +must use ``tensorrt_llm._torch.modules.kimi_kda.KimiKDALinearAttention``. + +Cache ownership +--------------- +KDA carries three short-convolution states (``conv_state_{q,k,v}``, HF +layout ``[B, D, W]`` bf16) and one delta-rule recurrent state +(``recurrent_state``, layout ``[B, HV, V, K]`` fp32, matching the optimized +kernel's transposed convention). These match the hybrid-cache ownership +pattern used by the mamba modules. The reference consumes and returns state +tensors directly so module-level tests can prove state roundtrip without +runtime cache-manager plumbing. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +from einops import rearrange +from fla.modules import FusedRMSNormGated, ShortConvolution +from fla.ops.kda import chunk_kda, fused_recurrent_kda +from torch import nn + +from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import KimiKDALinearAttention + + +def get_production_prefill_kernel_path(attention: KimiKDALinearAttention) -> str: + """Return the selected production prefill path for kernel-routing tests.""" + return attention._dispatch.prefill_kernel_path + + +def get_production_decode_kernel_path(attention: KimiKDALinearAttention) -> str: + """Return the selected production decode path for kernel-routing tests.""" + return attention._dispatch.decode_kernel_path + + +def _meta_safe_cast_dtype(module, dtype): + """``module.to(dtype=dtype)`` that also works under ``MetaInitMode``. + + ``Module.to`` dispatches ``aten._to_copy``, which MetaInitMode rejects + (it would silently fall back to full CPU construction of the model — + ~70 GB of host RAM per rank for Kimi K3). Under meta init the values + are garbage anyway, so a dtype-only re-allocation via ``empty_like`` + (an allowed init op) is equivalent; off meta this matches ``.to``. + """ + import torch as _torch + + def _cast(t): + if not t.is_floating_point(): + return t + if t.is_meta: + return _torch.empty_like(t, dtype=dtype) + return t.to(dtype=dtype) + + module._apply(_cast) + + +class _MetaSafeFusedRMSNormGated(FusedRMSNormGated): + """FusedRMSNormGated whose init survives the model loader's MetaInitMode. + + ``FusedRMSNormGated.reset_parameters`` uses ``nn.init.ones_`` (a plain + ``fill_``), which MetaInitMode rejects on meta tensors and which would + force the whole model construction to fall back to eager CPU init. + ``uniform_(1, 1)`` produces identical values and is on MetaInitMode's + random-init allowlist. + """ + + def reset_parameters(self) -> None: + if self.elementwise_affine: + with torch.no_grad(): + self.weight.uniform_(1.0, 1.0) + + +@dataclass +class KimiKDATestCachedState: + """Per-layer KDA cache tensors in HF layout. + + ``conv_state_*`` — shape ``[B, D, W]`` bf16 where ``W`` is + ``short_conv_kernel_size`` and the newest processed token sits at + position ``W-1``. ``D`` is ``num_heads * head_dim`` for q/k and + ``num_heads * head_dim`` for v (K3 uses HV == H). + + ``recurrent_state`` — shape ``[B, HV, V, K]`` fp32. This is the + transposed layout the optimized KDA kernels expect, and it is the same + layout HF stores when running with ``transpose_state_layout=True``. + ``None`` fields are treated as zero. + """ + + conv_state_q: Optional[torch.Tensor] + conv_state_k: Optional[torch.Tensor] + conv_state_v: Optional[torch.Tensor] + recurrent_state: Optional[torch.Tensor] + + +class KimiKDAReference(nn.Module): + """Standalone FLA Kimi K3 linear-attention parity reference. + + Parameters + ---------- + hidden_size : int + num_heads : int + head_dim : int + conv_kernel_size : int + use_full_rank_gate : bool + gate_lower_bound : Optional[float] + rms_norm_eps : float + dtype : Optional[torch.dtype] + layer_idx : int + """ + + def __init__( + self, + *, + hidden_size: int, + num_heads: int, + head_dim: int, + conv_kernel_size: int, + use_full_rank_gate: bool, + gate_lower_bound: Optional[float], + rms_norm_eps: float = 1e-5, + dtype: Optional[torch.dtype] = None, + layer_idx: int = 0, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = head_dim + self.head_k_dim = head_dim + self.num_k_heads = num_heads + self.conv_size = conv_kernel_size + self.use_full_rank_gate = use_full_rank_gate + self.gate_lower_bound = gate_lower_bound + self.rms_norm_eps = rms_norm_eps + self.layer_idx = layer_idx + + projection_k_size = self.head_k_dim * self.num_k_heads + projection_size = self.head_dim * self.num_heads + + self.q_proj = nn.Linear(hidden_size, projection_k_size, bias=False) + self.k_proj = nn.Linear(hidden_size, projection_k_size, bias=False) + self.v_proj = nn.Linear(hidden_size, projection_size, bias=False) + + self.q_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=conv_kernel_size, + activation="silu", + ) + self.k_conv1d = ShortConvolution( + hidden_size=projection_k_size, + kernel_size=conv_kernel_size, + activation="silu", + ) + self.v_conv1d = ShortConvolution( + hidden_size=projection_size, + kernel_size=conv_kernel_size, + activation="silu", + ) + + self.A_log = nn.Parameter( + torch.log(torch.empty(num_heads, dtype=torch.float32).uniform_(1, 16)) + ) + self.f_a_proj = nn.Linear(hidden_size, head_dim, bias=False) + self.f_b_proj = nn.Linear(head_dim, projection_size, bias=False) + self.dt_bias = nn.Parameter( + torch.empty(projection_size, dtype=torch.float32).uniform_( + math.log(1e-3), math.log(1e-1) + ) + ) + self.b_proj = nn.Linear(hidden_size, num_heads, bias=False) + + if use_full_rank_gate: + self.g_proj = nn.Linear(hidden_size, projection_size, bias=False) + else: + self.g_a_proj = nn.Linear(hidden_size, head_dim, bias=False) + self.g_b_proj = nn.Linear(head_dim, projection_size, bias=False) + + self.o_norm = _MetaSafeFusedRMSNormGated(head_dim, eps=rms_norm_eps, activation="sigmoid") + self.o_proj = nn.Linear(projection_size, hidden_size, bias=False) + + if dtype is not None: + _meta_safe_cast_dtype(self, dtype) + + # ------------------------------------------------------------------ + # Prefill entry (Goal 2.1 pass path). + # ------------------------------------------------------------------ + + def forward_prefill( + self, + hidden_states: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Prefill forward matching HF ``KimiDeltaAttention.forward`` in chunk mode. + + Parameters + ---------- + hidden_states : ``(B, T, hidden_size)`` for equal-length prefill or + ``(1, sum(seq_lens), hidden_size)`` when ``cu_seqlens`` is given. + cu_seqlens : optional cumulative sequence lengths for varlen inputs. + + Returns + ------- + ``(B, T, hidden_size)`` output tensor (equal-length case) or + ``(1, sum(seq_lens), hidden_size)`` (varlen case). + """ + if cu_seqlens is not None: + cu_seqlens = cu_seqlens.to(device=hidden_states.device, dtype=torch.long) + + q_proj_states = self.q_proj(hidden_states) + k_proj_states = self.k_proj(hidden_states) + v_proj_states = self.v_proj(hidden_states) + + q, _ = self.q_conv1d( + x=q_proj_states, + cache=None, + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + k, _ = self.k_conv1d( + x=k_proj_states, + cache=None, + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + v, _ = self.v_conv1d( + x=v_proj_states, + cache=None, + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + + g = self.f_b_proj(self.f_a_proj(hidden_states)) + g = rearrange(g, "... (h d) -> ... h d", d=self.head_dim) + beta = self.b_proj(hidden_states).float() + + q = rearrange(q, "... (h d) -> ... h d", d=self.head_k_dim) + k = rearrange(k, "... (h d) -> ... h d", d=self.head_k_dim) + v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) + + lower_bound = self.gate_lower_bound + safe_gate = lower_bound is not None + scale = self.head_k_dim**-0.5 + + o, _final_state = chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + scale=scale, + initial_state=None, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + safe_gate=safe_gate, + lower_bound=lower_bound, + state_v_first=True, + cu_seqlens=cu_seqlens, + ) + + if self.use_full_rank_gate: + g_out = self.g_proj(hidden_states) + else: + g_out = self.g_b_proj(self.g_a_proj(hidden_states)) + g_out = rearrange(g_out, "... (h d) -> ... h d", d=self.head_dim) + o = self.o_norm(o, g_out) + + o = rearrange(o, "b t h d -> b t (h d)") + o = self.o_proj(o) + return o + + def forward_decode( + self, + hidden_states: torch.Tensor, + cache: Optional[KimiKDATestCachedState] = None, + ) -> Tuple[torch.Tensor, KimiKDATestCachedState]: + """Run the FLA T=1 cached-decode reference. + + ``hidden_states`` shape ``(B, 1, hidden_size)``. Cache is + ``KimiKDATestCachedState`` in HF layout; ``None`` fields become zero + tensors. + """ + _, q_len, _ = hidden_states.shape + assert q_len == 1, f"KimiKDAReference.forward_decode expects T=1, got T={q_len}" + q_proj_states = self.q_proj(hidden_states) + k_proj_states = self.k_proj(hidden_states) + v_proj_states = self.v_proj(hidden_states) + + conv_q_in = cache.conv_state_q if cache is not None else None + conv_k_in = cache.conv_state_k if cache is not None else None + conv_v_in = cache.conv_state_v if cache is not None else None + recurrent_in = cache.recurrent_state if cache is not None else None + + q, new_conv_q = self.q_conv1d(x=q_proj_states, cache=conv_q_in, output_final_state=True) + k, new_conv_k = self.k_conv1d(x=k_proj_states, cache=conv_k_in, output_final_state=True) + v, new_conv_v = self.v_conv1d(x=v_proj_states, cache=conv_v_in, output_final_state=True) + + g_hidden = self.f_b_proj(self.f_a_proj(hidden_states)) + if self.use_full_rank_gate: + onorm_g_hidden = self.g_proj(hidden_states) + else: + onorm_g_hidden = self.g_b_proj(self.g_a_proj(hidden_states)) + beta = self.b_proj(hidden_states).float() + + g = rearrange(g_hidden, "... (h d) -> ... h d", d=self.head_dim) + q = rearrange(q, "... (h d) -> ... h d", d=self.head_k_dim) + k = rearrange(k, "... (h d) -> ... h d", d=self.head_k_dim) + v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim) + + o, new_recurrent = fused_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=self.A_log, + dt_bias=self.dt_bias, + initial_state=recurrent_in, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + lower_bound=self.gate_lower_bound, + state_v_first=True, + ) + + onorm_g = rearrange(onorm_g_hidden, "... (h d) -> ... h d", d=self.head_dim) + o = self.o_norm(o, onorm_g) + o = rearrange(o, "b t h d -> b t (h d)") + o = self.o_proj(o) + + new_cache = KimiKDATestCachedState( + conv_state_q=new_conv_q, + conv_state_k=new_conv_k, + conv_state_v=new_conv_v, + recurrent_state=new_recurrent, + ) + return o, new_cache diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_decode_op.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_decode_op.py index e4ab388a394a..aa62fc29f02d 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_decode_op.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_decode_op.py @@ -3,6 +3,7 @@ """Parity tests for the optimized Kimi K3 KDA decode op.""" import copy +from types import SimpleNamespace import pytest import torch @@ -10,10 +11,11 @@ pytest.importorskip("fla") -from tensorrt_llm._torch.modules.kimi_kda import _kda_decode # noqa: E402 -from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import ( # noqa: E402 - KimiKDACachedState, - KimiKDALinearAttention, +from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention, _kda_decode # noqa: E402 +from tests.unittest._torch.modules.kimi_kda.kimi_kda_test_utils import ( # noqa: E402 + KimiKDAReference, + KimiKDATestCachedState, + get_production_decode_kernel_path, ) # 73: deliberately odd and > 64 to cover non-power-of-two batched decode. @@ -36,7 +38,9 @@ def _has_supported_gpu() -> bool: ) -def _make_attention_pair() -> tuple[KimiKDALinearAttention, KimiKDALinearAttention]: +def _make_attention_pair( + *, finalize_decode_weights: bool = True +) -> tuple[KimiKDALinearAttention, KimiKDAReference]: common = { "hidden_size": HIDDEN_SIZE, "num_heads": NUM_HEADS, @@ -47,19 +51,32 @@ def _make_attention_pair() -> tuple[KimiKDALinearAttention, KimiKDALinearAttenti "rms_norm_eps": 1e-5, "dtype": torch.bfloat16, } - optimized = KimiKDALinearAttention(**common).to("cuda") - reference = KimiKDALinearAttention(**common, use_optimized_decode=False).to("cuda") + cfg = SimpleNamespace( + hidden_size=HIDDEN_SIZE, + rms_norm_eps=common["rms_norm_eps"], + linear_attn_config={ + "num_heads": NUM_HEADS, + "head_dim": HEAD_DIM, + "short_conv_kernel_size": CONV_KERNEL_SIZE, + "use_full_rank_gate": common["use_full_rank_gate"], + "gate_lower_bound": common["gate_lower_bound"], + }, + ) + optimized = KimiKDALinearAttention(cfg, layer_idx=0).to("cuda") + with torch.no_grad(): + optimized.dt_bias.zero_() + reference = KimiKDAReference(**common).to("cuda") reference.load_state_dict(optimized.state_dict()) + if finalize_decode_weights: + optimized.finalize_decode_weights() - assert optimized.decode_kernel_path == "optimized" - assert reference.decode_kernel_path == "fla" - assert reference.prefill_kernel_path == optimized.prefill_kernel_path + assert get_production_decode_kernel_path(optimized) == "optimized" return optimized, reference -def _make_cache(batch_size: int = BATCH_SIZE) -> KimiKDACachedState: +def _make_cache(batch_size: int = BATCH_SIZE) -> KimiKDATestCachedState: projection_size = NUM_HEADS * HEAD_DIM - return KimiKDACachedState( + return KimiKDATestCachedState( conv_state_q=( torch.randn( batch_size, @@ -104,6 +121,57 @@ def _make_cache(batch_size: int = BATCH_SIZE) -> KimiKDACachedState: ) +def _run_production_decode( + attention: KimiKDALinearAttention, + hidden_states: torch.Tensor, + initial_cache: KimiKDATestCachedState, + *, + conv_pool: torch.Tensor | None = None, + state_pool: torch.Tensor | None = None, + slot_indices: torch.Tensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + include_metadata: bool = True, +) -> tuple[torch.Tensor, KimiKDATestCachedState]: + batch_size = hidden_states.shape[0] + projection_size = NUM_HEADS * HEAD_DIM + if slot_indices is None: + slot_indices = torch.arange(batch_size, device="cuda", dtype=torch.long) + if conv_pool is None: + conv_pool = torch.cat( + [ + initial_cache.conv_state_q, + initial_cache.conv_state_k, + initial_cache.conv_state_v, + ], + dim=1, + ).clone() + if state_pool is None: + state_pool = initial_cache.recurrent_state.clone() + + metadata = ( + SimpleNamespace( + _arange_buffer=torch.arange(batch_size + 1, device="cuda", dtype=torch.int32) + ) + if include_metadata + else None + ) + output = attention.forward_decode( + hidden_states.squeeze(1), + conv_pool, + state_pool, + slot_indices, + metadata, + ssm_state_indices=ssm_state_indices, + ) + selected_conv = conv_pool.index_select(0, slot_indices) + return output.unsqueeze(1), KimiKDATestCachedState( + conv_state_q=selected_conv[:, :projection_size], + conv_state_k=selected_conv[:, projection_size : 2 * projection_size], + conv_state_v=selected_conv[:, 2 * projection_size :], + recurrent_state=state_pool.index_select(0, slot_indices), + ) + + def _assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None: actual_float = actual.float() expected_float = expected.float() @@ -132,9 +200,61 @@ def test_optimized_decode_matches_fla_reference(batch_size: int) -> None: ) initial_cache = _make_cache(batch_size) - actual_output, actual_cache = optimized.forward_decode( + actual_output, actual_cache = _run_production_decode( + optimized, hidden_states, copy.deepcopy(initial_cache) + ) + expected_output, expected_cache = reference.forward_decode( hidden_states, copy.deepcopy(initial_cache) ) + + _assert_close(actual_output, expected_output) + _assert_close(actual_cache.recurrent_state, expected_cache.recurrent_state) + _assert_close(actual_cache.conv_state_q, expected_cache.conv_state_q) + _assert_close(actual_cache.conv_state_k, expected_cache.conv_state_k) + _assert_close(actual_cache.conv_state_v, expected_cache.conv_state_v) + + +@torch.no_grad() +@pytest.mark.parametrize( + ("fallback_case", "batch_size"), + ( + ("unfused-projections", 1), + ("unfused-projections", BATCH_SIZE), + ("missing-metadata", 1), + ("capture-before-staging", 1), + ), + ids=( + "unfused-projections-b1", + "unfused-projections-odd-large-batch", + "missing-metadata", + "capture-before-staging", + ), +) +def test_decode_fallback_matches_fla_reference(monkeypatch, fallback_case, batch_size): + """Cover missing fused projections/metadata and capture-safe staging fallback.""" + torch.manual_seed(2) + finalize_weights = fallback_case != "unfused-projections" + optimized, reference = _make_attention_pair(finalize_decode_weights=finalize_weights) + if fallback_case == "capture-before-staging": + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + hidden_states = ( + torch.randn( + batch_size, + 1, + HIDDEN_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + * 0.05 + ) + initial_cache = _make_cache(batch_size) + actual_output, actual_cache = _run_production_decode( + optimized, + hidden_states, + copy.deepcopy(initial_cache), + include_metadata=fallback_case != "missing-metadata", + ) expected_output, expected_cache = reference.forward_decode( hidden_states, copy.deepcopy(initial_cache) ) @@ -144,7 +264,6 @@ def test_optimized_decode_matches_fla_reference(batch_size: int) -> None: _assert_close(actual_cache.conv_state_q, expected_cache.conv_state_q) _assert_close(actual_cache.conv_state_k, expected_cache.conv_state_k) _assert_close(actual_cache.conv_state_v, expected_cache.conv_state_v) - assert optimized.decode_kernel_source() @torch.no_grad() @@ -171,8 +290,8 @@ def test_optimized_decode_updates_indexed_recurrent_state_pool_in_place( ) initial_cache = _make_cache(batch_size) - local_output, local_cache = optimized.forward_decode( - hidden_states, copy.deepcopy(initial_cache) + local_output, local_cache = _run_production_decode( + optimized, hidden_states, copy.deepcopy(initial_cache) ) slots = batch_size + 3 @@ -198,23 +317,39 @@ def test_optimized_decode_updates_indexed_recurrent_state_pool_in_place( state_pool.index_copy_(0, slot_indices.long(), initial_cache.recurrent_state) unselected_indices = torch.arange(3, device="cuda") unselected_before = state_pool.index_select(0, unselected_indices).clone() - indexed_cache = KimiKDACachedState( - conv_state_q=initial_cache.conv_state_q.clone(), - conv_state_k=initial_cache.conv_state_k.clone(), - conv_state_v=initial_cache.conv_state_v.clone(), - recurrent_state=state_pool, + conv_pool = torch.randn( + slots, + 3 * NUM_HEADS * HEAD_DIM, + CONV_KERNEL_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + conv_pool.index_copy_( + 0, + slot_indices.long(), + torch.cat( + [ + initial_cache.conv_state_q, + initial_cache.conv_state_k, + initial_cache.conv_state_v, + ], + dim=1, + ), ) - indexed_output, indexed_cache = optimized.forward_decode( + indexed_output, indexed_cache = _run_production_decode( + optimized, hidden_states, - indexed_cache, + initial_cache, + conv_pool=conv_pool, + state_pool=state_pool, + slot_indices=slot_indices.long(), ssm_state_indices=slot_indices, ) - assert indexed_cache.recurrent_state is state_pool torch.testing.assert_close(indexed_output, local_output, rtol=0, atol=0) torch.testing.assert_close( - state_pool.index_select(0, slot_indices.long()), + indexed_cache.recurrent_state, local_cache.recurrent_state, rtol=0, atol=0, diff --git a/tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py similarity index 98% rename from tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py rename to tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py index abd128c3edc3..b7c10bfbc16d 100644 --- a/tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py @@ -19,7 +19,7 @@ 1. A pure-torch fp32 CPU golden (vendored from the kernel drop's ``cpu_reference`` self-check). 2. An FLA sequential reference built from the exact op sequence - ``KimiKDARuntime._forward_verify`` uses in-tree: per-step fp32 causal + ``KimiKDALinearAttention.forward_verify`` uses in-tree: per-step fp32 causal conv + SiLU followed by ``fla.ops.kda.fused_recurrent_kda`` with ``use_qk_l2norm/use_gate/use_beta_sigmoid`` in kernel, ``lower_bound``, ``state_v_first=True``. @@ -319,7 +319,7 @@ def cute_run(data, zero_accepted_hint=False): def _fla_sequential_reference(data, num_accepted): """Per-request sequential conv+SiLU (fp32 torch) + fused_recurrent_kda. - Mirrors ``KimiKDARuntime._forward_verify``'s op sequence, extended with + Mirrors ``KimiKDALinearAttention.forward_verify``'s op sequence, extended with the replay prefix: for request ``n`` with ``a = num_accepted[n]``, the processed token sequence is ``a`` cached raw tokens (re-convolved from the extended conv-cache slots) followed by the ``1 + M`` new tokens. @@ -452,7 +452,7 @@ def test_round1_zero_accepted(B, H): for name in ("qkg_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): _assert_close(f"{name}(cute vs cpu)", cute_out[name], cpu[name], atol=2e-2) - # Kernel vs the FLA sequential path (the in-tree _forward_verify math). + # Kernel vs the FLA sequential path (the in-tree fused verify math). _assert_close( "out(cute vs fla)", cute_out["out"][0, new_rows].float(), diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py index bc092f195b58..7198fc858ab7 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py @@ -2,12 +2,21 @@ # SPDX-License-Identifier: Apache-2.0 """Parity tests for the optimized Kimi K3 KDA prefill op.""" +from types import SimpleNamespace + import pytest import torch pytest.importorskip("fla") -from tensorrt_llm._torch.modules.kimi_kda.kimi_kda_mixer import KimiKDALinearAttention # noqa: E402 +from tensorrt_llm._torch.modules.kimi_kda import ( + KimiKDALinearAttention, # noqa: E402 + _kda_kernels, # noqa: E402 +) +from tests.unittest._torch.modules.kimi_kda.kimi_kda_test_utils import ( # noqa: E402 + KimiKDAReference, + get_production_prefill_kernel_path, +) NUM_HEADS = 96 HEAD_DIM = 128 @@ -25,7 +34,9 @@ def _has_supported_gpu() -> bool: ) -def _make_attention_pair() -> tuple[KimiKDALinearAttention, KimiKDALinearAttention]: +def _make_attention_pair( + expected_prefill_kernel_path: str = "optimized", +) -> tuple[KimiKDALinearAttention, KimiKDAReference]: common = { "hidden_size": HIDDEN_SIZE, "num_heads": NUM_HEADS, @@ -36,16 +47,83 @@ def _make_attention_pair() -> tuple[KimiKDALinearAttention, KimiKDALinearAttenti "rms_norm_eps": 1e-5, "dtype": torch.bfloat16, } - optimized = KimiKDALinearAttention(**common).to("cuda") - reference = KimiKDALinearAttention(**common, use_optimized_prefill=False).to("cuda") + cfg = SimpleNamespace( + hidden_size=HIDDEN_SIZE, + rms_norm_eps=common["rms_norm_eps"], + linear_attn_config={ + "num_heads": NUM_HEADS, + "head_dim": HEAD_DIM, + "short_conv_kernel_size": CONV_KERNEL_SIZE, + "use_full_rank_gate": common["use_full_rank_gate"], + "gate_lower_bound": common["gate_lower_bound"], + }, + ) + optimized = KimiKDALinearAttention(cfg, layer_idx=0).to("cuda") + with torch.no_grad(): + optimized.dt_bias.zero_() + reference = KimiKDAReference(**common).to("cuda") reference.load_state_dict(optimized.state_dict()) + optimized.finalize_decode_weights() - assert optimized.prefill_kernel_path == "optimized" - assert reference.prefill_kernel_path == "fla" - assert reference.decode_kernel_path == optimized.decode_kernel_path + assert get_production_prefill_kernel_path(optimized) == expected_prefill_kernel_path return optimized, reference +def _run_production_prefill( + attention: KimiKDALinearAttention, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + *, + conv_pool: torch.Tensor | None = None, + state_pool: torch.Tensor | None = None, + slot_indices: torch.Tensor | None = None, + has_initial_states: torch.Tensor | None = None, +) -> torch.Tensor: + batch_size, sequence_length, _ = hidden_states.shape + if cu_seqlens is None: + cu_seqlens = torch.arange(batch_size + 1, device="cuda", dtype=torch.long) * sequence_length + else: + batch_size = cu_seqlens.numel() - 1 + + projection_size = NUM_HEADS * HEAD_DIM + if conv_pool is None: + conv_pool = torch.zeros( + batch_size, + 3 * projection_size, + CONV_KERNEL_SIZE, + dtype=torch.bfloat16, + device="cuda", + ) + if state_pool is None: + state_pool = torch.zeros( + batch_size, + NUM_HEADS, + HEAD_DIM, + HEAD_DIM, + dtype=torch.float32, + device="cuda", + ) + if slot_indices is None: + slot_indices = torch.arange(batch_size, device="cuda", dtype=torch.long) + use_initial_states = has_initial_states is not None + if has_initial_states is None: + has_initial_states = torch.zeros(batch_size, device="cuda", dtype=torch.bool) + metadata = SimpleNamespace( + use_initial_states=use_initial_states, + has_initial_states=has_initial_states, + ) + output = attention.forward_prefill( + hidden_states.reshape(-1, HIDDEN_SIZE), + cu_seqlens, + metadata, + batch_size, + conv_pool, + state_pool, + slot_indices, + ) + return output.reshape_as(hidden_states) + + def _assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None: actual_float = actual.float() expected_float = expected.float() @@ -75,12 +153,12 @@ def test_optimized_prefill_matches_fla_reference() -> None: ) * 0.05 ) - actual = optimized.forward_prefill(hidden_states) + actual = _run_production_prefill(optimized, hidden_states) expected = reference.forward_prefill(hidden_states) _assert_close(actual, expected) hidden_states = torch.randn(1, 300, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 - actual = optimized.forward_prefill(hidden_states) + actual = _run_production_prefill(optimized, hidden_states) expected = reference.forward_prefill(hidden_states) _assert_close(actual, expected) @@ -100,10 +178,9 @@ def test_optimized_prefill_matches_fla_reference() -> None: ) * 0.05 ) - actual = optimized.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + actual = _run_production_prefill(optimized, hidden_states, cumulative_lengths) expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) _assert_close(actual, expected) - assert optimized.prefill_kernel_source() @torch.no_grad() @@ -193,12 +270,10 @@ def test_kda_prefill_op_empty_token_batch_variants(): @torch.no_grad() def test_kda_mixer_empty_prefill(): - """Runtime-shaped regression: the mixer dispatch with an empty token - payload (a crashing call shape observed at runtime) must run - end-to-end on the optimized path.""" + """The production mixer handles an empty token payload without raising.""" optimized, _ = _make_attention_pair() hidden_states = torch.empty(1, 0, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") - out = optimized.forward_prefill(hidden_states) + out = _run_production_prefill(optimized, hidden_states) assert out.shape == (1, 0, HIDDEN_SIZE) @@ -237,42 +312,127 @@ def test_kda_prefill_op_partial_final_chunk_large_batch(): ) * 0.05 ) - actual = optimized.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + actual = _run_production_prefill(optimized, hidden_states, cumulative_lengths) expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) _assert_close(actual, expected) @torch.no_grad() -def test_kda_prefill_op_small_varlen_batch(): - """Small varlen batches (short-prompt contexts) through the dispatch. - - The persistent K123 scheduler needs >= 4 total chunks, so the dispatch - routes NT < 4 batches to the FLA path ([6,12], [1,2,3], [30] here); - the [6,12,20,25] case carries exactly NT=4 with total T < 64, running - the optimized masked path where building the eqlen chunk-offset - scratch used to raise ``step must be nonzero`` (arange step - ``T // 64 == 0``) — its output is parity-checked against FLA. +@pytest.mark.parametrize( + "sequence_lengths", + ([30], [65], [6, 12], [1, 2, 3], [64, 64, 63], [6, 12, 20, 25]), + ids=( + "one-chunk", + "two-chunk-single-seq", + "two-chunk-varlen", + "three-short-seqs", + "three-chunk-boundary", + "four-chunk-optimized-boundary", + ), +) +def test_kda_prefill_small_varlen_dispatch_matches_fla_reference(sequence_lengths): + """Cover every small-varlen fallback chunk count and the optimized boundary. + + The persistent K123 scheduler needs at least four total chunks. The + one-, two-, and three-chunk configurations must use FLA; the final case + has exactly four chunks and verifies the optimized boundary. """ optimized, reference = _make_attention_pair() - for sequence_lengths in ([6, 12], [1, 2, 3], [30], [6, 12, 20, 25]): - cumulative_lengths = torch.tensor( - [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], - dtype=torch.long, + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + hidden_states = ( + torch.randn( + 1, + sum(sequence_lengths), + HIDDEN_SIZE, + dtype=torch.bfloat16, device="cuda", ) - hidden_states = ( - torch.randn( - 1, - sum(sequence_lengths), - HIDDEN_SIZE, - dtype=torch.bfloat16, - device="cuda", - ) - * 0.05 + * 0.05 + ) + actual = _run_production_prefill(optimized, hidden_states, cumulative_lengths) + expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + _assert_close(actual, expected) + + +@torch.no_grad() +def test_kda_prefill_unavailable_kernel_fallback_preserves_mixed_initial_states(monkeypatch): + """The FLA core preserves prefix and continuation pool states.""" + torch.manual_seed(3) + optimized, _ = _make_attention_pair() + monkeypatch.setattr(_kda_kernels, "is_intree_prefill_available", lambda: False) + fallback, _ = _make_attention_pair(expected_prefill_kernel_path="fla") + fallback.load_state_dict(optimized.state_dict()) + + sequence_lengths = [64, 64, 64, 64] + cumulative_lengths = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + hidden_states = ( + torch.randn(1, sum(sequence_lengths), HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") + * 0.05 + ) + slots = 5 + slot_indices = torch.tensor([3, 0, 4, 1], dtype=torch.long, device="cuda") + has_initial_states = torch.tensor([True, False, True, False], device="cuda") + projection_size = NUM_HEADS * HEAD_DIM + conv_seed = ( + torch.randn( + slots, + 3 * projection_size, + CONV_KERNEL_SIZE, + dtype=torch.bfloat16, + device="cuda", ) - actual = optimized.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) - expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) - _assert_close(actual, expected) + * 0.05 + ) + state_seed = ( + torch.randn( + slots, + NUM_HEADS, + HEAD_DIM, + HEAD_DIM, + dtype=torch.float32, + device="cuda", + ) + * 0.05 + ) + + optimized_conv, optimized_state = conv_seed.clone(), state_seed.clone() + optimized_output = _run_production_prefill( + optimized, + hidden_states, + cumulative_lengths, + conv_pool=optimized_conv, + state_pool=optimized_state, + slot_indices=slot_indices, + has_initial_states=has_initial_states, + ) + fallback_conv, fallback_state = conv_seed.clone(), state_seed.clone() + fallback_output = _run_production_prefill( + fallback, + hidden_states, + cumulative_lengths, + conv_pool=fallback_conv, + state_pool=fallback_state, + slot_indices=slot_indices, + has_initial_states=has_initial_states, + ) + + _assert_close(fallback_output, optimized_output) + _assert_close( + fallback_conv.index_select(0, slot_indices), + optimized_conv.index_select(0, slot_indices), + ) + _assert_close( + fallback_state.index_select(0, slot_indices), + optimized_state.index_select(0, slot_indices), + ) @torch.no_grad() @@ -318,6 +478,6 @@ def test_kda_prefill_op_shape_growth_and_cu_dtype_transitions(): ) * 0.05 ) - actual = optimized.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) + actual = _run_production_prefill(optimized, hidden_states, cumulative_lengths) expected = reference.forward_prefill(hidden_states, cu_seqlens=cumulative_lengths) _assert_close(actual, expected) diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py similarity index 93% rename from tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py rename to tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py index 934bae38d809..f38aa46d1174 100644 --- a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py @@ -12,16 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Runtime-level parity: KimiKDARuntime fused verify vs sequential verify. +"""Runtime-level parity: KimiKDALinearAttention fused verify vs sequential verify. Simulates two chained speculative-verification rounds through -``KimiKDARuntime._forward_verify`` in both worlds: +``KimiKDALinearAttention.forward_verify`` in both worlds: * Sequential world: the legacy intermediate-buffer path - (``_forward_verify_sequential``) plus the manager's legacy promotion + (``forward_verify_sequential``) plus the manager's legacy promotion (copy the accepted step's conv window / SSM state into the live pools). * Fused world: the ``trtllm::kda_mtp_decode`` replay path - (``_forward_verify_fused``) with per-slot replay caches, in-place state + (``forward_verify_fused``) with per-slot replay caches, in-place state commit after the golden token, and only the accepted-draft count recorded between rounds. @@ -50,7 +50,7 @@ # The model module transitively imports the optional deps above, so it # must stay behind the guard too or collection fails instead of skipping. from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig - from tensorrt_llm._torch.models.modeling_kimi_linear import KimiKDARuntime + from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention except ImportError as e: _HAVE_DEPS = False _DEP_ERR = str(e) @@ -96,7 +96,7 @@ def _make_runtime(seed, aux_stream=None): gate_lower_bound=LB, ), ) - rt = KimiKDARuntime(cfg, layer_idx=0, aux_stream=aux_stream).to("cuda") + rt = KimiKDALinearAttention(cfg, layer_idx=0, aux_stream=aux_stream).to("cuda") gen = torch.Generator(device="cuda").manual_seed(seed) for name, p in rt.named_parameters(): if name.endswith("A_log"): @@ -209,11 +209,11 @@ def tokens(scale=0.5): ok = True # ---- Round 1 (no pending drafts) ---- x1 = tokens() - out1_seq = rt_seq._forward_verify_sequential( + out1_seq = rt_seq.forward_verify_sequential( x1, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices ) with with_multi_stream(True): - out1_fused = rt_fused._forward_verify( + out1_fused = rt_fused.forward_verify( x1, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices ) print("round 1:") @@ -226,11 +226,11 @@ def tokens(scale=0.5): # ---- Round 2 (fused path replays the accepted drafts) ---- x2 = tokens() - out2_seq = rt_seq._forward_verify_sequential( + out2_seq = rt_seq.forward_verify_sequential( x2, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices ) with with_multi_stream(True): - out2_fused = rt_fused._forward_verify( + out2_fused = rt_fused.forward_verify( x2, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices ) print("round 2 (mixed replay):") diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py similarity index 91% rename from tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py rename to tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py index 37ff7a8c9768..6756e663e63c 100644 --- a/tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py @@ -3,7 +3,7 @@ """KDA runtime projection and speculative-verification parity. The verify path must produce, for every step t, exactly the state and output -that t sequential single-token _forward_decode calls would produce — the two +that t sequential single-token forward_decode calls would produce — the two paths call the same FLA kernels with the same [B, 1] shapes, so agreement is expected to near-bitwise tolerance. A real mismatch here means the verify implementation (state threading, conv stepping, intermediate writes) is @@ -23,8 +23,11 @@ pytest.importorskip("fla") -from tensorrt_llm._torch.models.modeling_kimi_linear import KimiKDARuntime +from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream +from tests.unittest._torch.modules.kimi_kda.kimi_kda_test_utils import ( + get_production_decode_kernel_path, +) class _Cfg: @@ -115,14 +118,14 @@ def test_kda_fused_prefill_matches_separate_projections(): d = h * head_dim w = lin["short_conv_kernel_size"] - runtime = KimiKDARuntime(cfg, layer_idx=0).to(device) - if runtime.mixer.decode_kernel_path != "optimized": + runtime = KimiKDALinearAttention(cfg, layer_idx=0).to(device) + if get_production_decode_kernel_path(runtime) != "optimized": pytest.skip("needs an SM100/SM103 GPU") for param in runtime.parameters(): if param.is_floating_point(): torch.nn.init.normal_(param, std=0.02) - reference = KimiKDARuntime(cfg, layer_idx=0).to(device) + reference = KimiKDALinearAttention(cfg, layer_idx=0).to(device) reference.load_state_dict(runtime.state_dict()) runtime.finalize_decode_weights() assert runtime._qkvg_proj_weight is not None @@ -162,14 +165,14 @@ def test_kda_qkvg_multistream_decode_matches_separate_projections(): d = h * head_dim w = lin["short_conv_kernel_size"] - runtime = KimiKDARuntime(cfg, layer_idx=0, aux_stream=torch.cuda.Stream()).to(device) - if runtime.mixer.decode_kernel_path != "optimized": + runtime = KimiKDALinearAttention(cfg, layer_idx=0, aux_stream=torch.cuda.Stream()).to(device) + if get_production_decode_kernel_path(runtime) != "optimized": pytest.skip("needs an SM100/SM103 GPU") for param in runtime.parameters(): if param.is_floating_point(): torch.nn.init.normal_(param, std=0.02) - reference = KimiKDARuntime(cfg, layer_idx=0).to(device) + reference = KimiKDALinearAttention(cfg, layer_idx=0).to(device) reference.load_state_dict(runtime.state_dict()) runtime.finalize_decode_weights() assert runtime._qkvg_proj_weight is not None @@ -209,11 +212,11 @@ def test_kda_verify_matches_sequential_decode(batch, t_steps): dim = h * lin["head_dim"] w = lin["short_conv_kernel_size"] - runtime = KimiKDARuntime(cfg, layer_idx=0).to(device) + runtime = KimiKDALinearAttention(cfg, layer_idx=0).to(device) # dt_bias is torch.empty at construction and only filled by # load_weights(); with random weights it holds heap garbage, and a # NaN/Inf bit pattern poisons both paths identically (nvbug 6599150). - torch.nn.init.normal_(runtime.mixer.dt_bias, std=0.1) + torch.nn.init.normal_(runtime.dt_bias, std=0.1) slots = batch + 2 # non-trivial slot mapping cache = _LayerCache(slots, 3 * dim, w, h, lin["head_dim"], lin["head_dim"], t_steps, device) slot_indices = torch.arange(2, 2 + batch, device=device, dtype=torch.long) @@ -228,7 +231,7 @@ def test_kda_verify_matches_sequential_decode(batch, t_steps): ref_ssm = cache.temporal.clone() ref_outs, ref_conv_steps, ref_ssm_steps = [], [], [] for t in range(t_steps): - out = runtime._forward_decode(x[:, t], ref_conv, ref_ssm, slot_indices) + out = runtime.forward_decode(x[:, t], ref_conv, ref_ssm, slot_indices) ref_outs.append(out) ref_conv_steps.append(ref_conv.index_select(0, slot_indices).clone()) ref_ssm_steps.append(ref_ssm.index_select(0, slot_indices).clone()) @@ -236,7 +239,7 @@ def test_kda_verify_matches_sequential_decode(batch, t_steps): # --- Verify path: one call, intermediates into the scratch buffers. --- pristine_conv = cache.conv.clone() pristine_ssm = cache.temporal.clone() - out_verify = runtime._forward_verify( + out_verify = runtime.forward_verify( x.reshape(batch * t_steps, cfg.hidden_size), t_steps, cache, diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py index 9a7e053082af..e7db11ec0145 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py @@ -299,7 +299,7 @@ def __init__(self, *args, **kwargs): def forward(self, hidden_states, *args, **kwargs): return hidden_states - monkeypatch.setattr(modeling_kimi_linear, "KimiKDARuntime", _IdentityAttention) + monkeypatch.setattr(modeling_kimi_linear, "KimiKDALinearAttention", _IdentityAttention) monkeypatch.setattr(distributed, "AllReduce", _IdentityAllReduce) mapping = Mapping(