diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 60b2b7c157a5..1f1f810286ea 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -986,8 +986,13 @@ def __init__( self._w_q_t = self._w_k_t = self._w_v_t = None self._A_log_f32 = self._dt_bias_f32 = self._onorm_w_f32 = None # Persistent batch-row-dense staging for the fused decode kernel's - # per-section conv windows (lazily sized to the pool slot count). + # per-section conv windows. Sized once, on the first decode call, + # to the conv pool's slot count and never reallocated (see + # ``_forward_decode``). self._cs_dense: Optional[torch.Tensor] = None + # fp32 [dim, W] conv weights for the fused verify kernel, prebuilt + # by ``_build_mtp_conv_weights()`` at weight-load finalize time. + self._mtp_conv_weights: Optional[Tuple[torch.Tensor, ...]] = None def finalize_decode_weights(self) -> None: """Build the decode fast-path constants (once, after weight load). @@ -1056,7 +1061,7 @@ def _build_decode_kernel_constants(self) -> None: self._onorm_w_f32 = mixer.o_norm.weight.detach().float().contiguous() # Build the fused-verify conv constants eagerly too, so the first # verify call never allocates (a capture-unsafe lazy allocation). - self._get_mtp_conv_weights() + self._build_mtp_conv_weights() def finalize_decode_weights_fp8(self) -> None: """FP8 counterpart of ``finalize_decode_weights()``. @@ -1343,12 +1348,13 @@ def _forward_decode( W = mixer.conv_size # Allocated ONCE at the pool slot count (== per-rank max batch on - # the Mixed manager) and never reallocated: captured CUDA graphs - # hold this pointer, so a later realloc would leave earlier graphs - # writing into freed memory. Footprint: slots x ~9(H=6)..222(H=96) - # KB per layer. + # the Mixed manager; ``slot_indices`` are distinct pool rows and + # this is the plain one-token-per-request path, so B never exceeds + # it) and never reallocated: captured CUDA graphs hold this + # pointer, so a realloc would leave earlier graphs writing into + # freed memory. Footprint: slots x ~9(H=6)..222(H=96) KB per layer. buf = self._cs_dense - if buf is None or buf.shape[1] < B: + if buf is None: if torch.cuda.is_current_stream_capturing(): # Never allocate inside CUDA graph capture; the reference # path is capture-safe (just slower). @@ -1359,6 +1365,15 @@ def _forward_decode( 3, max(conv_pool.shape[0], B), d, W - 1, dtype=torch.bfloat16, device=x2d.device ) self._cs_dense = buf + else: + # Fail loudly if the sizing invariant ever breaks: silently + # reallocating here would hand previously captured CUDA graphs + # a dangling pointer. + assert buf.shape[1] >= B, ( + f"KDA decode staging buffer holds {buf.shape[1]} rows but the " + f"decode batch is {B}; reallocating would corrupt previously " + f"captured CUDA graphs" + ) if self._in_proj_weight is not None: # One GEMV over [q | k | v | g | f_a | b]; slices below are views. @@ -1600,26 +1615,28 @@ def _forward_verify_fused( o = out.view(num_decodes, num_steps, H, mixer.head_dim) return self._output_gate_and_proj(x, o) - def _get_mtp_conv_weights(self): + def _build_mtp_conv_weights(self) -> None: + """Prebuild the fp32 ``[dim, W]`` conv weights for the fused verify + kernel (once, at weight-load finalize time). Building them lazily + at first use would allocate at runtime; under CUDA graph capture + that bakes capture-pool pointers into the cached tuple.""" + mixer = self.mixer + self._mtp_conv_weights = tuple( + conv.weight.detach().squeeze(1).float().contiguous() + for conv in (mixer.q_conv1d, mixer.k_conv1d, mixer.v_conv1d) + ) + + def _get_mtp_conv_weights(self) -> Tuple[torch.Tensor, ...]: """fp32 ``[dim, W]`` conv weights for the fused verify kernel, - computed once per runtime instance.""" - cached = getattr(self, "_mtp_conv_weights", None) + prebuilt by ``_build_mtp_conv_weights()``.""" + cached = self._mtp_conv_weights if cached is None: - if torch.cuda.is_current_stream_capturing(): - # Allocating inside CUDA graph capture would bake - # capture-pool pointers into the cached tuple; the constants - # are normally prebuilt by _build_decode_kernel_constants(). - raise RuntimeError( - "Kimi K3 fused-verify conv weights were not prebuilt " - "before CUDA graph capture; call finalize_decode_weights" - "() / finalize_decode_weights_fp8() after weight load." - ) - mixer = self.mixer - cached = tuple( - conv.weight.detach().squeeze(1).float().contiguous() - for conv in (mixer.q_conv1d, mixer.k_conv1d, mixer.v_conv1d) + raise RuntimeError( + "Kimi K3 fused-verify conv weights were not prebuilt; call " + "_build_mtp_conv_weights() (done by load_weights() and by " + "finalize_decode_weights() / finalize_decode_weights_fp8()) " + "after weight load and before the first verify step." ) - self._mtp_conv_weights = cached return cached def _forward_verify_sequential( @@ -2581,6 +2598,14 @@ def get_tensor(key): if not kda_fp8: layer.self_attn.finalize_decode_weights() num_kda_fused += int(layer.self_attn._in_proj_weight is not None) + # The fused-verify conv constants are needed on every + # configuration that can reach _forward_verify_fused, + # including ones where neither finalize variant runs (e.g. + # FP8 KDA weight read with the fused decode glue disabled), + # and are never computed lazily (a first verify under CUDA + # graph capture must not allocate). Build them + # unconditionally; three small fp32 tensors per layer. + layer.self_attn._build_mtp_conv_weights() logger.info( f"Kimi K3: loaded {len(param_jobs)} parameters and the expert " f"slices of {len(expert_jobs)} MoE layers; fused decode " diff --git a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py index 9006ed8fb7ee..37c6469b4ce6 100644 --- a/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py +++ b/tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py @@ -103,6 +103,10 @@ def _make_runtime(seed): * 0.03 ).to(p.dtype) ) + # The fused-verify conv constants are prebuilt at weight-load finalize + # time in production; the runtime never computes them lazily. Mirror + # that here (after the random init above, which they snapshot). + rt._build_mtp_conv_weights() return rt