diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 16d7a9742f..1fd4982e76 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -330,6 +330,7 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused( cu_seqlens=cu_seqlens_t, use_qk_l2norm_in_kernel=True, ) + output = torch.empty_like(v) new_o, new_ht = chunk_kda_with_fused_gate( q=q.clone(), k=k.clone(), @@ -343,8 +344,10 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused( output_final_state=True, cu_seqlens=cu_seqlens_t, use_qk_l2norm_in_kernel=True, + out=output, ) + assert new_o.data_ptr() == output.data_ptr() assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3) assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3) @@ -442,6 +445,7 @@ def test_packed_kda_decode_correctness( use_qk_l2norm_in_kernel=True, ) packed_state = state + packed_output = torch.empty_like(dense_out) packed_out, _ = PACKED_DECODE_IMPLS[impl]( mixed_qkv=mixed_qkv, raw_g=raw_g, @@ -451,8 +455,11 @@ def test_packed_kda_decode_correctness( lower_bound=lower_bound, initial_state=packed_state, state_indices=state_indices, + **({"out": packed_output} if impl == "nvidia" else {}), ) + if impl == "nvidia": + assert packed_out.data_ptr() == packed_output.data_ptr() assert_close("o", dense_out, packed_out, 1e-3, err_atol=1e-3) assert_close("ht", dense_state, packed_state, 1e-3, err_atol=1e-3) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index f344b27656..c3f423804a 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -512,6 +512,8 @@ def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata(): assert actual.non_spec_query_start_loc is None torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0, 1])) torch.testing.assert_close(actual.spec_token_indx, torch.tensor([2, 3, 4])) + assert actual.non_spec_token_start == 0 + assert actual.spec_token_start == 2 torch.testing.assert_close( actual.spec_query_start_loc, torch.tensor([0, 3], dtype=torch.int32), @@ -539,6 +541,8 @@ def test_mixed_regular_and_spec_decode_excludes_request_padding(): assert actual.non_spec_state_indices_tensor.shape == (1,) torch.testing.assert_close(actual.non_spec_token_indx, torch.tensor([0])) torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3])) + assert actual.non_spec_token_start == 0 + assert actual.spec_token_start == 1 @pytest.mark.parametrize("mamba_cache_mode", ["none", "align"]) diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 52039acab1..92326944f5 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -903,6 +903,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): non_spec_state_indices_tensor = m.non_spec_state_indices_tensor spec_token_indx = m.spec_token_indx non_spec_token_indx = m.non_spec_token_indx + spec_token_start = m.spec_token_start + non_spec_token_start = m.non_spec_token_start spec_state_indices_tensor = m.spec_state_indices_tensor spec_query_start_loc = m.spec_query_start_loc num_accepted_tokens = m.num_accepted_tokens @@ -981,6 +983,18 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): mixed_qkv_spec = mixed_qkv g1_spec, beta_spec = g1, beta mixed_qkv_ns = g1_ns = beta_ns = None + elif spec_token_start is not None: + assert non_spec_token_start is not None + spec_end = spec_token_start + m.num_spec_decode_tokens + non_spec_end = ( + non_spec_token_start + m.num_prefill_tokens + m.num_decode_tokens + ) + spec_slice = slice(spec_token_start, spec_end) + non_spec_slice = slice(non_spec_token_start, non_spec_end) + mixed_qkv_spec = mixed_qkv[spec_slice] + g1_spec, beta_spec = g1[:, spec_slice], beta[:, spec_slice] + mixed_qkv_ns = mixed_qkv[non_spec_slice] + g1_ns, beta_ns = g1[:, non_spec_slice], beta[:, non_spec_slice] else: assert spec_token_indx is not None assert non_spec_token_indx is not None @@ -1029,6 +1043,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): if m.num_prefills == 0 and m.num_decodes == 0 else None ) + if spec_token_start is not None: + spec_out = core_attn_out[:, spec_slice] if self.use_recoverssm: from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( kda_recoverssm_verify, @@ -1076,6 +1092,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): core_attn_out_non_spec = None if mixed_qkv_ns is not None: assert g1_ns is not None and beta_ns is not None + non_spec_out = None + if non_spec_token_start is not None: + non_spec_out = core_attn_out[:, non_spec_slice] if m.num_prefills > 0: q_ns, k_ns, v_ns = mixed_qkv_ns.split( self.local_projection_size, dim=-1 @@ -1126,6 +1145,8 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): flashkda_out = ( workspace_out if has_spec_decode else core_attn_out )[:, : q_ns.shape[1]] + if non_spec_out is not None: + flashkda_out = non_spec_out if checkpoint is not None: assert non_spec_query_start_loc is not None num_sequences = initial_state.shape[0] @@ -1211,7 +1232,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): if q_ns.shape[1] > initial_state.shape[0]: assert m.flashinfer_prefill_seq_order is not None flashinfer_out = core_attn_out[:, : q_ns.shape[1]] - if has_spec_decode: + if non_spec_out is not None: + flashinfer_out = non_spec_out + elif has_spec_decode: assert self._flashinfer_kda_output_spec is not None (workspace_out,) = current_workspace_manager().get_simultaneous( self._flashinfer_kda_output_spec @@ -1251,6 +1274,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): output_final_state=True, use_qk_l2norm_in_kernel=True, cu_seqlens=non_spec_query_start_loc, + out=non_spec_out, ) recurrent_state[non_spec_state_indices_tensor] = ( last_recurrent_state.to(recurrent_state.dtype) @@ -1284,12 +1308,18 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): lower_bound=self.gate_lower_bound, initial_state=recurrent_state, state_indices=decode_conv_indices, + out=non_spec_out, ) # Restore the scheduler's original token order for mixed batches. if core_attn_out_spec is not None and core_attn_out_non_spec is not None: - core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec) - core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) + if spec_token_start is None: + assert spec_token_indx is not None + assert non_spec_token_indx is not None + core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec) + core_attn_out.index_copy_( + 1, non_spec_token_indx, core_attn_out_non_spec + ) elif core_attn_out_non_spec is not None: if ( self.kda_prefill_backend not in ("flashkda", "flashinfer") diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 21fdbea0b1..0b395b0c6c 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -270,6 +270,8 @@ class KDACheckpointMetadata: @dataclass class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata): + spec_token_start: int | None = None + non_spec_token_start: int | None = None flashinfer_prefill_query_start_loc: torch.Tensor | None = None flashinfer_prefill_seq_order: torch.Tensor | None = None recoverssm_commit: KDARecoverSSMCommitMetadata | None = None @@ -422,6 +424,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): spec_sequence_masks_cpu = None spec_request_indices = None + spec_token_start = None + non_spec_token_start = None if num_spec_decodes == 0: # The runner orders ordinary decodes before prefills. num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( @@ -518,6 +522,13 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): non_spec_token_indx = index[:num_non_spec_tokens] spec_token_indx = index[num_non_spec_tokens:] + active_spec_mask = spec_sequence_masks_cpu[query_lens_cpu > 0] + # check if spec / non spec tokens are continuous + if (active_spec_mask[1:] != active_spec_mask[:-1]).sum().item() == 1: + spec_first = active_spec_mask[0].item() + spec_token_start = 0 if spec_first else num_non_spec_tokens + non_spec_token_start = num_spec_decode_tokens if spec_first else 0 + # Native spec uses one state slot per step. RecoverSSM keeps # only the current checkpoint slot. spec_state_indices_tensor = block_table_tensor[ @@ -759,6 +770,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): spec_token_indx=spec_token_indx, non_spec_token_indx=non_spec_token_indx, num_accepted_tokens=num_accepted_tokens, + spec_token_start=spec_token_start, + non_spec_token_start=non_spec_token_start, flashinfer_prefill_query_start_loc=flashinfer_prefill_query_start_loc, flashinfer_prefill_seq_order=flashinfer_prefill_seq_order, recoverssm_commit=recoverssm_commit, diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py index 01bc609128..416c316a70 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/chunk.py @@ -598,6 +598,7 @@ def _chunk_kda_fwd_with_cumulative_g( chunk_indices: torch.Tensor | None = None, chunk_size: int = FLA_CHUNK_SIZE, safe_gate: bool = False, + out: torch.Tensor | None = None, ): # `g` must already be chunk-local cumulatively-summed AND scaled by # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)). @@ -642,7 +643,7 @@ def _chunk_kda_fwd_with_cumulative_g( g=g, A=Aqk, h=h, - o=v, + o=v if out is None else out, scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, @@ -706,6 +707,7 @@ def chunk_kda_with_fused_gate_fwd( output_final_state: bool, lower_bound: float | None = None, cu_seqlens: torch.Tensor | None = None, + out: torch.Tensor | None = None, ): chunk_size = FLA_CHUNK_SIZE chunk_indices = ( @@ -736,6 +738,7 @@ def chunk_kda_with_fused_gate_fwd( chunk_indices=chunk_indices, chunk_size=chunk_size, safe_gate=lower_bound is not None, + out=out, ) @@ -787,6 +790,7 @@ def chunk_kda_with_fused_gate( lower_bound: float | None = None, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.Tensor | None = None, + out: torch.Tensor | None = None, **kwargs, ): """Run chunk KDA from raw gate and beta projections.""" @@ -810,6 +814,7 @@ def chunk_kda_with_fused_gate( output_final_state=output_final_state, lower_bound=lower_bound, cu_seqlens=cu_seqlens, + out=out, ) return o, final_state diff --git a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py index b4d35d8557..7f5f98400f 100644 --- a/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py +++ b/vllm/models/kimi_k3/nvidia/ops/third_party/kda/fused_recurrent.py @@ -591,6 +591,7 @@ def fused_recurrent_kda_packed_decode( initial_state: torch.Tensor, state_indices: torch.Tensor, scale: float | None = None, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Run one-token KDA decode directly from packed post-conv QKV.""" if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1: @@ -639,7 +640,8 @@ def fused_recurrent_kda_packed_decode( if scale is None: scale = K**-0.5 - out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device) + if out is None: + out = torch.empty((1, B, H, V), dtype=mixed_qkv.dtype, device=device) grid = (cdiv(V, BV), B * H) fused_recurrent_kda_packed_decode_kernel[grid]( mixed_qkv=mixed_qkv,