From e82c1acec2f284a6c567d875a7e828a50d7bbf2b Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Tue, 18 Aug 2026 17:11:17 +0800 Subject: [PATCH] [None][perf] eliminate KDA recurrent-state and metadata copies Signed-off-by: Jiagan Cheng --- .../_torch/attention_backend/interface.py | 15 +- .../custom_ops/cute_dsl_kimi_k3_custom_ops.py | 554 +++++++----------- .../blackwell/kimi_k3_kda/k4_persistent.py | 34 +- .../_torch/modules/kimi_kda/_kda_kernels.py | 124 +++- .../kimi_kda/kimi_k3_mamba_metadata.py | 93 +++ .../_torch/modules/kimi_kda/kimi_kda_mixer.py | 130 ++-- .../_torch/modules/mamba/gdn_mixer.py | 69 +-- .../_torch/modules/mamba/mamba2_metadata.py | 14 - .../modules/mamba/recurrent_state_cache.py | 87 +++ .../_torch/pyexecutor/model_engine.py | 11 +- .../test_lists/test-db/l0_b200.yml | 3 +- .../kimi_kda/test_kda_cache_soundness.py | 382 ------------ .../modules/kimi_kda/test_kda_prefill_op.py | 91 +-- .../kimi_kda/test_kda_prefill_state_parity.py | 397 ------------- .../kimi_kda/test_kimi_k3_mamba_metadata.py | 71 +++ .../test_kimi_kda_fused_verify_parity.py | 2 +- .../kimi_kda/test_kimi_kda_verify_parity.py | 42 +- .../mamba/test_gdn_kernel_optimizations.py | 8 +- .../mamba/test_recurrent_state_cache.py | 52 ++ 19 files changed, 803 insertions(+), 1376 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py create mode 100644 tensorrt_llm/_torch/modules/mamba/recurrent_state_cache.py delete mode 100644 tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py delete mode 100644 tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py create mode 100644 tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py create mode 100644 tests/unittest/_torch/modules/mamba/test_recurrent_state_cache.py diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 7a3977b910c6..adb9cdfabccc 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -43,6 +43,13 @@ ALLOWED_LAYER_TYPES as ALLOWED_ATTENTION_LAYER_TYPES +def _default_mamba_metadata_cls() -> Type[Any]: + """Import lazily to avoid an interface/Mamba metadata import cycle.""" + from ..modules.mamba.mamba2_metadata import Mamba2Metadata + + return Mamba2Metadata + + @dataclass class AttentionRuntimeFeatures: chunked_prefill: bool = False @@ -173,6 +180,8 @@ class AttentionMetadata: mamba_metadata: Optional[Any] = None mamba_chunk_size: int = 128 + mamba_metadata_cls: Type[Any] = field( + default_factory=_default_mamba_metadata_cls, init=False, repr=False) # The number of tokens in the padded sequence. padded_num_tokens: Optional[int] = None @@ -347,9 +356,9 @@ def _prepare_mamba_metadata(self): if self.mamba_metadata is None: if isinstance(self.kv_cache_manager, BaseMambaCacheManager): - from ..modules.mamba.mamba2_metadata import Mamba2Metadata - self.mamba_metadata = Mamba2Metadata(self.max_num_requests, - self.mamba_chunk_size) + metadata_cls = self.mamba_metadata_cls + self.mamba_metadata = metadata_cls(self.max_num_requests, + self.mamba_chunk_size) else: self.mamba_metadata = False return diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py index 5ff484f2d999..711b5e47e873 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py @@ -16,8 +16,7 @@ The implementation source-integrates the production orchestration from the KDA ``chunk_kda_fwd`` benchmark. The default pipeline launches fused K123, -Akk inverse, and persistent K4 kernels. Equal-length inputs may instead use -the fused K1234 kernel. +Akk inverse, and persistent K4 kernels. Supported: - Equal-length sequences (B ≥ 1, with the existing B=1 padding path) @@ -27,12 +26,12 @@ - use_gate_in_kernel=True with A_log - chunk_size=64 -The public ``trtllm::kda_prefill`` operator returns only the KDA output and -final recurrent state. Intermediate matrices remain private runner workspace. +The ``trtllm::kda_prefill`` operator updates selected recurrent-state pool rows +in place. Intermediate matrices remain private runner workspace. """ import weakref -from typing import Optional, Tuple +from typing import Optional import torch @@ -49,9 +48,6 @@ from ..cute_dsl_kernels.blackwell.kimi_k3_kda.fused_k123 import ( make_host_function as _fused_make_host, ) - from ..cute_dsl_kernels.blackwell.kimi_k3_kda.fused_k1234 import ( - make_host_fn as _fused_k1234_make_host, - ) from ..cute_dsl_kernels.blackwell.kimi_k3_kda.k4_persistent import ( BYTES_PER_TENSORMAP as _K4P_BTM, ) @@ -71,6 +67,51 @@ def _ct(t, etype): return r +def _fits_32bit_stride(tensor: torch.Tensor) -> bool: + int32_max = 2**31 - 1 + max_offset = int(tensor.storage_offset()) + for size, stride in zip(tensor.shape, tensor.stride()): + stride = abs(int(stride)) + if stride > int32_max: + return False + if size: + max_offset += (int(size) - 1) * stride + if max_offset > int32_max: + return False + return True + + +def _dynamic_dlpack_arg(tensor: torch.Tensor): + wrapper = from_dlpack( + tensor, + assumed_align=16, + use_32bit_stride=_fits_32bit_stride(tensor), + ) + for dim, stride in enumerate(tensor.stride()): + if stride == 1: + return wrapper.mark_layout_dynamic(dim) + return wrapper + + +def _layout_key(tensor: torch.Tensor): + return ( + tensor.dtype, + tuple(tensor.shape), + tuple(tensor.stride()), + _fits_32bit_stride(tensor), + ) + + +def _dynamic_vector_layout_key(tensor: torch.Tensor): + """Key a 1D CuTe argument whose sole extent is runtime dynamic.""" + assert tensor.ndim == 1 + return ( + tensor.dtype, + tuple(tensor.stride()), + _fits_32bit_stride(tensor), + ) + + def _current_cu_stream(device): """CUstream handle of torch's CURRENT stream, queried per call. @@ -259,17 +300,6 @@ def _get_varlen_k4_inputs(cu_seqlens, BT): _k4p_cache = {} _k4p_tm_ws = {} -# Side-stream cache for the initial-state copy overlapping K123. -_side_streams = {} - - -def _get_side_stream(dev): - idx = dev.index or 0 - if idx not in _side_streams: - _side_streams[idx] = torch.cuda.Stream(device=dev) - return _side_streams[idx] - - # Buffer cache: avoid re-allocating ~67us of intermediate tensors per call. # Also caches cute.Tensor wrappers (saves ~7us each call from from_dlpack). # LRU-bounded: entries are keyed by (B, T, ...) shapes, and the runtime @@ -375,11 +405,6 @@ def _with_slack(ctor, *shape, dtype): A_qk = _with_slack(torch.zeros, B, T_alloc, H, BT, dtype=dtype_k) A_kk = _with_slack(torch.zeros, B, T_alloc, H, BT, dtype=dtype_k) O_flat = _with_slack(torch.empty, B, T_alloc, H, V_dim, dtype=dtype_k) - # K4 reads initial state from S_out (caller copies it in) and writes - # the final state back into the same buffer — no separate s_4d, no - # extra D2D memcpy at the end of the K4 launcher. Layout matches - # caller's [N_seqs, H, K, V] contig fp32 convention. - S_out = torch.empty(N_seqs, H, K_dim, V_dim, device=dev, dtype=torch.float32) cu_eqlen = torch.arange(0, (B + 1) * T, T, dtype=torch.int32, device=dev) # co_eqlen is consumed only on the eqlen path (varlen K4 derives its # chunk offsets from cu_seqlens). Eqlen inputs are pre-padded to a @@ -425,22 +450,6 @@ def _wrap(t, etype): akk_out_view = from_dlpack(A_kk, assumed_align=16).mark_layout_dynamic() akk_out_view.element_type = fp32 - # K4 state wrapper points directly at S_out — K4 reads/writes in - # place. Only N_seqs (the outer compact mode) varies with the runtime - # context batch. Keep the inner H/K/V shape and strides static so the - # DSL can prove each (batch, head) state tile remains 16-byte aligned. - s_ct = from_dlpack(S_out, assumed_align=16).mark_compact_shape_dynamic( - mode=0, - stride_order=(0, 1, 2, 3), - ) - s_ct.element_type = fp32 - - # Side stream is process-long, safe to cache. The MAIN stream must - # NOT be cached here: the executor creates the buffers under one - # stream (warmup) and calls under another (execution_stream), so it - # is re-queried per call in _chunk_kda_fwd. - side_stream_cached = _get_side_stream(dev) - cute_wrappers = dict( a_ct=a_ct, aqc_ct=aqc_ct, @@ -453,13 +462,10 @@ def _wrap(t, etype): co_eqlen_ct=co_eqlen_ct, akk_in_view=akk_in_view, akk_out_view=akk_out_view, - s_ct=s_ct, - side_stream=side_stream_cached, # Filled lazily on first launch — saves cache_key tuple build + # outer dict lookup on subsequent calls. _k123_fns={}, _akk_inv_fn=None, - _k4_fn=None, ) while len(_buf_cache) >= _BUF_CACHE_MAX_ENTRIES: @@ -472,7 +478,6 @@ def _wrap(t, etype): A_qk, A_kk, O_flat, - S_out, cu_eqlen, co_eqlen, cute_wrappers, @@ -486,15 +491,14 @@ def _wrap(t, etype): def _launch_k4_persistent( cute_wrappers, v_beta, - S_in, - S_out, + state_pool, + state_indices, cu_seqlens, chunk_offsets, cu_eqlen_passed=False, num_sm=148, H=None, V_dim=None, - use_fast_sync=False, ): """Launch persistent K4 with cached CuTe wrappers. @@ -505,7 +509,7 @@ def _launch_k4_persistent( runtime never emits zero-length sequences. No fast-launch (args-tuple) cache here: such a cache pins the per-call - v/initial-state tensors via their cute wrappers (the wrapper holds the + activation and state-pool tensors via their CuTe wrappers (the wrapper holds the storage, so the keyed object never dies and weakref pruning never fires) — a per-call GPU memory leak in the executor runtime. The per-call cost is re-wrapping a handful of tensors (~10us each). @@ -515,7 +519,7 @@ def _launch_k4_persistent( dev = v_beta.device dev_idx = dev.index or 0 - s_ct = cute_wrappers["s_ct"] + s_ct = _dynamic_dlpack_arg(state_pool) a_ct = cute_wrappers["a_ct"] b_ct = cute_wrappers["ks_ct"] # raw k_scaled (beta absorbed in akk_inv) q_ct = cute_wrappers["qs_ct"] @@ -536,6 +540,8 @@ def _launch_k4_persistent( else: cu_ct, co_ct = _get_k4_varlen_cu_co(cu_seqlens, chunk_offsets) + state_indices_ct = _dynamic_dlpack_arg(state_indices) + # Allocate for the maximum persistent grid once. The scheduler applies # min(N_seqs * H, num_sm) at runtime, so pre-minimizing here would turn # every distinct request count into another compiled host function. @@ -555,41 +561,43 @@ def _launch_k4_persistent( # Launch on torch's CURRENT stream — see _current_cu_stream. stream = _current_cu_stream(dev) - k4_fn = cute_wrappers.get("_k4_fn") + # H and the state V/K dims MUST be in the key: s_ct bakes the inner + # [H, V, K] shape and all strides at compile time, so the compiled + # function is head-count- and layout-specific. N_seqs and the + # token/chunk counts remain runtime values. + cache_key = ( + dev_idx, + num_sm, + H, + state_pool.shape[-1], + V_dim, + _layout_key(state_pool), + # The sole dimension of this vector is runtime dynamic. Excluding + # its extent lets warmup serve full and underfilled context batches. + _dynamic_vector_layout_key(state_indices), + ) + k4_fn = _k4p_cache.get(cache_key) if k4_fn is None: - # H and the state K/V dims MUST be in the key: s_ct bakes the inner - # [H, K, V] shape and all strides at compile time - # (mark_compact_shape_dynamic(mode=0) in _get_buffers), so the - # compiled function is head-count-specific. Reusing an H=96-compiled - # kernel for a smaller H misaddresses every (seq, head) state tile - # (baked mode-0 stride H*K*V) — silently corrupting final_state and, - # through the recurrence, O. Observed as the small-H parity - # regression whenever another head count compiled first in the same - # process (isolated processes were clean). N_seqs and the token/chunk - # counts stay runtime — only layout-baked dims belong here. - cache_key = (dev_idx, num_sm, H, S_out.shape[-2], V_dim) - k4_fn = _k4p_cache.get(cache_key) - if k4_fn is None: - host_fn = _k4p_make_host(num_sm=num_sm) - k4_fn = cute.compile( - host_fn, - a_ct, - b_ct, - v_ct, - q_ct, - aqc_ct, - kg_ct, - o_ct, - gk_ct, - s_ct, - cu_ct, - co_ct, - tm_ct, - N_seqs, - stream, - ) - _k4p_cache[cache_key] = k4_fn - cute_wrappers["_k4_fn"] = k4_fn + host_fn = _k4p_make_host(num_sm=num_sm) + k4_fn = cute.compile( + host_fn, + a_ct, + b_ct, + v_ct, + q_ct, + aqc_ct, + kg_ct, + o_ct, + gk_ct, + s_ct, + state_indices_ct, + cu_ct, + co_ct, + tm_ct, + N_seqs, + stream, + ) + _k4p_cache[cache_key] = k4_fn args = ( a_ct, @@ -601,6 +609,7 @@ def _launch_k4_persistent( o_ct, gk_ct, s_ct, + state_indices_ct, cu_ct, co_ct, tm_ct, @@ -634,6 +643,7 @@ def _launch_fused_k123_inv( akk_out_view=None, cute_wrappers=None, varlen_pure_override=False, + varlen_is_aligned=None, ): """Persistent K1+K2 (writes A_kk in I+L format with diag=1) chained with BF16 akk_inv (in-place inversion). Final A_kk_inv = (I+L)^-1.""" @@ -649,15 +659,17 @@ def _launch_fused_k123_inv( T_padded = T if is_varlen else None has_bias = dt_bias is not None - # Auto-detect VARLEN_PURE eligibility, cached by id(cu_seqlens). - # First call with a given cu_seqlens object pays one GPU->CPU sync; later - # calls with the same tensor object are a dict lookup (~100 ns). + # Prefer the host-derived alignment property supplied by the model runtime. + # Direct low-level callers that omit it retain the cached device inference + # fallback for backward compatibility. varlen_pure = False if is_varlen and cu_seqlens is not None: if varlen_pure_override: # Caller (Phase 2.1 single-seq path) sentinel-padded the data for # this call; the cached per-object verdict stays untouched. varlen_pure = True + elif varlen_is_aligned is not None: + varlen_pure = varlen_is_aligned else: _vp_key = id(cu_seqlens) if _vp_key not in _varlen_pure_cache: @@ -813,95 +825,15 @@ def _launch_fused_k123_inv( akk_fn(*akk_args) -# ========== Fused K1234 compilation cache ========== -_fused_k1234_cache = {} -_BT = 64 - - -def _launch_fused_k1234( - q, - k, - v, - g, - A_log, - beta, - scale, - initial_state, - output_final_state, - dt_bias=None, - safe_gate=False, - lower_bound=None, -): - """Launch single fused K1+K2+K3+K4 SMEM-direct kernel (eqlen only).""" - B, T, H, K = q.shape - V_dim = v.shape[-1] - device = q.device - dev = device.index or 0 - NC = T // _BT - has_bias = dt_bias is not None - - # B only changes the outer extent and launch grid. T, H, and K change - # compiled tensor layouts, so they must select distinct artifacts. - cache_key = (dev, has_bias, safe_gate, T, H, K) - - # Buffers - BH = B * H - N_seqs = B - S_fp32 = torch.empty(BH, K, V_dim, dtype=torch.float32, device=device) - if initial_state is None: - S_fp32.zero_() - else: - S_fp32.copy_(initial_state.reshape(BH, K, V_dim)) - O_out = torch.empty(B, T, H, V_dim, dtype=q.dtype, device=device) - clocks = torch.empty(2, BH, dtype=torch.int64, device=device) - - bf16 = cutlass.BFloat16 - fp32 = cutlass.Float32 - if has_bias: - bias_ct = _ct(dt_bias.float().contiguous().view(H, K), fp32) - else: - bias_ct = _ct(torch.empty(1, 1, dtype=torch.float32, device=device), fp32) - lb_val = float(lower_bound) if lower_bound is not None else 0.0 - - ct_args = ( - _ct(q.contiguous(), bf16), - _ct(k.contiguous(), bf16), - _ct(g.contiguous(), bf16), - _ct(A_log.float().contiguous(), fp32), - _ct(beta.contiguous(), bf16), - float(scale), - _ct(v.contiguous(), bf16), - _ct(O_out, bf16), - _ct(S_fp32, fp32), - bias_ct, - lb_val, - _ct(clocks, cutlass.Int64), - NC, - H, - B, - # Launch on torch's CURRENT stream — see _current_cu_stream. - _current_cu_stream(device), - ) - - if cache_key not in _fused_k1234_cache: - host_fn = _fused_k1234_make_host(has_bias=has_bias, use_safe_gate=safe_gate) - _fused_k1234_cache[cache_key] = cute.compile(host_fn, *ct_args) - _fused_k1234_cache[cache_key](*ct_args) - - S_out = S_fp32.reshape(N_seqs, H, K, V_dim) - final_state = S_out if output_final_state else None - return O_out, final_state - - def _chunk_kda_fwd( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, scale: float, - initial_state: torch.Tensor, - output_final_state: bool, cu_seqlens: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, chunk_size: int = 64, @@ -910,55 +842,17 @@ def _chunk_kda_fwd( use_gate_in_kernel: bool = False, A_log: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - cp_context=None, - use_fused_k1234: bool = False, -): - """KDA forward — optimized, FLA-compatible interface.""" + varlen_is_aligned: bool | None = None, + single_sequence_length: int | None = None, +) -> torch.Tensor: + """Run optimized KDA prefill against indexed recurrent-state rows.""" if safe_gate and lower_bound is None: lower_bound = -5.0 is_varlen = cu_seqlens is not None if q.shape[1] == 0: - # Zero-token call: the runtime can emit a context batch whose token - # payload is empty (observed under the overlap scheduler + logprobs - # flows; the FLA fallback tolerates it). No tokens means no output - # rows and an unchanged recurrent state — return before any kernel - # or buffer setup (whose arange(step=T) would raise "step must be - # nonzero"). - B, _, H, K = q.shape - V_dim = v.shape[-1] - n_seqs = B if cu_seqlens is None else cu_seqlens.shape[0] - 1 - o = v.new_empty(B, 0, H, V_dim) - if not output_final_state: - final_state = None - elif initial_state is not None: - # clone: the caller index_copy_s final_state back into the state - # pool the initial state may alias. - final_state = initial_state.to(torch.float32).clone() - else: - final_state = torch.zeros(n_seqs, H, K, V_dim, dtype=torch.float32, device=q.device) - return (o, final_state, None, None, None, None, None, None, None, None, None, initial_state) - - # ===== Fused K1234 path (eqlen only, single kernel launch) ===== - if use_fused_k1234 and not is_varlen: - o, final_state = _launch_fused_k1234( - q, - k, - v, - g, - A_log, - beta, - scale, - initial_state, - output_final_state, - dt_bias=dt_bias, - safe_gate=safe_gate, - lower_bound=lower_bound, - ) - return (o, final_state, None, None, None, None, None, None, None, None, None, initial_state) + raise ValueError("Indexed KDA prefill does not support an empty token batch.") B, T, H, K = q.shape V_dim = v.shape[-1] @@ -1021,17 +915,27 @@ def _chunk_kda_fwd( # Multi-seq optimization needs per-seq dynamic tensormap (Phase 2.2). needs_varlen_single_pad = False if is_varlen and cu_seqlens is not None and cu_seqlens.shape[0] == 2 and B == 1: - _vl_key = id(cu_seqlens) - if _vl_key not in _varlen_pure_cache: - cu_cpu = cu_seqlens.cpu().tolist() - sl = cu_cpu[1] - cu_cpu[0] - _varlen_pure_cache[_vl_key] = sl % BT == 0 - _varlen_single_seqlen_cache[_vl_key] = sl - _prune_on_gc(_varlen_pure_cache, _vl_key, cu_seqlens) - _prune_on_gc(_varlen_single_seqlen_cache, _vl_key, cu_seqlens) - if not _varlen_pure_cache[_vl_key]: - real_T = _varlen_single_seqlen_cache[_vl_key] - needs_varlen_single_pad = True + if single_sequence_length is not None: + single_sequence_is_aligned = single_sequence_length % BT == 0 + if varlen_is_aligned is not None and varlen_is_aligned != single_sequence_is_aligned: + raise ValueError("varlen_is_aligned disagrees with single_sequence_length") + varlen_is_aligned = single_sequence_is_aligned + if not single_sequence_is_aligned: + real_T = single_sequence_length + needs_varlen_single_pad = True + else: + _vl_key = id(cu_seqlens) + if _vl_key not in _varlen_pure_cache: + cu_cpu = cu_seqlens.cpu().tolist() + sl = cu_cpu[1] - cu_cpu[0] + _varlen_pure_cache[_vl_key] = sl % BT == 0 + _varlen_single_seqlen_cache[_vl_key] = sl + _prune_on_gc(_varlen_pure_cache, _vl_key, cu_seqlens) + _prune_on_gc(_varlen_single_seqlen_cache, _vl_key, cu_seqlens) + varlen_is_aligned = _varlen_pure_cache[_vl_key] + if not varlen_is_aligned: + real_T = _varlen_single_seqlen_cache[_vl_key] + needs_varlen_single_pad = True varlen_pure_override = False if needs_varlen_single_pad: # q/k/v/beta already zero-padded by caller (FLA convention). Re-build g @@ -1091,32 +995,11 @@ def _chunk_kda_fwd( A_qk, A_kk, O_flat, - S_out, cu_eqlen, co_eqlen, cute_wrappers, ) = _get_buffers(device, k.dtype, B, T, H, K, V_dim, NT, N_seqs, BT, varlen=is_varlen) - # ===== State copy on side stream, parallel with K123 ===== - # K4 needs S_out populated with initial_state. By doing this copy on a - # side stream BEFORE K123 launches, the D2D memcpy overlaps with K123's - # compute. Especially big win for high-N_seqs varlen where state is huge - # (192MB for N=32, ~64us memcpy) — would otherwise serialize before K4. - if initial_state is None: - S_in = torch.zeros(N_seqs, H, K, V_dim, dtype=torch.float32, device=device) - else: - S_in = initial_state - # Current stream per call — never the buffer-creation-time stream (the - # executor creates buffers under warmup's stream and calls under - # execution_stream; syncing against a stale stream un-orders the copy). - main_stream = torch.cuda.current_stream(device) - side_stream = cute_wrappers["side_stream"] - needs_copy = S_in.data_ptr() != S_out.data_ptr() - if needs_copy: - side_stream.wait_stream(main_stream) - with torch.cuda.stream(side_stream): - S_out.copy_(S_in) - # Beta is fused entirely in akk_inv kernel epilogue (post-inv column-scale). # No host v*beta and no K1 k_scaled*beta any more. _launch_fused_k123_inv( @@ -1143,6 +1026,7 @@ def _chunk_kda_fwd( akk_out_view=cute_wrappers["akk_out_view"], cute_wrappers=cute_wrappers, varlen_pure_override=varlen_pure_override, + varlen_is_aligned=varlen_is_aligned, ) # ===== K4: persistent kernel (eqlen + varlen via cu_seqlens) ===== @@ -1153,49 +1037,24 @@ def _chunk_kda_fwd( cu_for_k4 = cu_eqlen chunk_offsets_for_k4 = co_eqlen - # Wait for the side-stream state copy to finish before K4 reads S_out. - if needs_copy: - main_stream.wait_stream(side_stream) - _launch_k4_persistent( cute_wrappers, v, - S_in, - S_out, + state_pool, + state_indices, cu_for_k4, chunk_offsets_for_k4, cu_eqlen_passed=(not is_varlen), H=H, V_dim=V_dim, - use_fast_sync=(not is_varlen), ) o = O_flat - A_qk_out = A_qk - A_kk_out = A_kk if needs_eqlen_pad: # Caller called with original T = real_T; their downstream code expects # outputs at that shape. Slice the padded scratch tail back off. o = o[:, :real_T] - A_qk_out = A_qk_out[:, :real_T] - A_kk_out = A_kk_out[:, :real_T] - # multiseq_info is always None here (Phase 2.2 disabled — see above) - final_state = S_out if output_final_state else None - - return ( - o, - final_state, - None, - A_qk_out, - A_kk_out, - None, - None, - None, - None, - None, - None, - initial_state, - ) + return o class KdaPrefillRunner: @@ -1208,9 +1067,9 @@ def forward( v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, scale: float, - initial_state: Optional[torch.Tensor], - output_final_state: bool, cu_seqlens: Optional[torch.Tensor] = None, chunk_indices: Optional[torch.Tensor] = None, chunk_size: int = 64, @@ -1219,31 +1078,26 @@ def forward( use_gate_in_kernel: bool = False, A_log: Optional[torch.Tensor] = None, dt_bias: Optional[torch.Tensor] = None, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - use_fused_k1234: bool = False, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """Run KDA prefill and return output plus final recurrent state.""" + varlen_is_aligned: Optional[bool] = None, + single_sequence_length: Optional[int] = None, + ) -> torch.Tensor: + """Run KDA prefill and update selected recurrent-state rows.""" if not IS_CUTLASS_DSL_AVAILABLE: raise RuntimeError("Kimi K3 KDA prefill requires NVIDIA CUTLASS DSL") if chunk_size != 64: raise ValueError(f"Kimi K3 KDA prefill requires chunk_size=64, got {chunk_size}") - # Zero-token calls take the early return in _chunk_kda_fwd and - # never touch A_log; the runtime emits such batches (overlap - # scheduler + logprobs flows), so only require A_log when there - # is work to do. - if A_log is None and q.shape[1] != 0: + if A_log is None: raise ValueError("Kimi K3 KDA prefill requires A_log") - result = _chunk_kda_fwd( + return _chunk_kda_fwd( q=q, k=k, v=v, g=g, beta=beta, + state_pool=state_pool, + state_indices=state_indices, scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, chunk_size=chunk_size, @@ -1252,25 +1106,74 @@ def forward( use_gate_in_kernel=use_gate_in_kernel, A_log=A_log, dt_bias=dt_bias, - disable_recompute=disable_recompute, - return_intermediate_states=return_intermediate_states, - use_fused_k1234=use_fused_k1234, + varlen_is_aligned=varlen_is_aligned, + single_sequence_length=single_sequence_length, ) - return result[0], result[1] + + +def _validate_indexed_state_pool( + q: torch.Tensor, + v: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], +) -> None: + num_sequences = q.shape[0] if cu_seqlens is None else cu_seqlens.shape[0] - 1 + H, K, V = q.shape[2], q.shape[3], v.shape[3] + if state_pool.dtype != torch.float32 or state_pool.ndim != 4: + raise ValueError( + "Expected state_pool to be a rank-4 fp32 tensor with shape [slots, H, V, K]." + ) + if state_pool.shape[1:] != (H, V, K): + raise ValueError( + f"Expected state_pool shape [slots, {H}, {V}, {K}], got {tuple(state_pool.shape)}." + ) + expected_inner_strides = (V * K, K, 1) + if state_pool.stride()[1:] != expected_inner_strides: + raise ValueError( + "Indexed KDA prefill requires dense H/V/K inner dimensions; " + f"expected strides {expected_inner_strides}, got " + f"{state_pool.stride()[1:]}." + ) + if state_pool.stride(0) < H * V * K or state_pool.stride(0) % 4 != 0: + raise ValueError( + "Indexed KDA prefill requires a non-overlapping, 16-byte-aligned slot stride." + ) + if state_pool.data_ptr() % 16 != 0: + raise ValueError("Indexed KDA prefill requires a 16-byte-aligned state pool.") + if state_indices.ndim != 1 or state_indices.shape[0] != num_sequences: + raise ValueError( + f"Expected state_indices shape [{num_sequences}], got {tuple(state_indices.shape)}." + ) + if state_indices.dtype not in (torch.int32, torch.int64): + raise ValueError("state_indices must have dtype int32 or int64.") + if not state_indices.is_contiguous(): + raise ValueError("state_indices must be contiguous.") + for name, tensor in ( + ("q", q), + ("v", v), + ("state_indices", state_indices), + ): + if tensor.device != state_pool.device: + raise ValueError(f"{name} and state_pool must be on the same device.") if IS_CUTLASS_DSL_AVAILABLE: - @torch.library.custom_op("trtllm::kda_prefill", mutates_args=(), device_types="cuda") + @torch.library.custom_op( + "trtllm::kda_prefill", + mutates_args=("state_pool",), + device_types="cuda", + ) def kda_prefill( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, scale: float, - initial_state: Optional[torch.Tensor], - output_final_state: bool, cu_seqlens: Optional[torch.Tensor] = None, chunk_indices: Optional[torch.Tensor] = None, chunk_size: int = 64, @@ -1279,41 +1182,37 @@ def kda_prefill( use_gate_in_kernel: bool = False, A_log: Optional[torch.Tensor] = None, dt_bias: Optional[torch.Tensor] = None, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - use_fused_k1234: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Run Kimi K3 KDA chunked prefill on Blackwell GPUs. - - Returns ``(output, final_state)``. ``final_state`` is an empty - tensor when ``output_final_state`` is False — - ``torch.library.custom_op`` cannot infer a schema for - ``Optional[Tensor]`` returns. - """ - output, final_state = KdaPrefillRunner.forward( + varlen_is_aligned: Optional[bool] = None, + single_sequence_length: Optional[int] = None, + ) -> torch.Tensor: + """Run KDA prefill directly against a V-first recurrent-state pool.""" + _validate_indexed_state_pool( + q, + v, + state_pool, + state_indices, + cu_seqlens, + ) + return KdaPrefillRunner.forward( q=q, k=k, v=v, g=g, beta=beta, + state_pool=state_pool, + state_indices=state_indices, A_log=A_log, scale=scale, dt_bias=dt_bias, - initial_state=initial_state, - output_final_state=output_final_state, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, chunk_size=chunk_size, safe_gate=safe_gate, lower_bound=lower_bound, use_gate_in_kernel=use_gate_in_kernel, - disable_recompute=disable_recompute, - return_intermediate_states=return_intermediate_states, - use_fused_k1234=use_fused_k1234, + varlen_is_aligned=varlen_is_aligned, + single_sequence_length=single_sequence_length, ) - if final_state is None: - final_state = q.new_empty(0, dtype=torch.float32) - return output, final_state @kda_prefill.register_fake def _( @@ -1322,9 +1221,9 @@ def _( v: torch.Tensor, g: torch.Tensor, beta: torch.Tensor, + state_pool: torch.Tensor, + state_indices: torch.Tensor, scale: float, - initial_state: Optional[torch.Tensor], - output_final_state: bool, cu_seqlens: Optional[torch.Tensor] = None, chunk_indices: Optional[torch.Tensor] = None, chunk_size: int = 64, @@ -1333,20 +1232,11 @@ def _( use_gate_in_kernel: bool = False, A_log: Optional[torch.Tensor] = None, dt_bias: Optional[torch.Tensor] = None, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - use_fused_k1234: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor]: - del k, g, beta, A_log, scale, dt_bias, initial_state - del chunk_indices, chunk_size, safe_gate, lower_bound - del use_gate_in_kernel, disable_recompute, return_intermediate_states - del use_fused_k1234 - num_sequences = q.shape[0] if cu_seqlens is None else cu_seqlens.shape[0] - 1 - output = v.new_empty(v.shape) - if output_final_state: - final_state = q.new_empty( - (num_sequences, q.shape[2], q.shape[3], v.shape[3]), dtype=torch.float32 - ) - else: - final_state = q.new_empty(0, dtype=torch.float32) - return output, final_state + varlen_is_aligned: Optional[bool] = None, + single_sequence_length: Optional[int] = None, + ) -> torch.Tensor: + del q, k, g, beta, state_pool, state_indices + del scale, cu_seqlens, chunk_indices, chunk_size, safe_gate + del lower_bound, use_gate_in_kernel, A_log, dt_bias + del varlen_is_aligned, single_sequence_length + return v.new_empty(v.shape) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py index 7c118a3051b6..019c09efe0ff 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/k4_persistent.py @@ -517,6 +517,7 @@ def k4_persistent_kernel( nv_b_mn_sl: cute.ComposedLayout, mGkLastExp: cute.Tensor, mS_fp32: cute.Tensor, + state_indices: cute.Tensor, cu_seqlens: cute.Tensor, chunk_offsets: cute.Tensor, mA: cute.Tensor, @@ -685,13 +686,13 @@ def _make_umma_pipe(): tRT_tCtState = thr_state_r2t.partition_D(tCtState_mn) atom_state_g2r = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=128 + cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=32 ) tiled_state_g2r = cute.make_tiled_copy_S(atom_state_g2r, tiled_state_r2t) thr_state_g2r = tiled_state_g2r.get_slice(warpgroup_tidx) atom_state_r2g = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=128 + cute.nvgpu.CopyUniversalOp(), acc_dtype, num_bits_per_copy=32 ) tiled_state_r2g = cute.make_tiled_copy_D(atom_state_r2g, tiled_state_t2r) thr_state_r2g = tiled_state_r2g.get_slice(warpgroup_tidx) @@ -717,6 +718,7 @@ def _make_umma_pipe(): while work.is_valid_tile: batch_idx, head_idx, _ = work.tile_idx + state_idx = state_indices[batch_idx] batch_start = cu_seqlens[batch_idx] batch_end = cu_seqlens[batch_idx + 1] num_chunks = cute.ceil_div(batch_end - batch_start, M) @@ -724,9 +726,12 @@ def _make_umma_pipe(): # batch_start // M only works when all seq lengths are multiples of M. chunk_base = chunk_offsets[batch_idx] - gS_init = cute.flat_divide(mS_fp32[batch_idx, head_idx, None, None], (M6, N6))[ - None, None, 0, 0 - ] + state_vk = mS_fp32[state_idx, head_idx, None, None] + state_kv = cute.make_tensor( + state_vk.iterator, + cute.make_layout((M6, N6), stride=(1, M6)), + ) + gS_init = cute.flat_divide(state_kv, (M6, N6))[None, None, 0, 0] tGR_tCgState_in = thr_state_g2r.partition_S(gS_init) tGR_tCrState_in = thr_state_g2r.retile(tRrState) cute.copy(tiled_state_g2r, tGR_tCgState_in, tGR_tCrState_in) @@ -774,18 +779,23 @@ def _make_umma_pipe(): global_chunk = global_chunk + 1 cute.arch.mbarrier_wait(mma6_done_mbar, phase=(global_chunk - 1) % 2) - gS_out = cute.flat_divide(mS_fp32[batch_idx, head_idx, None, None], (M6, N6))[ - None, None, 0, 0 - ] - tGR_tCgState_out = thr_state_r2g.partition_D(gS_out) - tGR_tCrState_out = thr_state_r2g.retile(tRrState) num_state_subs_final = tRrState.shape[2] for sub in cutlass.range(num_state_subs_final): cute.copy(tiled_state_t2r, tTR_tCtState[None, 0, sub], tRrState[None, 0, sub]) cute.arch.fence_view_async_tmem_load() + state_vk = mS_fp32[state_idx, head_idx, None, None] + state_kv = cute.make_tensor( + state_vk.iterator, + cute.make_layout((M6, N6), stride=(1, M6)), + ) + gS_out = cute.flat_divide(state_kv, (M6, N6))[None, None, 0, 0] + tGR_tCgState_out = thr_state_r2g.partition_D(gS_out) + tGR_tCrState_out = thr_state_r2g.retile(tRrState) for sub in cutlass.range(num_state_subs_final): cute.copy( - tiled_state_r2g, tGR_tCrState_out[None, 0, sub], tGR_tCgState_out[None, 0, sub] + tiled_state_r2g, + tGR_tCrState_out[None, 0, sub], + tGR_tCgState_out[None, 0, sub], ) cute.arch.mbarrier_arrive(final_state_done_mbar) @@ -1316,6 +1326,7 @@ def host_fn( o_raw: cute.Tensor, gk_last_exp: cute.Tensor, s_fp32: cute.Tensor, + state_indices: cute.Tensor, cu_seqlens: cute.Tensor, chunk_offsets: cute.Tensor, tm_workspace: cute.Tensor, @@ -1480,6 +1491,7 @@ def host_fn( sl_nv_b_mn, gk_last_exp, s_fp32, + state_indices, cu_seqlens, chunk_offsets, a, diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py index f5bbf0390d88..9ea8ded6bee8 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -226,15 +226,69 @@ def can_use_optimized_prefill( self, *, cu_seqlens: Optional[torch.Tensor], + num_tokens: int, + chunk_indices: Optional[torch.Tensor] = None, chunk_size: int = 64, ) -> bool: """Return whether this batch can stay on the Blackwell prefill path.""" - if self.prefill_kernel_path != "optimized": + if self.prefill_kernel_path != "optimized" or num_tokens == 0: return False if cu_seqlens is not None: - from fla.ops.utils.index import prepare_chunk_indices + if chunk_indices is None: + from fla.ops.utils.index import prepare_chunk_indices - return prepare_chunk_indices(cu_seqlens, chunk_size).shape[0] >= 4 + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + return chunk_indices.shape[0] >= 4 + return True + + def can_use_indexed_prefill( + self, + *, + state_pool: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + num_tokens: int, + chunk_indices: Optional[torch.Tensor] = None, + chunk_size: int = 64, + ) -> bool: + """Return whether prefill can update this V-first pool directly.""" + if not self.can_use_optimized_prefill( + cu_seqlens=cu_seqlens, + num_tokens=num_tokens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + ): + return False + if state_pool.dtype != torch.float32 or state_pool.ndim != 4: + return False + _, H, V, K = state_pool.shape + if state_pool.stride()[1:] != (V * K, K, 1): + return False + if state_pool.stride(0) < H * V * K or state_pool.stride(0) % 4: + return False + if state_pool.data_ptr() % 16: + return False + num_sequences = state_indices.shape[0] if cu_seqlens is None else cu_seqlens.shape[0] - 1 + if ( + state_indices.ndim != 1 + or state_indices.shape[0] != num_sequences + or state_indices.dtype not in (torch.int32, torch.int64) + or not state_indices.is_contiguous() + ): + return False + if ( + has_initial_states.ndim != 1 + or has_initial_states.shape[0] != num_sequences + or has_initial_states.dtype != torch.bool + or not has_initial_states.is_contiguous() + ): + return False + if ( + state_pool.device != state_indices.device + or state_pool.device != has_initial_states.device + ): + return False return True def prefill_chunk_kda( @@ -252,7 +306,13 @@ def prefill_chunk_kda( safe_gate: bool, lower_bound: Optional[float], cu_seqlens: Optional[torch.Tensor], + chunk_indices: Optional[torch.Tensor] = None, chunk_size: int = 64, + state_pool: Optional[torch.Tensor] = None, + state_indices: Optional[torch.Tensor] = None, + varlen_is_aligned: Optional[bool] = None, + single_sequence_length: Optional[int] = None, + force_fallback: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Run KDA chunked prefill. @@ -265,20 +325,27 @@ def prefill_chunk_kda( On the FLA path it calls ``fla.ops.kda.chunk_kda`` directly with the matching flags so the semantics are byte-equivalent. - State layout contract (both paths): ``initial_state`` is consumed - and ``final_state`` returned in the V-first ``[N, H, V, K]`` layout — - the layout of the executor's ssm pool and of the fused decode - kernel. The in-tree prefill op natively uses the FLA-default K-first - ``[N, H, K, V]`` layout, so the optimized path transposes at both - boundaries (K == V == 128 for Kimi K3, so shapes alone cannot catch - a mix-up — the transpose is semantic). + The optimized kernel always reads and writes indexed V-first + ``[slots, H, V, K]`` state-pool rows. Callers provide the executor + pool and reset fresh rows before entering the kernel. The FLA fallback + retains the batch-dense ``initial_state``/``final_state`` contract. + + ``chunk_indices``, ``varlen_is_aligned``, and + ``single_sequence_length`` are prepared once from Kimi K3 runtime + metadata. They keep the optimized path from reading ``cu_seqlens`` + back from the GPU in every KDA layer. """ - use_optimized = self.prefill_kernel_path == "optimized" - chunk_indices = None + use_indexed_state = state_pool is not None + if use_indexed_state and (initial_state is not None or state_indices is None): + raise ValueError( + "Indexed KDA prefill requires state_indices and does not accept initial_state." + ) + use_optimized = self.prefill_kernel_path == "optimized" and not force_fallback if use_optimized and cu_seqlens is not None: - from fla.ops.utils.index import prepare_chunk_indices + if chunk_indices is None: + from fla.ops.utils.index import prepare_chunk_indices - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) # The persistent K123 scheduler needs at least 4 total chunks # (cgs_per_head = NT // 4 cooperative groups per head). The # eqlen path guarantees this by padding to a 256-token multiple @@ -291,6 +358,14 @@ def prefill_chunk_kda( if chunk_indices.shape[0] < 4: use_optimized = False + if use_indexed_state and not use_optimized: + raise RuntimeError( + "Indexed KDA prefill requires the optimized prefill path; " + "use the FLA state path for this batch." + ) + if use_optimized and not use_indexed_state: + raise RuntimeError("Optimized KDA prefill requires an indexed state pool.") + if use_optimized: import torch.nn.functional as F from fla.modules.l2norm import l2norm_fwd @@ -320,20 +395,15 @@ def prefill_chunk_kda( dt_bias_kernel = dt_bias.detach() if dt_bias is not None else None _load_prefill_module() # registers trtllm::kda_prefill - - if initial_state is not None: - # Pool V-first [N, H, V, K] -> op K-first [N, H, K, V]. - initial_state = initial_state.transpose(-1, -2).contiguous() - - out, final_state = torch.ops.trtllm.kda_prefill( + out = torch.ops.trtllm.kda_prefill( q=q, k=k, v=v, g=g, beta=beta, + state_pool=state_pool, + state_indices=state_indices, scale=scale, - initial_state=initial_state, - output_final_state=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, chunk_size=chunk_size, @@ -342,16 +412,12 @@ def prefill_chunk_kda( use_gate_in_kernel=True, A_log=A_log_kernel, dt_bias=dt_bias_kernel, + varlen_is_aligned=varlen_is_aligned, + single_sequence_length=single_sequence_length, ) if out.shape[1] != real_T: out = out[:, :real_T] - if final_state is not None and final_state.numel() > 0: - # Op K-first [N, H, K, V] -> pool V-first [N, H, V, K]. - # .contiguous() also detaches the result from the op's - # shared per-shape S_out scratch, which the next same-shape - # call overwrites. - final_state = final_state.transpose(-1, -2).contiguous() - return out, final_state + return out, None from fla.ops.kda import chunk_kda diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py new file mode 100644 index 000000000000..d95a99c7c0a1 --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_k3_mamba_metadata.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3-specific recurrent-state metadata.""" + +from dataclasses import dataclass, field + +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata + +KDA_PREFILL_CHUNK_SIZE = 64 + + +@torch.compile(dynamic=True, fullgraph=True) +def _prepare_kda_chunk_indices( + sequence_lengths: torch.Tensor, + chunk_indices: torch.Tensor, +) -> torch.Tensor: + """Fill FLA-compatible chunk indices without reading CUDA counts on CPU.""" + chunk_counts = (sequence_lengths + KDA_PREFILL_CHUNK_SIZE - 1).div( + KDA_PREFILL_CHUNK_SIZE, rounding_mode="floor" + ) + num_chunks = chunk_indices.shape[0] + # Supplying output_size keeps repeat_interleave from synchronizing to read + # sum(chunk_counts) from the device. + sequence_indices = torch.repeat_interleave( + torch.arange(sequence_lengths.shape[0], dtype=torch.long, device=sequence_lengths.device), + chunk_counts, + output_size=num_chunks, + ) + chunk_starts = torch.cumsum(chunk_counts, dim=0) - chunk_counts + local_chunk_indices = ( + torch.arange(num_chunks, dtype=torch.long, device=sequence_lengths.device) + - chunk_starts[sequence_indices] + ) + chunk_indices[:, 0].copy_(sequence_indices) + chunk_indices[:, 1].copy_(local_chunk_indices) + return chunk_indices + + +class KimiK3MambaMetadata(Mamba2Metadata): + """Mamba metadata extended with Kimi K3 KDA preparation.""" + + def __init__(self, max_batch_size: int, chunk_size: int) -> None: + super().__init__(max_batch_size, chunk_size) + self._kda_chunk_indices = torch.empty((0, 2), dtype=torch.long, device="cuda") + self.kda_chunk_indices: torch.Tensor | None = None + self.kda_varlen_is_aligned: bool | None = None + self.kda_single_sequence_length: int | None = None + + def prepare(self, attn_metadata: AttentionMetadata) -> None: + super().prepare(attn_metadata) + + context_lengths = attn_metadata.seq_lens[: attn_metadata.num_contexts].tolist() + self.kda_varlen_is_aligned = ( + all(length % KDA_PREFILL_CHUNK_SIZE == 0 for length in context_lengths) + if context_lengths + else None + ) + self.kda_single_sequence_length = context_lengths[0] if len(context_lengths) == 1 else None + + if not context_lengths: + self.kda_chunk_indices = None + return + + chunk_counts = [ + (length + KDA_PREFILL_CHUNK_SIZE - 1) // KDA_PREFILL_CHUNK_SIZE + for length in context_lengths + ] + num_chunks = sum(chunk_counts) + if self._kda_chunk_indices.shape[0] < num_chunks: + self._kda_chunk_indices = torch.empty((num_chunks, 2), dtype=torch.long, device="cuda") + self.kda_chunk_indices = self._kda_chunk_indices[:num_chunks] + if num_chunks == 0: + return + + _prepare_kda_chunk_indices( + attn_metadata.seq_lens_cuda[: len(context_lengths)], + self.kda_chunk_indices, + ) + + +@dataclass(kw_only=True) +class KimiK3AttentionMetadata(TrtllmAttentionMetadata): + """TRT-LLM attention metadata with Kimi K3 recurrent-state metadata.""" + + mamba_metadata_cls: type[Mamba2Metadata] = field( + default=KimiK3MambaMetadata, + init=False, + repr=False, + ) 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 36681da67b9f..fd2558a7f171 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -27,7 +27,8 @@ from ...attention_backend import AttentionMetadata from ...distributed import AllReduce, AllReduceStrategy -from ...modules.multi_stream_utils import maybe_execute_in_parallel +from ..mamba.recurrent_state_cache import reset_recurrent_state_rows +from ..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" @@ -310,12 +311,7 @@ def forward( 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() + state_indices = mamba_metadata.state_indices[:batch_size] cu_seqlens = mamba_metadata.query_start_loc_long[: num_prefills + 1] num_decodes = batch_size - num_prefills @@ -419,7 +415,14 @@ def forward_prefill( slot_indices, layer_cache=None, ) -> torch.Tensor: - if not self._dispatch.can_use_optimized_prefill(cu_seqlens=cu_seqlens): + chunk_indices = getattr(mamba_metadata, "kda_chunk_indices", None) + varlen_is_aligned = getattr(mamba_metadata, "kda_varlen_is_aligned", None) + single_sequence_length = getattr(mamba_metadata, "kda_single_sequence_length", None) + if not self._dispatch.can_use_optimized_prefill( + cu_seqlens=cu_seqlens, + num_tokens=x2d.shape[0], + chunk_indices=chunk_indices, + ): return self.forward_prefill_fallback( x2d, cu_seqlens, @@ -460,15 +463,36 @@ def forward_prefill( # 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. + has_init = mamba_metadata.has_initial_states[:num_prefills] + if not self._dispatch.can_use_indexed_prefill( + state_pool=ssm_pool, + state_indices=slot_indices, + has_initial_states=has_init, + cu_seqlens=cu_seqlens, + num_tokens=x2d.shape[0], + chunk_indices=chunk_indices, + ): + return self.forward_prefill_fallback( + x2d, + cu_seqlens, + mamba_metadata, + num_prefills, + conv_pool, + ssm_pool, + slot_indices, + layer_cache, + ) + reset_recurrent_state_rows( + ssm_pool, + slot_indices, + has_init, + conv_pool if mamba_metadata.use_initial_states else None, + ) + slot_indices_long = slot_indices.long() 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 + cs = conv_pool.index_select(0, slot_indices_long) 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 @@ -494,11 +518,11 @@ def forward_prefill( k = rearrange(k, "... (h d) -> ... h d", d=self.head_k_dim) v = rearrange(v, "... (h d) -> ... h d", d=self.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. + # The optimized kernel reads and writes the V-first pool directly. + # Unsupported layouts and batches are routed to the FLA fallback + # before reaching this point. lower_bound = self.gate_lower_bound - o, final_state = self._dispatch.prefill_chunk_kda( + o, _ = self._dispatch.prefill_chunk_kda( q=q, k=k, v=v, @@ -507,22 +531,27 @@ def forward_prefill( A_log=self.A_log, dt_bias=self.dt_bias, scale=self.head_k_dim**-0.5, - initial_state=recurrent_in, + initial_state=None, safe_gate=lower_bound is not None, lower_bound=lower_bound, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_pool=ssm_pool, + state_indices=slot_indices, + varlen_is_aligned=varlen_is_aligned, + single_sequence_length=single_sequence_length, ) # 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) + 0, + slot_indices_long, + 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) + self._sync_kda_replay_conv_window(layer_cache, slot_indices_long, conv_q, conv_k, conv_v) return self._output_gate_and_proj(x, o, onorm_g) @@ -541,6 +570,7 @@ def forward_prefill_fallback( from einops import rearrange d = self.proj_size + slot_indices_long = slot_indices.long() x = x2d.unsqueeze(0) q_proj = self.q_proj(x) k_proj = self.k_proj(x) @@ -549,10 +579,10 @@ def forward_prefill_fallback( conv_q_in = conv_k_in = conv_v_in = recurrent_in = None if mamba_metadata.use_initial_states: has_init = mamba_metadata.has_initial_states[:num_prefills] - conv_state = conv_pool.index_select(0, slot_indices) + conv_state = conv_pool.index_select(0, slot_indices_long) conv_state[~has_init] = 0 conv_q_in, conv_k_in, conv_v_in = _kda_split_conv_sections(conv_state, d) - recurrent_in = ssm_pool.index_select(0, slot_indices) + recurrent_in = ssm_pool.index_select(0, slot_indices_long) recurrent_in[~has_init] = 0 q, conv_q = self.q_conv1d( @@ -587,16 +617,17 @@ def forward_prefill_fallback( safe_gate=self.gate_lower_bound is not None, lower_bound=self.gate_lower_bound, cu_seqlens=cu_seqlens, + force_fallback=True, ) conv_pool.index_copy_( 0, - slot_indices, + slot_indices_long, 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)) - self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_q, conv_k, conv_v) + ssm_pool.index_copy_(0, slot_indices_long, final_state.to(ssm_pool.dtype)) + self._sync_kda_replay_conv_window(layer_cache, slot_indices_long, conv_q, conv_k, conv_v) return self._output_gate_and_proj(x, o) def forward_decode( @@ -708,16 +739,19 @@ def _project_bfa_and_fb() -> tuple[torch.Tensor, torch.Tensor]: ) x_qkv = qkvg[:, : 3 * d] onorm_g = qkvg[:, 3 * d : 4 * d] + slot_indices_long = slot_indices.long() # 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 = conv_pool.index_select(0, slot_indices_long) # [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) + ssm_pool + if ssm_state_indices is not None + else ssm_pool.index_select(0, slot_indices_long) ) o = self._dispatch.decode_kda( @@ -751,18 +785,22 @@ def _project_bfa_and_fb() -> tuple[torch.Tensor, torch.Tensor]: update_conv_cache=False, ) if ssm_state_indices is None: - ssm_pool.index_copy_(0, slot_indices, state) + ssm_pool.index_copy_(0, slot_indices_long, 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) + conv_pool.index_copy_(0, slot_indices_long, 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 :] + layer_cache, + slot_indices_long, + new_win[:, :d], + new_win[:, d : 2 * d], + new_win[:, 2 * d :], ) return self.o_proj(o.view(B, d)) @@ -774,11 +812,14 @@ def forward_decode_fallback( from einops import rearrange d = self.proj_size + slot_indices_long = slot_indices.long() x = x2d.unsqueeze(1) # [B, 1, hidden] - cs = conv_pool.index_select(0, slot_indices) + cs = conv_pool.index_select(0, slot_indices_long) 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) + ssm_pool + if ssm_state_indices is not None + else ssm_pool.index_select(0, slot_indices_long) ) q_proj = self.q_proj(x) @@ -873,11 +914,11 @@ def _kernel_input(value: torch.Tensor) -> torch.Tensor: conv_pool.index_copy_( 0, - slot_indices, + slot_indices_long, 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)) + ssm_pool.index_copy_(0, slot_indices_long, 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 @@ -887,7 +928,7 @@ def _kernel_input(value: torch.Tensor) -> torch.Tensor: # so drafted batches always take the verify path. self._sync_kda_replay_conv_window( layer_cache, - slot_indices, + slot_indices_long, new_conv_q, new_conv_k, new_conv_v, @@ -1026,9 +1067,9 @@ def forward_verify_fused( 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 + pending = layer_cache.prev_num_accepted_tokens[ + slot_indices + ] # 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 ) @@ -1054,7 +1095,7 @@ def forward_verify_fused( 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), + ssm_state_indices=slot_indices, cu_seqlens=cu_seqlens, num_spec=num_spec, num_accepted_tokens=pending, @@ -1123,9 +1164,10 @@ def forward_verify_sequential( # Gathered copies — mutated across steps, never written back to the # live pools. - cs = conv_pool.index_select(0, slot_indices) + slot_indices_long = slot_indices.long() + cs = conv_pool.index_select(0, slot_indices_long) conv_q, conv_k, conv_v = _kda_split_conv_sections(cs, d) - state = ssm_pool.index_select(0, slot_indices) + state = ssm_pool.index_select(0, slot_indices_long) step_outputs: List[torch.Tensor] = [] for t in range(num_steps): diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index 331bfe442976..5a0ff6dce73f 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -46,6 +46,7 @@ from .layernorm_gated import RMSNorm as RMSNormGated from .layernorm_gated import rms_norm_gated_token_major from .mamba2_metadata import Mamba2Metadata +from .recurrent_state_cache import reset_recurrent_state_rows # FlashInfer GDN prefill is ON by default; set TLLM_USE_FLASHINFER_GDN_PREFILL=0 @@ -90,70 +91,6 @@ def _extract_gdn_extra_attrs(layer_idx: str): return metadata, gdn_layer, extra_attrs.get("spec_metadata", None) -@triton.jit -def _reset_gdn_states_kernel( - ssm_states, - conv_states, - state_indices, - has_initial_states, - ssm_state_stride, - conv_state_stride, - NUM_CACHE_LINES: tl.constexpr, - SSM_STATE_SIZE: tl.constexpr, - CONV_STATE_SIZE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - request_idx = tl.program_id(0) - offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - state_idx = tl.load(state_indices + request_idx).to(tl.int64) - needs_reset = ~tl.load(has_initial_states + request_idx).to(tl.int1) - valid_state = (state_idx >= 0) & (state_idx < NUM_CACHE_LINES) - ssm_row_offset = state_idx * ssm_state_stride.to(tl.int64) - conv_row_offset = state_idx * conv_state_stride.to(tl.int64) - - tl.store( - ssm_states + ssm_row_offset + offsets, - 0.0, - mask=needs_reset & valid_state & (offsets < SSM_STATE_SIZE), - ) - tl.store( - conv_states + conv_row_offset + offsets, - 0.0, - mask=needs_reset & valid_state & (offsets < CONV_STATE_SIZE), - ) - - -def _reset_gdn_states( - ssm_states: torch.Tensor, - conv_states: torch.Tensor, - state_indices: torch.Tensor, - has_initial_states: torch.Tensor, -) -> None: - num_requests = state_indices.shape[0] - if num_requests == 0: - return - - ssm_state_size = ssm_states.numel() // ssm_states.shape[0] - conv_state_size = conv_states.numel() // conv_states.shape[0] - block_size = 256 - grid = ( - num_requests, - triton.cdiv(max(ssm_state_size, conv_state_size), block_size), - ) - _reset_gdn_states_kernel[grid]( - ssm_states, - conv_states, - state_indices, - has_initial_states, - ssm_states.stride(0), - conv_states.stride(0), - ssm_states.shape[0], - ssm_state_size, - conv_state_size, - block_size, - ) - - @torch.library.custom_op("trtllm::gdn_custom_op_inplace", mutates_args=("output",)) def gdn_custom_op_inplace( mixed_qkv: torch.Tensor, @@ -1107,11 +1044,11 @@ def forward_core( if num_prefills > 0: # PyExecutor guarantees prefill requests are placed before decode requests has_initial_states_p = has_initial_states[:num_prefills] - _reset_gdn_states( + reset_recurrent_state_rows( ssm_states, - conv_states, state_indices_p, has_initial_states_p, + conv_states, ) is_target_verify = ( diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 96ef5223b7bb..faa73250b774 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -261,14 +261,6 @@ def __init__(self, max_batch_size: int, chunk_size: int): self.state_indices = torch.zeros(max_batch_size, dtype=torch.int32, device="cuda") - # int64 mirror of state_indices, refreshed once per prepare() so - # per-layer consumers that need long indices (index_select / - # index_copy_) do not each launch an int32->int64 cast kernel - # inside the decode CUDA graph (69 KDA layers x ~1.7us for Kimi K3). - self._state_indices_long = torch.zeros(max_batch_size, - dtype=torch.long, - device="cuda") - self.state_indices_long = self._state_indices_long[:0] # Stable data_ptr() of the CUDA tensor we alias (if any) — used to # detect cache-manager buffer reallocation that would silently break # CUDA graph replays. @@ -417,12 +409,6 @@ def prepare(self, attn_metadata: AttentionMetadata): self.state_indices[:batch_size].copy_( self.state_indices_cpu[:batch_size], non_blocking=True) - # Refresh the int64 mirror once per step (outside the decode graph) - # so layers can index pools without a per-layer cast kernel. - self._state_indices_long[:batch_size].copy_( - self.state_indices[:batch_size]) - self.state_indices_long = self._state_indices_long[:batch_size] - self._prepare_replay_work_items(kv_cache_manager, batch_size, num_contexts) diff --git a/tensorrt_llm/_torch/modules/mamba/recurrent_state_cache.py b/tensorrt_llm/_torch/modules/mamba/recurrent_state_cache.py new file mode 100644 index 000000000000..aada6121dc40 --- /dev/null +++ b/tensorrt_llm/_torch/modules/mamba/recurrent_state_cache.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared indexed state-cache operations for recurrent attention modules.""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _reset_recurrent_state_rows_kernel( + recurrent_states, + conv_states, + state_indices, + has_initial_states, + recurrent_state_stride, + conv_state_stride, + NUM_CACHE_LINES: tl.constexpr, + RECURRENT_STATE_SIZE: tl.constexpr, + CONV_STATE_SIZE: tl.constexpr, + RESET_CONV: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + request_idx = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + state_idx = tl.load(state_indices + request_idx).to(tl.int64) + needs_reset = ~tl.load(has_initial_states + request_idx).to(tl.int1) + valid_state = (state_idx >= 0) & (state_idx < NUM_CACHE_LINES) + recurrent_row_offset = state_idx * recurrent_state_stride.to(tl.int64) + tl.store( + recurrent_states + recurrent_row_offset + offsets, + 0.0, + mask=needs_reset & valid_state & (offsets < RECURRENT_STATE_SIZE), + ) + if RESET_CONV: + conv_row_offset = state_idx * conv_state_stride.to(tl.int64) + tl.store( + conv_states + conv_row_offset + offsets, + 0.0, + mask=needs_reset & valid_state & (offsets < CONV_STATE_SIZE), + ) + + +def reset_recurrent_state_rows( + recurrent_states: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, + conv_states: Optional[torch.Tensor] = None, +) -> None: + """Zero fresh-request rows in recurrent and optional convolution pools. + + Invalid or negative slot indices are ignored, matching the cache-manager + reset contract used by GDN. The launch is ordered on PyTorch's current + stream, so a following indexed recurrence observes the zeroed rows. + """ + num_requests = state_indices.shape[0] + if num_requests == 0: + return + if recurrent_states.shape[0] == 0: + raise ValueError("recurrent state pool must have at least one cache line") + if conv_states is not None and conv_states.shape[0] != recurrent_states.shape[0]: + raise ValueError("recurrent and convolution pools must have the same slot count") + if has_initial_states.shape != state_indices.shape: + raise ValueError("has_initial_states and state_indices must have matching shapes") + + recurrent_state_size = recurrent_states.numel() // recurrent_states.shape[0] + conv_state_size = conv_states.numel() // conv_states.shape[0] if conv_states is not None else 0 + block_size = 256 + grid = ( + num_requests, + triton.cdiv(max(recurrent_state_size, conv_state_size), block_size), + ) + _reset_recurrent_state_rows_kernel[grid]( + recurrent_states, + conv_states if conv_states is not None else recurrent_states, + state_indices, + has_initial_states, + recurrent_states.stride(0), + conv_states.stride(0) if conv_states is not None else 0, + recurrent_states.shape[0], + recurrent_state_size, + conv_state_size, + conv_states is not None, + block_size, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b4de2dba6935..439db06b145c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -84,7 +84,7 @@ set_per_request_prefill_cuda_graph_flag, set_torch_compiling, with_model_extra_attrs) from .breakable_cuda_graph_runner import BreakableCUDAGraphRunner -from .config_utils import is_mla +from .config_utils import is_kimi_linear, is_mla from .cuda_graph_runner import (ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM, CUDAGraphRunner, CUDAGraphRunnerConfig, EncoderCUDAGraphRunner, @@ -3217,6 +3217,15 @@ def _set_up_attn_metadata( assert self.attn_metadata.kv_cache_manager is kv_cache_manager return self.attn_metadata + if is_kimi_linear(config): + if metadata_cls is not TrtllmAttentionMetadata: + raise ValueError( + "Kimi K3 requires the TRTLLM attention backend") + from ..modules.kimi_kda.kimi_k3_mamba_metadata import \ + KimiK3AttentionMetadata + + metadata_cls = KimiK3AttentionMetadata + self.attn_metadata = metadata_cls( max_num_requests=self.batch_size, max_num_tokens=self.max_num_tokens, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a7b49d1c70df..0f165e6fbaeb 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -97,9 +97,8 @@ l0_b200: # ------------- Kimi K3 (KimiLinear) unit tests --------------- - 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/modules/kimi_kda/test_kda_cache_soundness.py + - unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py - unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py - - unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py # GPU KDA disagg transfer + peer-validation (cpu_only cases skipped by the # stage's "not cpu_only" markexpr and run on the CPU-Generic stage instead). diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py deleted file mode 100644 index fdf5b7a8b961..000000000000 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py +++ /dev/null @@ -1,382 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Runtime-integration soundness tests for the optimized KDA prefill op. - -Covers the two classes of bug that op-math parity tests cannot see: - -1. STREAM soundness. The executor runs the model on a dedicated non-blocking - ``torch.cuda.Stream`` (``py_executor.execution_stream``); the CuTe DSL - kernels must launch on torch's current stream, not the DSL default stream. - A default-stream launch races with the projections producing q/k/v/g/beta - and with consumers of the output — silent, intermittent corruption e2e - (GSM8K 68.99 vs 97.01 on the FLA control) while single-stream unit tests - all pass. ``test_nondefault_stream_parity`` reproduces this - deterministically by delaying the execution stream with ``_sleep`` so any - kernel launched on the default stream reads not-yet-written inputs. - -2. id()-keyed cache soundness under the runtime's fresh-tensor-per-call - pattern. CPython recycles an object's address (= its id) as soon as it is - freed; the ``_prune_on_gc`` weakref guard must pop the cache entry at - dealloc, before a new tensor can alias the key. The recycled-id tests - engineer exactly that aliasing: cache a verdict under id(t1), free t1, - re-allocate until a new tensor lands on the same id, refill it with data - for which the stale verdict would be WRONG, and parity-check against FLA. - ``test_single_seq_same_object_repeated`` additionally covers the - Phase 2.1 poisoning case (same cu_seqlens object, two calls). -""" - -import gc - -import pytest -import torch - -pytest.importorskip("fla") - -from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import KDAKernelDispatch # noqa: E402 - -NUM_HEADS = 96 -HEAD_K_DIM = 128 -LOWER_BOUND = -5.0 - - -def _has_supported_gpu() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {(10, 0), (10, 3)} - - -pytestmark = pytest.mark.skipif( - not _has_supported_gpu(), - reason="Kimi K3 is supported only on Blackwell (SM100/SM103)", -) - - -def _op_module(): - from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops - - return cute_dsl_kimi_k3_custom_ops - - -@pytest.fixture(scope="module") -def dispatch_pair(): - optimized = KDAKernelDispatch(use_optimized_prefill=True, use_optimized_decode=False) - assert optimized.prefill_kernel_path == "optimized" - reference = KDAKernelDispatch(use_optimized_prefill=False, use_optimized_decode=False) - assert reference.prefill_kernel_path == "fla" - return optimized, reference - - -@pytest.fixture(scope="module") -def gate_params(): - torch.manual_seed(0) - a_log = torch.randn(NUM_HEADS, dtype=torch.float32, device="cuda") * 0.5 - dt_bias = torch.randn(NUM_HEADS * HEAD_K_DIM, dtype=torch.float32, device="cuda") * 0.1 - return a_log, dt_bias - - -def _make_inputs(total_t: int, seed: int): - gen = torch.Generator(device="cuda").manual_seed(seed) - h, k = NUM_HEADS, HEAD_K_DIM - - def rnd(*shape, dtype=torch.bfloat16): - return torch.randn(*shape, generator=gen, dtype=torch.float32, device="cuda").to(dtype) - - q = rnd(1, total_t, h, k) - key = rnd(1, total_t, h, k) - v = rnd(1, total_t, h, k) - g = rnd(1, total_t, h, k) - beta = rnd(1, total_t, h, dtype=torch.float32) - return q, key, v, g, beta - - -def _run(dispatch, gate_params, q, k, v, g, beta, cu): - """No .clone() on cu — these tests exercise object identity on purpose.""" - a_log, dt_bias = gate_params - return dispatch.prefill_chunk_kda( - q=q.clone(), - k=k.clone(), - v=v.clone(), - g=g.clone(), - beta=beta.clone(), - A_log=a_log, - dt_bias=dt_bias, - scale=HEAD_K_DIM**-0.5, - initial_state=None, - safe_gate=True, - lower_bound=LOWER_BOUND, - cu_seqlens=cu, - ) - - -def _assert_close(name, actual, expected): - actual, expected = actual.float(), expected.float() - cos = torch.nn.functional.cosine_similarity(actual.flatten(), expected.flatten(), dim=0).item() - rel = ((actual - expected).norm() / (expected.norm() + 1e-12)).item() - assert cos > 0.999 and rel < 3e-2, f"{name}: cos={cos:.6f} rel_l2={rel:.3e}" - - -def _make_cu(lens): - return torch.tensor( - [0] + torch.cumsum(torch.tensor(lens), 0).tolist(), dtype=torch.long, device="cuda" - ) - - -def _flush_tensor_cache_pins(): - """Evict external pins on cu_seqlens tensors. - - Both the dispatcher (fla.ops.utils.index) and the op - (tensorrt_llm._torch.modules.fla.index) run cu_seqlens through - ``@tensor_cache`` helpers that keep the 4 most recent (args, result) - tuples alive. Until those entries are evicted, a cu_seqlens object - stays pinned after the call — which is SOUND (the id cannot be - recycled while our id-keyed entry exists) but makes the recycled-id - scenario unreachable. Churn the caches with fresh dummy tensors so the - pins drop and the finalizer-prune path can be exercised. - """ - from fla.ops.utils.index import prepare_chunk_indices as fla_pci - from fla.ops.utils.index import prepare_chunk_offsets as fla_pco - - from tensorrt_llm._torch.modules.fla.index import prepare_chunk_indices as intree_pci - - for _ in range(5): - dummy = torch.tensor([0, 64], dtype=torch.long, device="cuda") - fla_pci(dummy, 64) - fla_pco(dummy, 64) - intree_pci(dummy, 64) - - -def _release_cu(mod, cu_holder): - """Free the cu_seqlens tensor held in the one-element list ``cu_holder`` - and return (old_id, pruned). ``pruned`` False means something still pins - the tensor — sound, but the recycled-id scenario is unreachable; the - caller must then prove the id cannot be recycled and skip. - """ - old_id = id(cu_holder[0]) - cu_holder.clear() - _flush_tensor_cache_pins() - gc.collect() - return old_id, old_id not in mod._varlen_pure_cache - - -def _alloc_with_recycled_id(target_id, make, attempts=512): - """Allocate via ``make()`` until an object lands on ``target_id``. - - Wrong-address candidates are kept alive so the allocator cannot hand the - same wrong slot back; pymalloc free-lists normally return the freed - address on the first attempt. Returns None if unattainable. - """ - hold = [] - for _ in range(attempts): - cand = make() - if id(cand) == target_id: - return cand - hold.append(cand) - return None - - -# --------------------------------------------------------------------------- -# 1. Stream soundness — deterministic repro of the e2e accuracy regression. -# --------------------------------------------------------------------------- - - -def _make_eqlen_inputs(batch, t, seed): - gen = torch.Generator(device="cuda").manual_seed(seed) - h, k = NUM_HEADS, HEAD_K_DIM - - def rnd(*shape, dtype=torch.bfloat16): - return torch.randn(*shape, generator=gen, dtype=torch.float32, device="cuda").to(dtype) - - return ( - rnd(batch, t, h, k), - rnd(batch, t, h, k), - rnd(batch, t, h, k), - rnd(batch, t, h, k), - rnd(batch, t, h, dtype=torch.float32), - ) - - -@pytest.mark.parametrize("regime", ["eqlen", "varlen"]) -@torch.no_grad() -def test_nondefault_stream_parity(dispatch_pair, gate_params, regime): - """Run the optimized prefill on a fresh non-default stream (like the - executor's ``execution_stream``) whose queue is held back by a GPU sleep, - with the inputs produced BEHIND that sleep, and read the outputs on that - same stream right after the call. A kernel the op launches on the default - stream instead of the current stream reads inputs before they exist - (eqlen regime — the op has no host sync there, so the repro is - deterministic) and/or its output is consumed before it is written - (varlen regime — the ``out * 1.0`` capture below races a stray - default-stream K4).""" - optimized, reference = dispatch_pair - if regime == "eqlen": - lens = None - make = lambda seed: _make_eqlen_inputs(2, 256, seed) # noqa: E731 - else: - lens = [100, 257, 300] # multi-seq masked path, like eval traffic - make = lambda seed: _make_inputs(sum(lens), seed) # noqa: E731 - - q, k, v, g, beta = make(31337) - cu = _make_cu(lens) if lens else None - out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, cu) - # Warmup compile of the optimized variants on the default stream so the - # streamed run below launches immediately (a ~100 s JIT inside the - # streamed region would let the sleep expire and hide the race). - _run(optimized, gate_params, q, k, v, g, beta, _make_cu(lens) if lens else None) - torch.cuda.synchronize() - - exec_stream = torch.cuda.Stream() # non-blocking, like the executor's - with torch.cuda.stream(exec_stream): - torch.cuda._sleep(1 << 28) # hold the stream back ~100 ms - # Produced on exec_stream, pending behind the sleep: a default-stream - # launch would read these buffers before they are written. - q2, k2, v2 = q * 1.0, k * 1.0, v * 1.0 - g2, beta2 = g * 1.0, beta * 1.0 - out_opt, state_opt = _run( - optimized, gate_params, q2, k2, v2, g2, beta2, _make_cu(lens) if lens else None - ) - # Consume on the execution stream immediately, exactly like the - # runtime's output-gate matmul and ssm_pool.index_copy_ do. - out_opt = out_opt * 1.0 - state_opt = state_opt * 1.0 - torch.cuda.synchronize() - - _assert_close(f"stream_{regime}/out", out_opt, out_ref) - _assert_close(f"stream_{regime}/state", state_opt, state_ref) - - -# --------------------------------------------------------------------------- -# 2. id()-keyed cache soundness under recycled ids / reused objects. -# --------------------------------------------------------------------------- - - -@torch.no_grad() -def test_recycled_id_varlen_pure_cache(dispatch_pair, gate_params): - """Aligned lens cache VARLEN_PURE=True under id(cu). Free cu, land a new - cu_seqlens on the recycled id with NON-aligned lens: a stale hit would - run the mask-free compile variant on partial chunks.""" - optimized, reference = dispatch_pair - mod = _op_module() - - aligned = [128, 256, 192] - cu_holder = [_make_cu(aligned)] - q, k, v, g, beta = _make_inputs(sum(aligned), seed=11) - _run(optimized, gate_params, q, k, v, g, beta, cu_holder[0]) - torch.cuda.synchronize() - assert mod._varlen_pure_cache.get(id(cu_holder[0])) is True, ( - "test setup: expected VARLEN_PURE=True cached under id(cu1)" - ) - - old_id, pruned = _release_cu(mod, cu_holder) - nonaligned = [100, 257, 219] # same n_seqs so the entry shape matches - if not pruned: - # Entry outlived our release: something still pins the tensor. That - # is sound ONLY if no new tensor can land on its id. - clash = _alloc_with_recycled_id(old_id, lambda: _make_cu(nonaligned), attempts=64) - assert clash is None, ( - "UNSOUND: id recycled while a stale _varlen_pure_cache entry for it is still present" - ) - pytest.skip("cu_seqlens still pinned externally; recycled-id scenario unreachable this run") - - cu2 = _alloc_with_recycled_id(old_id, lambda: _make_cu(nonaligned)) - if cu2 is None: - pytest.skip("could not obtain a recycled id for cu_seqlens") - q, k, v, g, beta = _make_inputs(sum(nonaligned), seed=12) - out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, cu2) - out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, _make_cu(nonaligned)) - _assert_close("recycled_pure/out", out_opt, out_ref) - _assert_close("recycled_pure/state", state_opt, state_ref) - assert mod._varlen_pure_cache.get(id(cu2)) is False - - -@torch.no_grad() -def test_recycled_id_single_seq_caches(dispatch_pair, gate_params): - """Phase 2.1 caches: a single ALIGNED seq caches (pure=True, seqlen=256) - under id(cu). Recycle the id with a single NON-aligned seq: a stale hit - would skip the g sentinel pad / use the wrong single-seq padding.""" - optimized, reference = dispatch_pair - mod = _op_module() - - cu_holder = [_make_cu([256])] - q, k, v, g, beta = _make_inputs(256, seed=21) - _run(optimized, gate_params, q, k, v, g, beta, cu_holder[0]) - torch.cuda.synchronize() - assert mod._varlen_pure_cache.get(id(cu_holder[0])) is True - - old_id, pruned = _release_cu(mod, cu_holder) - if not pruned: - clash = _alloc_with_recycled_id(old_id, lambda: _make_cu([300]), attempts=64) - assert clash is None, ( - "UNSOUND: id recycled while stale Phase 2.1 cache entries for it are still present" - ) - pytest.skip("cu_seqlens still pinned externally; recycled-id scenario unreachable this run") - assert old_id not in mod._varlen_single_seqlen_cache - - cu2 = _alloc_with_recycled_id(old_id, lambda: _make_cu([300])) - if cu2 is None: - pytest.skip("could not obtain a recycled id for cu_seqlens") - q, k, v, g, beta = _make_inputs(300, seed=22) - out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, cu2) - out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, _make_cu([300])) - _assert_close("recycled_single/out", out_opt, out_ref) - _assert_close("recycled_single/state", state_opt, state_ref) - - -@torch.no_grad() -def test_single_seq_same_object_repeated(dispatch_pair, gate_params): - """Two calls with the SAME cu_seqlens object and a non-aligned single - seq. The Phase 2.1 path must sentinel-pad g on BOTH calls; poisoning - _varlen_pure_cache with the per-call override made the second call skip - the pad while still compiling the mask-free variant.""" - optimized, reference = dispatch_pair - cu = _make_cu([300]) - for i in range(2): - q, k, v, g, beta = _make_inputs(300, seed=40 + i) - out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, cu) - out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, _make_cu([300])) - _assert_close(f"sameobj_call{i}/out", out_opt, out_ref) - _assert_close(f"sameobj_call{i}/state", state_opt, state_ref) - - -@torch.no_grad() -def test_recycled_id_input_wrap_cache(): - """_ct_cached wrappers are keyed by (id(tensor), etype); the wrapper pins - the tensor's storage, so entries for live tensors are immortal by - design. Soundness requires the entry to be popped when the keyed tensor - dies before its id can be recycled — otherwise a new tensor on the - recycled id would silently reuse a wrapper over freed storage.""" - mod = _op_module() - import cutlass - - t1 = torch.zeros(64, 64, dtype=torch.bfloat16, device="cuda") - w1 = mod._ct_cached(t1, cutlass.BFloat16) - assert mod._ct_cached(t1, cutlass.BFloat16) is w1 # cache hit path - old_id = id(t1) - key = (old_id, cutlass.BFloat16) - assert key in mod._input_wrap_cache - del t1 - gc.collect() - - if key in mod._input_wrap_cache: - # The wrapper (still referenced by the cache) pins the tensor object, - # so the entry survives — sound only if the id can NOT be recycled - # while the entry is present. Prove no new tensor lands on old_id. - t2 = _alloc_with_recycled_id( - old_id, - lambda: torch.zeros(64, 64, dtype=torch.bfloat16, device="cuda"), - attempts=64, - ) - assert t2 is None, ( - "UNSOUND: id recycled while a stale _input_wrap_cache entry for " - "it is still present — a new tensor would reuse a wrapper over " - "another tensor's storage" - ) - # Dropping the last wrapper ref must let the finalizer prune. - del w1 - mod._input_wrap_cache.pop(key, None) - gc.collect() - else: - # Finalizer fired at t1's dealloc — entry pruned before any recycle. - del w1 - - # Fresh tensor (recycled id or not) must get a fresh wrapper. - t3 = torch.zeros(64, 64, dtype=torch.bfloat16, device="cuda") - w3 = mod._ct_cached(t3, cutlass.BFloat16) - assert mod._input_wrap_cache[(id(t3), cutlass.BFloat16)] is w3 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 66dc16e9149a..1aecf629339e 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 @@ -181,91 +181,6 @@ def test_optimized_prefill_matches_fla_reference() -> None: _assert_close(actual, expected) -@torch.no_grad() -def test_kda_prefill_op_empty_token_batch(): - """T=0 call: no output rows, recurrent state passes through unchanged. - - The runtime can emit a context batch with an empty token payload - (observed under the overlap scheduler + logprobs flows). The op used - to raise ``RuntimeError: step must be nonzero`` from its - ``arange(step=T)`` buffer setup; the FLA fallback tolerates the call. - """ - num_heads, head_k, head_v = 4, 128, 128 - q = torch.empty(1, 0, num_heads, head_k, dtype=torch.bfloat16, device="cuda") - k = torch.empty_like(q) - g = torch.empty(1, 0, num_heads, head_k, dtype=torch.float32, device="cuda") - v = torch.empty(1, 0, num_heads, head_v, dtype=torch.bfloat16, device="cuda") - beta = torch.empty(1, 0, num_heads, dtype=torch.float32, device="cuda") - initial_state = torch.randn(1, num_heads, head_k, head_v, dtype=torch.float32, device="cuda") - - output, final_state = torch.ops.trtllm.kda_prefill( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=head_k**-0.5, - initial_state=initial_state, - output_final_state=True, - ) - assert output.shape == (1, 0, num_heads, head_v) - torch.testing.assert_close(final_state, initial_state) - # State must not alias the input (the caller copies it back into the - # pool the initial state may be a view of). - assert final_state.data_ptr() != initial_state.data_ptr() - - _, final_state_zero = torch.ops.trtllm.kda_prefill( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=head_k**-0.5, - initial_state=None, - output_final_state=True, - ) - torch.testing.assert_close(final_state_zero, torch.zeros_like(initial_state)) - - -@torch.no_grad() -def test_kda_prefill_op_empty_token_batch_variants(): - """Regression coverage for the T=0 guard's other entry shapes. - - - varlen (cu_seqlens present): n_seqs derives from cu_seqlens, not B - - output_final_state=False: the op's empty-tensor final-state fallback - - use_fused_k1234=True: the guard must precede the fused-path branch - """ - num_heads, head_k, head_v = 4, 128, 128 - q = torch.empty(1, 0, num_heads, head_k, dtype=torch.bfloat16, device="cuda") - k = torch.empty_like(q) - g = torch.empty(1, 0, num_heads, head_k, dtype=torch.float32, device="cuda") - v = torch.empty(1, 0, num_heads, head_v, dtype=torch.bfloat16, device="cuda") - beta = torch.empty(1, 0, num_heads, dtype=torch.float32, device="cuda") - common = dict(q=q, k=k, v=v, g=g, beta=beta, scale=head_k**-0.5) - - # Varlen: two zero-length sequences -> final_state per sequence. - cu_seqlens = torch.tensor([0, 0, 0], dtype=torch.long, device="cuda") - _, final_state = torch.ops.trtllm.kda_prefill( - **common, initial_state=None, output_final_state=True, cu_seqlens=cu_seqlens - ) - assert final_state.shape == (2, num_heads, head_k, head_v) - assert (final_state == 0).all() - - # output_final_state=False: op returns the empty placeholder tensor. - output, final_state = torch.ops.trtllm.kda_prefill( - **common, initial_state=None, output_final_state=False - ) - assert output.shape == (1, 0, num_heads, head_v) - assert final_state.numel() == 0 - - # Fused path: the guard must fire before _launch_fused_k1234. - output, final_state = torch.ops.trtllm.kda_prefill( - **common, initial_state=None, output_final_state=True, use_fused_k1234=True - ) - assert output.shape == (1, 0, num_heads, head_v) - assert final_state.shape == (1, num_heads, head_k, head_v) - - @torch.no_grad() def test_kda_mixer_empty_prefill(): """The production mixer handles an empty token payload without raising.""" @@ -357,8 +272,8 @@ def test_kda_prefill_small_varlen_dispatch_matches_fla_reference(sequence_length @torch.no_grad() -def test_kda_prefill_opt_out_fallback_preserves_mixed_initial_states(monkeypatch): - """The explicit FLA fallback preserves prefix/continuation pool states.""" +def test_kda_prefill_state_pool_matches_fallback_with_mixed_initial_states(monkeypatch): + """The indexed pool path matches FLA for prefix and continuation states.""" torch.manual_seed(3) optimized, _ = _make_attention_pair() monkeypatch.setenv("TLLM_KDA_ENABLE_OPT_PREFILL", "0") @@ -376,7 +291,7 @@ def test_kda_prefill_opt_out_fallback_preserves_mixed_initial_states(monkeypatch * 0.05 ) slots = 5 - slot_indices = torch.tensor([3, 0, 4, 1], dtype=torch.long, device="cuda") + slot_indices = torch.tensor([3, 0, 4, 1], dtype=torch.int32, device="cuda") has_initial_states = torch.tensor([True, False, True, False], device="cuda") projection_size = NUM_HEADS * HEAD_DIM conv_seed = ( diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py deleted file mode 100644 index b712bbbe760a..000000000000 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py +++ /dev/null @@ -1,397 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""State parity tests: optimized KDA prefill dispatch vs FLA, kernel level. - -Complements test_kda_prefill_op.py (module outputs, no cache) with the cases -the executor runtime actually exercises through KDAKernelDispatch: - - * final_state parity in the pool's V-first [N, H, V, K] layout — K == V == - 128 for Kimi K3, so a layout mix-up is invisible to shape checks and only - caught numerically; - * a LARGE carried initial_state — the K4 sign error fixed by dev-tech - commit e45ae259 (NV = U - W@S, was U + W@S) is masked by tiny/zero - initial states and random inputs, and only blows up when a real-magnitude - state is carried in; - * non-64-aligned varlen (single-seq pad path and multi-seq masked path); - * fresh input tensors on every call — the runtime never reuses tensor - objects, so this exercises per-call wrapper rebuild instead of the - benchmark-style stable-object path, and would catch both stale-cache - corruption and per-call pinning leaks. -""" - -import pytest -import torch - -pytest.importorskip("fla") - -from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import KDAKernelDispatch # noqa: E402 - -NUM_HEADS = 96 -HEAD_K_DIM = 128 -LOWER_BOUND = -5.0 - - -def _has_supported_gpu() -> bool: - return torch.cuda.is_available() and torch.cuda.get_device_capability(0) in {(10, 0), (10, 3)} - - -pytestmark = pytest.mark.skipif( - not _has_supported_gpu(), - reason="Kimi K3 is supported only on Blackwell (SM100/SM103)", -) - - -@pytest.fixture(scope="module") -def dispatch_pair(): - optimized = KDAKernelDispatch(use_optimized_prefill=True, use_optimized_decode=False) - assert optimized.prefill_kernel_path == "optimized" - reference = KDAKernelDispatch(use_optimized_prefill=False, use_optimized_decode=False) - assert reference.prefill_kernel_path == "fla" - return optimized, reference - - -@pytest.fixture(scope="module") -def gate_params(): - torch.manual_seed(0) - a_log = torch.randn(NUM_HEADS, dtype=torch.float32, device="cuda") * 0.5 - dt_bias = torch.randn(NUM_HEADS * HEAD_K_DIM, dtype=torch.float32, device="cuda") * 0.1 - return a_log, dt_bias - - -def _make_inputs(batch: int, total_t: int, h0_scale, seed: int): - gen = torch.Generator(device="cuda").manual_seed(seed) - h, k = NUM_HEADS, HEAD_K_DIM - - def rnd(*shape, dtype=torch.bfloat16, scale=1.0): - return ( - torch.randn(*shape, generator=gen, dtype=torch.float32, device="cuda").to(dtype) * scale - ) - - q = rnd(batch, total_t, h, k) - key = rnd(batch, total_t, h, k) - v = rnd(batch, total_t, h, k) - g = rnd(batch, total_t, h, k) - beta = rnd(batch, total_t, h, dtype=torch.float32) - h0 = None - if h0_scale is not None: - # Pool V-first [N, H, V, K] layout, deliberately non-symmetric in the - # last two dims (row-index modulation) so a [K,V]<->[V,K] transpose - # mix-up cannot cancel out. - h0 = rnd(batch, h, k, k, dtype=torch.float32, scale=h0_scale) - h0 = h0 * torch.linspace(0.5, 1.5, k, device="cuda").view(1, 1, k, 1) - return q, key, v, g, beta, h0 - - -def _run(dispatch, gate_params, q, k, v, g, beta, h0, cu): - a_log, dt_bias = gate_params - return dispatch.prefill_chunk_kda( - q=q.clone(), - k=k.clone(), - v=v.clone(), - g=g.clone(), - beta=beta.clone(), - A_log=a_log, - dt_bias=dt_bias, - scale=HEAD_K_DIM**-0.5, - initial_state=h0.clone() if h0 is not None else None, - safe_gate=True, - lower_bound=LOWER_BOUND, - cu_seqlens=cu.clone() if cu is not None else None, - ) - - -def _assert_close(name, actual, expected): - actual, expected = actual.float(), expected.float() - cos = torch.nn.functional.cosine_similarity(actual.flatten(), expected.flatten(), dim=0).item() - rel = ((actual - expected).norm() / (expected.norm() + 1e-12)).item() - assert cos > 0.999 and rel < 3e-2, f"{name}: cos={cos:.6f} rel_l2={rel:.3e}" - - -CASES = [ - # (label, batch, eqlen T or None, varlen seq lens or None, h0_scale) - ("eqlen_b2_t256_no_state", 2, 256, None, None), - ("eqlen_b2_t256_large_state", 2, 256, None, 1.0), - ("eqlen_b1_t300_pad_large_state", 1, 300, None, 1.0), - ("varlen_aligned_large_state", 1, None, [128, 256, 192], 1.0), - ("varlen_nonaligned_large_state", 1, None, [100, 257, 64], 1.0), - ("varlen_single_nonaligned_large_state", 1, None, [300], 1.0), -] - - -@pytest.mark.parametrize("label,batch,eqlen_t,lens,h0_scale", CASES, ids=[c[0] for c in CASES]) -@torch.no_grad() -def test_state_parity(dispatch_pair, gate_params, label, batch, eqlen_t, lens, h0_scale): - optimized, reference = dispatch_pair - if lens is not None: - total_t = sum(lens) - cu = torch.tensor( - [0] + torch.cumsum(torch.tensor(lens), 0).tolist(), dtype=torch.long, device="cuda" - ) - n_seqs = len(lens) - else: - total_t, cu, n_seqs = eqlen_t, None, batch - case_idx = [c[0] for c in CASES].index(label) - q, k, v, g, beta, h0 = _make_inputs( - 1 if lens is not None else batch, total_t, h0_scale, seed=100 + case_idx - ) - if h0 is not None and lens is not None: - h0 = h0[:1].expand(n_seqs, -1, -1, -1).contiguous() * torch.linspace( - 0.5, 1.5, n_seqs, device="cuda" - ).view(n_seqs, 1, 1, 1) - - out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, h0, cu) - out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, h0, cu) - _assert_close(f"{label}/out", out_opt, out_ref) - _assert_close(f"{label}/state", state_opt, state_ref) - - -@torch.no_grad() -def test_fresh_tensors_every_call(dispatch_pair, gate_params): - """Same shape, new tensor objects + contents per call. Catches stale - id-keyed cache hits (wrong data); a persistently climbing allocation - here would indicate per-call activation pinning.""" - optimized, reference = dispatch_pair - for i in range(8): - q, k, v, g, beta, h0 = _make_inputs(2, 256, 1.0, seed=1000 + i) - out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, h0, None) - out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, h0, None) - _assert_close(f"iter{i}/out", out_opt, out_ref) - _assert_close(f"iter{i}/state", state_opt, state_ref) - - -# Eval-scale coverage (2026-07-23): GSM8K on the optimized prefill scored -# 68.99 vs 97.01 on the FLA path while every existing unit test passed — -# the gap is the eval regime: varlen batches packing dozens of -# mixed-length sequences, and chunked prefill chaining a sequence's state -# across calls. These tests parity-check exactly those two regimes. - -_EVAL_SCALE_LENS = [(97 * (i + 3)) % 911 + 45 for i in range(24)] - - -@pytest.mark.parametrize("h0_scale", [None, 1.0], ids=["no_state", "state"]) -@torch.no_grad() -def test_eval_scale_packed_varlen_parity(dispatch_pair, gate_params, h0_scale): - """24 packed mixed-length sequences (45..955 tokens), like one - max_num_tokens=8192 eval context batch.""" - optimized, reference = dispatch_pair - lens = _EVAL_SCALE_LENS - n_seqs, total_t = len(lens), sum(lens) - cu = torch.tensor( - [0] + torch.cumsum(torch.tensor(lens), 0).tolist(), dtype=torch.long, device="cuda" - ) - q, k, v, g, beta, h0 = _make_inputs(1, total_t, h0_scale, seed=4242) - if h0 is not None: - h0 = h0[:1].expand(n_seqs, -1, -1, -1).contiguous() * torch.linspace( - 0.5, 1.5, n_seqs, device="cuda" - ).view(n_seqs, 1, 1, 1) - out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, h0, cu) - out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, h0, cu) - _assert_close("evalscale/out", out_opt, out_ref) - _assert_close("evalscale/state", state_opt, state_ref) - - -# GSM8K-shaped coverage (2026-07-23, second regression pass): the stream -# fix lifted GSM8K 68.99 -> 77.79, still below the 97.01 FLA control, and -# the partial-score trajectory declined with completion length — the -# residual defect skews toward LONG sequences (8-shot GSM8K prompts are -# ~1-1.2k tokens = 16-20 chunks/seq; the eval-scale test above caps at -# 955). These cases pin per-seq NT at 16-20. The weak-gate variants -# (g_scale=0.05) reduce the per-chunk decay so state accumulates across -# many chunks — accumulation-order/precision defects that strong random -# gates (effectively local memory) cannot surface. - -_LONG_MIXED_LENS = [1150, 1200, 980, 1100, 1279, 1024, 1216, 1090] - - -@pytest.mark.parametrize("g_scale", [1.0, 0.05], ids=["g_normal", "g_weak"]) -@pytest.mark.parametrize( - "lens,h0_scale", - [([1150], None), ([1150], 1.0), (_LONG_MIXED_LENS, None), (_LONG_MIXED_LENS, 1.0)], - ids=["single_long", "single_long_state", "mixed_long", "mixed_long_state"], -) -@torch.no_grad() -def test_long_sequence_parity(dispatch_pair, gate_params, lens, h0_scale, g_scale): - optimized, reference = dispatch_pair - n_seqs, total_t = len(lens), sum(lens) - cu = torch.tensor( - [0] + torch.cumsum(torch.tensor(lens), 0).tolist(), dtype=torch.long, device="cuda" - ) - q, k, v, g, beta, h0 = _make_inputs(1, total_t, h0_scale, seed=9000 + len(lens)) - g = (g.float() * g_scale).to(g.dtype) - if h0 is not None: - h0 = h0[:1].expand(n_seqs, -1, -1, -1).contiguous() * torch.linspace( - 0.5, 1.5, n_seqs, device="cuda" - ).view(n_seqs, 1, 1, 1) - out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, h0, cu) - out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, h0, cu) - _assert_close("long/out", out_opt, out_ref) - _assert_close("long/state", state_opt, state_ref) - - -# Poisoned-scratch FINAL-STATE parity (2026-07-23, review follow-up on the -# fused_k123 beta-guard lines ~1520): the store guard (t < ci_eos) leaves -# k_scaled/kg/q_scaled scratch rows past each seq's end UNWRITTEN for -# partial final chunks. K4 is supposed to neutralize those rows via its -# per-seq bounded TMA (token extent = cu[s+1], hardware zero-fill); these -# cases PROVE that with garbage actually present in the stale rows: a -# preceding call with the SAME buffer-cache key (same B/T/NT/n_seqs) and -# huge-magnitude inputs fills the scratch, then the victim call's final -# state is asserted against FLA. Every eval prompt ends in a partial -# chunk and the state feeds all decode steps, so state corruption here is -# invisible to output-only checks. -# -# Victim/poison pairs share (T_total, NT_total, n_seqs) so they hit the -# same _buf_cache entry; the poison split is chosen so its guarded stores -# leave different rows written than the victim's layout expects. - -_POISON_STATE_CASES = [ - # (victim lens, poison lens) — equal sum and equal total chunk count - ([8191, 1], [8127, 65]), - ([8000, 150, 42], [7999, 129, 64]), - ([1186, 30], [1150, 66]), -] - - -def _nt(lens): - return sum((n + 63) // 64 for n in lens) - - -@pytest.mark.parametrize( - "victim,poison", _POISON_STATE_CASES, ids=["v8191_1", "v8000_150_42", "v1186_30"] -) -@torch.no_grad() -def test_final_state_parity_poisoned_scratch(dispatch_pair, gate_params, victim, poison): - optimized, reference = dispatch_pair - assert ( - sum(victim) == sum(poison) and _nt(victim) == _nt(poison) and len(victim) == len(poison) - ), "pairs must share the buffer-cache key" - total_t = sum(victim) - - # Poison pass: same buffer-cache key, huge magnitudes. - q, k, v, g, beta, _ = _make_inputs(1, total_t, None, seed=31) - cu_p = torch.tensor( - [0] + torch.cumsum(torch.tensor(poison), 0).tolist(), dtype=torch.long, device="cuda" - ) - _run(optimized, gate_params, q * 100, k * 100, v * 1000, g, beta * 10, None, cu_p) - torch.cuda.synchronize() - - # Victim pass: fresh inputs, partial final chunks per seq. - q, k, v, g, beta, h0 = _make_inputs(1, total_t, 1.0, seed=32) - n_seqs = len(victim) - h0 = h0[:1].expand(n_seqs, -1, -1, -1).contiguous() * torch.linspace( - 0.5, 1.5, n_seqs, device="cuda" - ).view(n_seqs, 1, 1, 1) - cu_v = torch.tensor( - [0] + torch.cumsum(torch.tensor(victim), 0).tolist(), dtype=torch.long, device="cuda" - ) - out_opt, state_opt = _run(optimized, gate_params, q, k, v, g, beta, h0, cu_v) - out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, h0, cu_v) - _assert_close("poisoned/out", out_opt, out_ref) - _assert_close("poisoned/state", state_opt, state_ref) - - -def _run_headcount_case(dispatch_pair, heads): - optimized, reference = dispatch_pair - h, k = heads, HEAD_K_DIM - gen = torch.Generator(device="cuda").manual_seed(4711 + heads) - a_log = torch.randn(h, generator=gen, dtype=torch.float32, device="cuda") * 0.5 - dt_bias = torch.randn(h * k, generator=gen, dtype=torch.float32, device="cuda") * 0.1 - lens = [1150, 731, 1024, 987] - total_t = sum(lens) - cu = torch.tensor( - [0] + torch.cumsum(torch.tensor(lens), 0).tolist(), dtype=torch.long, device="cuda" - ) - - def rnd(*shape, dtype=torch.bfloat16): - return torch.randn(*shape, generator=gen, dtype=torch.float32, device="cuda").to(dtype) - - q = rnd(1, total_t, h, k) - key = rnd(1, total_t, h, k) - v = rnd(1, total_t, h, k) - g = rnd(1, total_t, h, k) - beta = rnd(1, total_t, h, dtype=torch.float32) - h0 = rnd(len(lens), h, k, k, dtype=torch.float32) * torch.linspace( - 0.5, 1.5, k, device="cuda" - ).view(1, 1, k, 1) - - def run(dispatch): - return dispatch.prefill_chunk_kda( - q=q.clone(), - k=key.clone(), - v=v.clone(), - g=g.clone(), - beta=beta.clone(), - A_log=a_log, - dt_bias=dt_bias, - scale=k**-0.5, - initial_state=h0.clone(), - safe_gate=True, - lower_bound=LOWER_BOUND, - cu_seqlens=cu.clone(), - ) - - out_opt, state_opt = run(optimized) - out_ref, state_ref = run(reference) - _assert_close(f"h{heads}/out", out_opt, out_ref) - _assert_close(f"h{heads}/state", state_opt, state_ref) - - -@pytest.mark.parametrize("heads", [6, 12], ids=["tp16_heads6", "tp8_heads12"]) -@torch.no_grad() -def test_long_sequence_parity_tp_headcount(dispatch_pair, heads): - """Per-rank head count under tensor parallelism: the e2e eval runs - tp_size=16 -> H = 96/16 = 6 heads per rank, but every other test here - compiles the DSL kernels at H=96 (H is a compile-time specializer). - Parity-check the H=6 / H=12 compiles on a GSM8K-shaped batch.""" - _run_headcount_case(dispatch_pair, heads) - - -@torch.no_grad() -def test_headcount_recompile_parity(dispatch_pair): - """Cross-head-count compile-cache isolation, order-explicit: run two - different per-rank head counts back-to-back in the same process and - parity-check the SECOND one. Regression test for the K4 persistent - kernel cache key losing H (0e44bf64a6 follow-up): s_ct bakes the - [H, K, V] state shape/strides at compile time, so reusing the - first-compiled head count's kernel for a different H misaddresses - every (seq, head) state tile. The tp_headcount params above only catch - this via pytest execution order; this test pins the order even when - run in isolation. Heads 8 then 4 avoid cache hits from other tests.""" - _run_headcount_case(dispatch_pair, 8) - _run_headcount_case(dispatch_pair, 4) - - -@torch.no_grad() -def test_chunked_continuation_parity(dispatch_pair, gate_params): - """Chunked prefill: run [3000] as [2048] then [952] with the first - call's final state carried into the second (what the runtime's - chunked-prefill continuation does through the pool), and compare the - second chunk's output and final state against a single full-sequence - reference run.""" - optimized, reference = dispatch_pair - total_t, split = 3000, 2048 - q, k, v, g, beta, _ = _make_inputs(1, total_t, None, seed=777) - cu_full = torch.tensor([0, total_t], dtype=torch.long, device="cuda") - out_ref, state_ref = _run(reference, gate_params, q, k, v, g, beta, None, cu_full) - - cu1 = torch.tensor([0, split], dtype=torch.long, device="cuda") - cu2 = torch.tensor([0, total_t - split], dtype=torch.long, device="cuda") - sl1 = slice(0, split) - sl2 = slice(split, total_t) - out1, state1 = _run( - optimized, gate_params, q[:, sl1], k[:, sl1], v[:, sl1], g[:, sl1], beta[:, sl1], None, cu1 - ) - out2, state2 = _run( - optimized, - gate_params, - q[:, sl2], - k[:, sl2], - v[:, sl2], - g[:, sl2], - beta[:, sl2], - state1, - cu2, - ) - _assert_close("chunked/out_chunk1", out1, out_ref[:, sl1]) - _assert_close("chunked/out_chunk2", out2, out_ref[:, sl2]) - _assert_close("chunked/state", state2, state_ref) diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py new file mode 100644 index 000000000000..fbe39e5acd22 --- /dev/null +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for Kimi K3-specific Mamba metadata.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.modules.kimi_kda.kimi_k3_mamba_metadata import KimiK3MambaMetadata + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _make_attention_metadata(context_lengths: list[int]) -> SimpleNamespace: + seq_lens = torch.tensor(context_lengths or [1], dtype=torch.int) + num_contexts = len(context_lengths) + return SimpleNamespace( + seq_lens=seq_lens, + seq_lens_cuda=seq_lens.cuda(), + num_contexts=num_contexts, + num_ctx_tokens=sum(context_lengths), + kv_cache_manager=None, + request_ids=None, + kv_cache_params=SimpleNamespace( + num_cached_tokens_per_seq=torch.zeros(num_contexts, dtype=torch.int), + ), + ) + + +def test_k3_prepare_metadata_match_chunk_indices() -> None: + pytest.importorskip("fla") + from fla.ops.utils.index import prepare_chunk_indices + + metadata = KimiK3MambaMetadata(max_batch_size=4, chunk_size=128) + cases = ( + ([128, 256], True, None), + ([1, 64, 65, 129], False, None), + ([300], False, 300), + ) + for context_lengths, expected_alignment, expected_single_length in cases: + metadata.prepare(_make_attention_metadata(context_lengths)) + + assert metadata.kda_varlen_is_aligned is expected_alignment + assert metadata.kda_single_sequence_length == expected_single_length + cu_seqlens = metadata.query_start_loc_long[: len(context_lengths) + 1] + legacy_chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size=64) + torch.testing.assert_close(metadata.kda_chunk_indices, legacy_chunk_indices) + + +def test_prepare_refreshes_and_clears_kda_prefill_metadata() -> None: + metadata = KimiK3MambaMetadata(max_batch_size=4, chunk_size=128) + + metadata.prepare(_make_attention_metadata([1, 64, 65, 129])) + metadata.prepare(_make_attention_metadata([300])) + + expected = torch.stack( + ( + torch.zeros(5, dtype=torch.long, device="cuda"), + torch.arange(5, dtype=torch.long, device="cuda"), + ), + dim=1, + ) + assert metadata.kda_varlen_is_aligned is False + assert metadata.kda_single_sequence_length == 300 + torch.testing.assert_close(metadata.kda_chunk_indices, expected) + + metadata.prepare(_make_attention_metadata([])) + assert metadata.kda_chunk_indices is None + assert metadata.kda_varlen_is_aligned is None + assert metadata.kda_single_sequence_length is None diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py index f38aa46d1174..da2a0781d267 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py @@ -191,7 +191,7 @@ def test_fused_vs_sequential_two_rounds(): rt_fused.finalize_decode_weights() assert rt_fused._qkvg_proj_weight is not None assert rt_fused._bfa_proj_weight is not None - slot_indices = torch.arange(B, dtype=torch.long, device="cuda") + slot_indices = torch.arange(B, dtype=torch.int32, device="cuda") conv_pool_seq, ssm_pool_seq = _make_pools(B, seed=2) conv_pool_fused = conv_pool_seq.clone() diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py index 6756e663e63c..a9b299d804e1 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py @@ -54,6 +54,44 @@ class _K3Cfg: } +def test_forward_preserves_native_int32_state_indices(monkeypatch): + attention = KimiKDALinearAttention(_Cfg(), layer_idx=0) + state_indices = torch.tensor([3, 1], dtype=torch.int32) + captured = {} + + def forward_decode(hidden_states, conv_pool, ssm_pool, slot_indices, *args, **kwargs): + captured["slot_indices"] = slot_indices + captured["ssm_state_indices"] = kwargs["ssm_state_indices"] + return torch.zeros_like(hidden_states) + + monkeypatch.setattr(attention, "forward_decode", forward_decode) + layer_cache = SimpleNamespace( + conv=torch.empty(4, 3 * attention.proj_size, attention.conv_size), + temporal=torch.empty( + 4, + attention.num_heads, + attention.head_dim, + attention.head_k_dim, + ), + ) + metadata = SimpleNamespace( + mamba_metadata=SimpleNamespace( + state_indices=state_indices, + query_start_loc_long=torch.zeros(1, dtype=torch.long), + ), + num_contexts=0, + num_ctx_tokens=0, + seq_lens=torch.tensor([1, 1]), + kv_cache_manager=SimpleNamespace(mamba_layer_cache=lambda _: layer_cache), + ) + + output = attention(torch.empty(2, _Cfg.hidden_size), metadata) + + assert output.shape == (2, _Cfg.hidden_size) + assert captured["slot_indices"].data_ptr() == state_indices.data_ptr() + assert captured["ssm_state_indices"].data_ptr() == state_indices.data_ptr() + + class _LayerCache: def __init__(self, slots, dim3, w, h, v, k, t_max, device): self.conv = torch.zeros(slots, dim3, w, dtype=torch.bfloat16, device=device) @@ -70,7 +108,6 @@ def _decode_metadata(layer_cache: SimpleNamespace, slot_indices: torch.Tensor) - batch = slot_indices.shape[0] mamba_metadata = SimpleNamespace( state_indices=slot_indices.to(torch.int32), - state_indices_long=slot_indices, query_start_loc_long=torch.zeros(1, device=slot_indices.device, dtype=torch.long), _arange_buffer=torch.arange(batch + 1, device=slot_indices.device, dtype=torch.int32), ) @@ -90,8 +127,9 @@ def _prefill_metadata( batch = slot_indices.shape[0] mamba_metadata = SimpleNamespace( state_indices=slot_indices.to(torch.int32), - state_indices_long=slot_indices, + query_start_loc=cu_seqlens.to(torch.int32), query_start_loc_long=cu_seqlens, + has_initial_states=torch.zeros(batch, dtype=torch.bool, device=slot_indices.device), use_initial_states=False, ) cache_manager = SimpleNamespace(mamba_layer_cache=lambda _: layer_cache) diff --git a/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py index c0d0f1b19098..d8fb41c2ce7a 100644 --- a/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py +++ b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py @@ -516,8 +516,8 @@ def test_pack_gdn_decode_qkv( @skip_no_cuda -def test_reset_gdn_states_preserves_initialized_and_invalid_slots(): - from tensorrt_llm._torch.modules.mamba.gdn_mixer import _reset_gdn_states +def test_recurrent_state_reset_supports_gdn_pool_views(): + from tensorrt_llm._torch.modules.mamba.recurrent_state_cache import reset_recurrent_state_rows num_slots = 4 ssm_state_size = 6 @@ -537,11 +537,11 @@ def test_reset_gdn_states_preserves_initialized_and_invalid_slots(): state_indices = torch.tensor([1, 3, -1], device="cuda", dtype=torch.int32) has_initial_states = torch.tensor([False, True, False], device="cuda", dtype=torch.bool) - _reset_gdn_states( + reset_recurrent_state_rows( ssm_states, - conv_states, state_indices, has_initial_states, + conv_states, ) torch.testing.assert_close(state_pool[1], torch.zeros_like(state_pool[1])) diff --git a/tests/unittest/_torch/modules/mamba/test_recurrent_state_cache.py b/tests/unittest/_torch/modules/mamba/test_recurrent_state_cache.py new file mode 100644 index 000000000000..75d23a1b341c --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/test_recurrent_state_cache.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for shared recurrent-state cache operations.""" + +import pytest +import torch + +from tensorrt_llm._torch.modules.mamba.recurrent_state_cache import reset_recurrent_state_rows + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@torch.no_grad() +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_reset_recurrent_state_rows(index_dtype: torch.dtype) -> None: + slots = 5 + recurrent_states = torch.arange( + slots * 2 * 3 * 4, + dtype=torch.float32, + device="cuda", + ).reshape(slots, 2, 3, 4) + conv_states = torch.arange( + slots * 3 * 2, + dtype=torch.bfloat16, + device="cuda", + ).reshape(slots, 3, 2) + recurrent_before = recurrent_states.clone() + conv_before = conv_states.clone() + conv_expected = conv_states.clone() + conv_expected[4].zero_() + + state_indices = torch.tensor([4, -1, 1, slots], dtype=index_dtype, device="cuda") + has_initial_states = torch.tensor([False, False, True, False], device="cuda") + reset_recurrent_state_rows( + recurrent_states, + state_indices, + has_initial_states, + conv_states, + ) + + assert torch.count_nonzero(recurrent_states[4]).item() == 0 + assert torch.count_nonzero(conv_states[4]).item() == 0 + torch.testing.assert_close(recurrent_states[:4], recurrent_before[:4]) + torch.testing.assert_close(conv_states[:4], conv_before[:4]) + + reset_recurrent_state_rows( + recurrent_states, + torch.tensor([2], dtype=index_dtype, device="cuda"), + torch.tensor([False], device="cuda"), + ) + assert torch.count_nonzero(recurrent_states[2]).item() == 0 + torch.testing.assert_close(conv_states, conv_expected)