diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index c55028f1e522..8fbf106d5bdc 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -974,6 +974,10 @@ def __init__( lin = cfg.linear_attn_config self.layer_idx = layer_idx self._use_indexed_ssm_pool = _KDA_INDEXED_STATE_POOL_ENABLED + # Opt-in while prefill token-shape performance is evaluated. When + # fused weights are unavailable, _forward_prefill still falls back + # to the original per-projection calls. + self._use_fused_prefill_proj = os.getenv("TLLM_KDA_ENABLE_FUSED_PREFILL_PROJ", "0") == "1" # Attention-family TP semantics (Qwen3-Next GatedDeltaNet pattern, # gdn_mixer.py): replicated under attention-DP — each rank runs # its own batch with the full head set — and head-sharded across @@ -1010,11 +1014,11 @@ def __init__( use_optimized_decode=True, ) self.proj_size = (num_heads // self._kda_tp_size) * lin["head_dim"] - # Decode fast-path constants, built once by + # Fused projection and decode-kernel constants, built once by # ``finalize_decode_weights()`` after checkpoint load: # fused in-projection weight + kernel-layout conv weights + fp32 - # copies of the small parameters. ``None`` routes decode through - # the reference (module-level) path. + # copies of the small parameters. ``None`` keeps projection calls + # separate and routes decode through the reference module path. self._in_proj_weight: Optional[torch.Tensor] = None # FP8 variant (``finalize_decode_weights_fp8``): q/k/v/g stay in the # mixer's fused FP8 ``qkvg_proj`` GEMM and only the small [f_a | b] @@ -1032,9 +1036,9 @@ def finalize_decode_weights(self) -> None: 1. Fused in-projection ``[q | k | v | g | f_a | b]``: all six projections consume the same hidden state, so one GEMV replaces five GEMV+splitK-reduce pairs plus a cublas dot per layer per - decode step. The source parameters are repointed to row views - of the fused buffer (no extra memory; prefill/verify paths keep - using them unchanged). + decode step. The fused verify path and the opt-in prefill path + reuse the same projection; source parameters remain row views for + fallback paths. 2. Kernel-layout constants that ``_decode_via_optimized`` used to rebuild with ~6 device kernels per layer per decode step: transposed conv weights (bf16 ``[W, D]``) and fp32 copies of @@ -1092,6 +1096,29 @@ def _build_decode_kernel_constants(self) -> None: self._dt_bias_f32 = mixer.dt_bias.detach().float().contiguous() self._onorm_w_f32 = mixer.o_norm.weight.detach().float().contiguous() + def _fused_input_projections( + self, x: torch.Tensor + ) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]: + """Project shared KDA inputs through the finalized fused weights.""" + mixer = self.mixer + d = self.proj_size + hd = mixer.head_dim + H = mixer.num_heads + + if self._in_proj_weight is not None: + proj = torch.nn.functional.linear(x, self._in_proj_weight) + x_qkv, onorm_g, f_a, beta = proj.split([3 * d, d, hd, H], dim=-1) + elif self._in_proj_small_weight is not None: + qkvg = mixer.qkvg_proj(x) + x_qkv, onorm_g = qkvg.split([3 * d, d], dim=-1) + small = torch.nn.functional.linear(x, self._in_proj_small_weight) + f_a, beta = small.split([hd, H], dim=-1) + else: + return None + + g = mixer.f_b_proj(f_a) + return x_qkv, onorm_g, g, beta + def finalize_decode_weights_fp8(self) -> None: """FP8 counterpart of ``finalize_decode_weights()``. @@ -1252,15 +1279,24 @@ def _forward_prefill( d = self.proj_size x = x2d.unsqueeze(0) # [1, T, hidden] - fused_qkvg = getattr(mixer, "qkvg_proj", None) - if fused_qkvg is not None: - q_proj_states, k_proj_states, v_proj_states = fused_qkvg(x)[..., : 3 * d].split( - d, dim=-1 - ) + projections = self._fused_input_projections(x) if self._use_fused_prefill_proj else None + if projections is not None: + qkv, onorm_g, g, beta = projections + q_proj_states, k_proj_states, v_proj_states = qkv.split(d, dim=-1) + beta = beta.float() else: - q_proj_states = mixer.q_proj(x) - k_proj_states = mixer.k_proj(x) - v_proj_states = mixer.v_proj(x) + fused_qkvg = getattr(mixer, "qkvg_proj", None) + if fused_qkvg is not None: + q_proj_states, k_proj_states, v_proj_states = fused_qkvg(x)[..., : 3 * d].split( + d, dim=-1 + ) + else: + q_proj_states = mixer.q_proj(x) + k_proj_states = mixer.k_proj(x) + v_proj_states = mixer.v_proj(x) + onorm_g = None + g = mixer.f_b_proj(mixer.f_a_proj(x)) + beta = mixer.b_proj(x).float() # Initial states: present for continuation chunks (chunked prefill) # and for prefix-cache hits (block reuse), where the previous @@ -1286,9 +1322,7 @@ def _forward_prefill( ) del q_proj_states, k_proj_states, v_proj_states - g = mixer.f_b_proj(mixer.f_a_proj(x)) g = rearrange(g, "... (h d) -> ... h d", d=mixer.head_dim) - beta = mixer.b_proj(x).float() q = rearrange(q, "... (h d) -> ... h d", d=mixer.head_k_dim) k = rearrange(k, "... (h d) -> ... h d", d=mixer.head_k_dim) @@ -1323,7 +1357,7 @@ def _forward_prefill( # are zero for a fresh request, so the tail columns are unused). self._sync_kda_replay_conv_window(layer_cache, slot_indices, conv_q, conv_k, conv_v) - return self._output_gate_and_proj(x, o) + return self._output_gate_and_proj(x, o, onorm_g) def _forward_decode( self, @@ -1401,24 +1435,9 @@ def _forward_decode( ) self._cs_dense = buf - if self._in_proj_weight is not None: - # One GEMV over [q | k | v | g | f_a | b]; slices below are views. - proj = torch.nn.functional.linear(x2d, self._in_proj_weight) - x_qkv = proj[:, : 3 * d] - onorm_g = proj[:, 3 * d : 4 * d] - f_a = proj[:, 4 * d : 4 * d + hd] - beta = proj[:, 4 * d + hd : 4 * d + hd + H] - else: - # FP8 weight read (KIMI_K3_KDA_GLUE_FP8=1): the loader's fused - # FP8 [q | k | v | g] GEMM plus one BF16 GEMV over [f_a | b]; - # slices below are views. - qkvg = mixer.qkvg_proj(x2d) - small = torch.nn.functional.linear(x2d, self._in_proj_small_weight) - x_qkv = qkvg[:, : 3 * d] - onorm_g = qkvg[:, 3 * d : 4 * d] - f_a = small[:, :hd] - beta = small[:, hd : hd + H] - g = mixer.f_b_proj(f_a) # [B, d] + projections = self._fused_input_projections(x2d) + assert projections is not None + x_qkv, onorm_g, g, beta = projections # Gather the HF-layout conv windows once, then repack the # historical W-1 columns into the kernel's dense per-section @@ -1588,13 +1607,25 @@ def _forward_verify_fused( x = x2d.view(num_decodes, num_steps, -1) # [B, T, hidden] T_total = num_decodes * num_steps - x_q = mixer.q_proj(x).view(1, T_total, H, K) - x_k = mixer.k_proj(x).view(1, T_total, H, K) - x_v = mixer.v_proj(x).view(1, T_total, H, mixer.head_dim) + projections = self._fused_input_projections(x) + if projections is None: + x_q = mixer.q_proj(x) + x_k = mixer.k_proj(x) + x_v = mixer.v_proj(x) + onorm_g = None + g = mixer.f_b_proj(mixer.f_a_proj(x)) + beta = mixer.b_proj(x) + else: + x_qkv, onorm_g, g, beta = projections + x_q, x_k, x_v = x_qkv.split(self.proj_size, dim=-1) + + x_q = x_q.view(1, T_total, H, K) + x_k = x_k.view(1, T_total, H, K) + x_v = x_v.view(1, T_total, H, mixer.head_dim) # Raw gate / beta: the kernel applies dt_bias, A_log, the # lower-bound sigmoid gate, and the beta sigmoid itself. - g = mixer.f_b_proj(mixer.f_a_proj(x)).view(1, T_total, H, K) - beta = mixer.b_proj(x).view(1, T_total, H) + g = g.view(1, T_total, H, K) + beta = beta.view(1, T_total, H) w_q, w_k, w_v = self._get_mtp_conv_weights() lower_bound = ( @@ -1639,7 +1670,7 @@ def _forward_verify_fused( scale=mixer.head_k_dim**-0.5, ) o = out.view(num_decodes, num_steps, H, mixer.head_dim) - return self._output_gate_and_proj(x, o) + return self._output_gate_and_proj(x, o, onorm_g) def _get_mtp_conv_weights(self): """fp32 ``[dim, W]`` conv weights for the fused verify kernel, @@ -1734,14 +1765,20 @@ def _forward_verify_sequential( o = torch.cat(step_outputs, dim=1) # [B, T, H, V] return self._output_gate_and_proj(x, o) - def _output_gate_and_proj(self, x: torch.Tensor, o: torch.Tensor) -> torch.Tensor: + def _output_gate_and_proj( + self, + x: torch.Tensor, + o: torch.Tensor, + g_out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: from einops import rearrange mixer = self.mixer - if mixer.use_full_rank_gate: - g_out = mixer.g_proj(x) - else: - g_out = mixer.g_b_proj(mixer.g_a_proj(x)) + if g_out is None: + if mixer.use_full_rank_gate: + g_out = mixer.g_proj(x) + else: + g_out = mixer.g_b_proj(mixer.g_a_proj(x)) g_out = rearrange(g_out, "... (h d) -> ... h d", d=mixer.head_dim) o = mixer.o_norm(o, g_out) o = rearrange(o, "b t h d -> (b t) (h d)") 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 646e233a8a69..161190ca1749 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 @@ -41,8 +41,8 @@ _HAVE_DEPS = True _DEP_ERR = None try: - import cutlass # noqa: F401 import cuda.bindings.driver # noqa: F401 + import cutlass # noqa: F401 from fla.ops.kda import fused_recurrent_kda # noqa: F401 except ImportError as e: _HAVE_DEPS = False @@ -68,11 +68,11 @@ def _is_blackwell(): W = 4 M = 2 # draft tokens per round LB = -5.0 +FUSED_PREFILL_ENV = "TLLM_KDA_ENABLE_FUSED_PREFILL_PROJ" def _make_runtime(seed): - from tensorrt_llm._torch.models.modeling_kimi_linear import \ - KimiKDARuntime + from tensorrt_llm._torch.models.modeling_kimi_linear import KimiKDARuntime cfg = SimpleNamespace( hidden_size=HIDDEN, @@ -90,24 +90,32 @@ def _make_runtime(seed): with torch.no_grad(): for name, p in rt.named_parameters(): if name.endswith("A_log"): - p.copy_(torch.randn(p.shape, generator=gen, device="cuda", - dtype=torch.float32) * 0.5) + p.copy_( + torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.5 + ) elif name.endswith("dt_bias"): - p.copy_(torch.randn(p.shape, generator=gen, device="cuda", - dtype=torch.float32) * 0.1) + p.copy_( + torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) * 0.1 + ) else: - p.copy_((torch.randn(p.shape, generator=gen, device="cuda", - dtype=torch.float32) * 0.03).to(p.dtype)) + p.copy_( + ( + torch.randn(p.shape, generator=gen, device="cuda", dtype=torch.float32) + * 0.03 + ).to(p.dtype) + ) + rt.finalize_decode_weights() + assert rt._in_proj_weight is not None return rt def _make_pools(B, seed): gen = torch.Generator(device="cuda").manual_seed(seed) d = H * K - conv_pool = (torch.randn(B, 3 * d, W, generator=gen, device="cuda", - dtype=torch.float32) * 0.5).to(torch.bfloat16) - ssm_pool = torch.randn(B, H, K, K, generator=gen, device="cuda", - dtype=torch.float32) + conv_pool = ( + torch.randn(B, 3 * d, W, generator=gen, device="cuda", dtype=torch.float32) * 0.5 + ).to(torch.bfloat16) + ssm_pool = torch.randn(B, H, K, K, generator=gen, device="cuda", dtype=torch.float32) ssm_pool *= torch.linspace(0.5, 1.5, K, device="cuda").view(1, 1, K, 1) return conv_pool, ssm_pool @@ -120,23 +128,18 @@ def _make_fused_layer_cache(B, conv_pool): S = W - 1 + M def _conv_cache(section): - cache = torch.zeros(B, S, d, device="cuda", - dtype=torch.float32).transpose(-1, -2) - cache[:, :, :W - 1] = conv_pool[:, section * d:(section + 1) * d, - 1:].float() + cache = torch.zeros(B, S, d, device="cuda", dtype=torch.float32).transpose(-1, -2) + cache[:, :, : W - 1] = conv_pool[:, section * d : (section + 1) * d, 1:].float() return cache return SimpleNamespace( kda_conv_q=_conv_cache(0), kda_conv_k=_conv_cache(1), kda_conv_v=_conv_cache(2), - kda_qkg_cache=torch.zeros(B, M, 3, d, device="cuda", - dtype=torch.float32), + kda_qkg_cache=torch.zeros(B, M, 3, d, device="cuda", dtype=torch.float32), kda_v_cache=torch.zeros(B, M, d, device="cuda", dtype=torch.float32), - kda_beta_cache=torch.zeros(B, M, H, device="cuda", - dtype=torch.float32), - prev_num_accepted_tokens=torch.zeros(B, dtype=torch.int32, - device="cuda"), + kda_beta_cache=torch.zeros(B, M, H, device="cuda", dtype=torch.float32), + prev_num_accepted_tokens=torch.zeros(B, dtype=torch.int32, device="cuda"), intermediate_conv_window=None, intermediate_ssm=None, ) @@ -146,11 +149,10 @@ def _make_seq_layer_cache(B): d = H * K return SimpleNamespace( kda_qkg_cache=None, - intermediate_conv_window=torch.zeros(B, M + 1, 3 * d, W, - device="cuda", - dtype=torch.bfloat16), - intermediate_ssm=torch.zeros(B, M + 1, H, K, K, device="cuda", - dtype=torch.float32), + intermediate_conv_window=torch.zeros( + B, M + 1, 3 * d, W, device="cuda", dtype=torch.bfloat16 + ), + intermediate_ssm=torch.zeros(B, M + 1, H, K, K, device="cuda", dtype=torch.float32), ) @@ -164,13 +166,49 @@ def _promote_sequential(layer_cache, conv_pool, ssm_pool, accept): def _rep(name, a, b): a, b = a.float(), b.float() - cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), - dim=0).item() + cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() rel = ((a - b).norm() / (b.norm() + 1e-12)).item() print(f" {name}: cos={cos:.6f} rel_l2={rel:.3e}") return cos > 0.999 and rel < 3e-2 +def test_fused_prefill_projection_matches_separate_gemms(monkeypatch): + B = 2 + lengths = (5, 3) + T = sum(lengths) + + monkeypatch.delenv(FUSED_PREFILL_ENV, raising=False) + sequential = _make_runtime(seed=11) + assert not sequential._use_fused_prefill_proj + monkeypatch.setenv(FUSED_PREFILL_ENV, "1") + fused = _make_runtime(seed=11) + assert fused._use_fused_prefill_proj + + gen = torch.Generator(device="cuda").manual_seed(12) + x = (torch.randn(T, HIDDEN, generator=gen, device="cuda", dtype=torch.float32) * 0.5).to( + torch.bfloat16 + ) + cu_seqlens = torch.tensor([0, lengths[0], T], dtype=torch.int64, device="cuda") + slot_indices = torch.arange(B, dtype=torch.long, device="cuda") + metadata = SimpleNamespace(use_initial_states=False) + layer_cache = SimpleNamespace(kda_qkg_cache=None) + + conv_seq, ssm_seq = _make_pools(B, seed=13) + conv_fused = conv_seq.clone() + ssm_fused = ssm_seq.clone() + with torch.no_grad(): + out_seq = sequential._forward_prefill( + x, cu_seqlens, metadata, B, conv_seq, ssm_seq, slot_indices, layer_cache + ) + out_fused = fused._forward_prefill( + x, cu_seqlens, metadata, B, conv_fused, ssm_fused, slot_indices, layer_cache + ) + + assert _rep("prefill out", out_fused, out_seq) + assert _rep("prefill conv", conv_fused, conv_seq) + assert _rep("prefill ssm", ssm_fused, ssm_seq) + + def test_fused_vs_sequential_two_rounds(): torch.manual_seed(0) B = 4 @@ -187,18 +225,20 @@ def test_fused_vs_sequential_two_rounds(): gen = torch.Generator(device="cuda").manual_seed(3) def tokens(scale=0.5): - return (torch.randn(B * T, HIDDEN, generator=gen, device="cuda", - dtype=torch.float32) * scale).to(torch.bfloat16) + return ( + torch.randn(B * T, HIDDEN, generator=gen, device="cuda", dtype=torch.float32) * scale + ).to(torch.bfloat16) ok = True with torch.no_grad(): # ---- Round 1 (no pending drafts) ---- x1 = tokens() - out1_seq = rt._forward_verify_sequential(x1, T, cache_seq, - conv_pool_seq, ssm_pool_seq, - slot_indices) - out1_fused = rt._forward_verify(x1, T, cache_fused, conv_pool_fused, - ssm_pool_fused, slot_indices) + out1_seq = rt._forward_verify_sequential( + x1, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices + ) + out1_fused = rt._forward_verify( + x1, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices + ) print("round 1:") ok &= _rep("out", out1_fused, out1_seq) @@ -209,19 +249,21 @@ def tokens(scale=0.5): # ---- Round 2 (fused path replays the accepted drafts) ---- x2 = tokens() - out2_seq = rt._forward_verify_sequential(x2, T, cache_seq, - conv_pool_seq, ssm_pool_seq, - slot_indices) - out2_fused = rt._forward_verify(x2, T, cache_fused, conv_pool_fused, - ssm_pool_fused, slot_indices) + out2_seq = rt._forward_verify_sequential( + x2, T, cache_seq, conv_pool_seq, ssm_pool_seq, slot_indices + ) + out2_fused = rt._forward_verify( + x2, T, cache_fused, conv_pool_fused, ssm_pool_fused, slot_indices + ) print("round 2 (mixed replay):") ok &= _rep("out", out2_fused, out2_seq) # Committed pool state cross-check: fused pool holds the state after # round-2's golden token; reproduce it in the sequential world by # promoting with accept=0 (golden only). - _promote_sequential(cache_seq, conv_pool_seq, ssm_pool_seq, - torch.zeros(B, dtype=torch.long, device="cuda")) + _promote_sequential( + cache_seq, conv_pool_seq, ssm_pool_seq, torch.zeros(B, dtype=torch.long, device="cuda") + ) ok &= _rep("committed ssm", ssm_pool_fused, ssm_pool_seq) assert ok @@ -229,4 +271,5 @@ def tokens(scale=0.5): if __name__ == "__main__": import sys + sys.exit(pytest.main([__file__, "-v", "-x", "-s"]))