diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index be916f6c0d0f..2480ae1f2a0c 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -15,6 +15,9 @@ from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( gather_initial_states, ) +from vllm.models.kimi_k3.amd.ops.third_party.kda import ( + chunk_kda_with_fused_gate as amd_chunk_kda_with_fused_gate, +) from vllm.models.kimi_k3.nvidia.kda import ( is_flashkda_supported, is_fused_kda_decode_supported, @@ -270,6 +273,84 @@ def test_chunk_kda_fused_gate_cumsum_matches_unfused( assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3) +@torch.inference_mode() +def test_chunk_kda_none_matches_zero_initial_state(): + H, D = 2, 128 + cu_seqlens = torch.tensor([0, 17, 49], dtype=torch.int32, device=DEVICE) + T = 49 + N = cu_seqlens.numel() - 1 + torch.manual_seed(123) + + kwargs = { + "q": torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE), + "k": torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE), + "v": torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE), + "raw_g": torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE), + "raw_beta": torch.randn(1, T, H, dtype=torch.bfloat16, device=DEVICE), + "A_log": torch.randn(H, dtype=torch.float32, device=DEVICE), + "g_bias": torch.randn(H * D, dtype=torch.float32, device=DEVICE), + "output_final_state": True, + "cu_seqlens": cu_seqlens, + "use_qk_l2norm_in_kernel": True, + } + zero_state = torch.zeros( + N, + H, + D, + D, + dtype=torch.float32, + device=DEVICE, + ) + + def run(initial_state: torch.Tensor | None, **extra_kwargs): + return amd_chunk_kda_with_fused_gate( + **{ + key: value.clone() if isinstance(value, torch.Tensor) else value + for key, value in kwargs.items() + }, + initial_state=initial_state, + **extra_kwargs, + ) + + output_with_zero, state_with_zero = run(zero_state) + output_without_state, state_without_state = run(None) + + torch.testing.assert_close(output_without_state, output_with_zero) + torch.testing.assert_close(state_without_state, state_with_zero) + + num_cache_rows = 5 + row_stride = H * D * D + 17 + cache_storage = torch.full( + (num_cache_rows * row_stride,), + torch.nan, + dtype=torch.float32, + device=DEVICE, + ) + final_state_cache = torch.as_strided( + cache_storage, + (num_cache_rows, H, D, D), + (row_stride, D * D, D, 1), + ) + final_state_indices = torch.tensor( + [3, 1], + dtype=torch.int32, + device=DEVICE, + ) + output_direct, returned_state = run( + None, + final_state_cache=final_state_cache, + final_state_indices=final_state_indices, + ) + + assert returned_state is None + torch.testing.assert_close(output_direct, output_with_zero) + torch.testing.assert_close( + final_state_cache[final_state_indices.long()], + state_with_zero, + ) + assert torch.isnan(final_state_cache[[0, 2, 4]]).all() + + @pytest.mark.parametrize("num_seqs", [1, 8, 32]) @pytest.mark.parametrize("lower_bound", [-5.0, None]) @pytest.mark.parametrize("state_indices_stride", [1, 8]) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 5352ef7a7a64..ae333d5ef8b9 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -188,6 +188,45 @@ def test_kimi_k3_kda_metadata_matches_shared_gdn( _assert_matches_shared_gdn(reference, actual) +@pytest.mark.parametrize( + ("batch", "expected"), + [ + pytest.param( + BatchSpec(seq_lens=[16, 32], query_lens=[16, 32]), + True, + id="all-fresh", + ), + pytest.param( + BatchSpec(seq_lens=[20, 32], query_lens=[4, 32]), + False, + id="one-resumed", + ), + pytest.param( + BatchSpec(seq_lens=[16, 0], query_lens=[16, 0]), + False, + id="zero-length-padding", + ), + ], +) +@pytest.mark.parametrize( + "builder_cls", + [GDNAttentionMetadataBuilder, KimiK3KDAMetadataBuilder], +) +def test_kda_metadata_marks_all_fresh_prefills( + batch: BatchSpec, + expected: bool, + builder_cls: type[AttentionMetadataBuilder], +): + common_attn_metadata = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE) + metadata = _make_builder( + builder_cls, + num_speculative_tokens=0, + full_cuda_graph=False, + ).build(0, common_attn_metadata) + + assert metadata.all_initial_states_fresh is expected + + def test_mixed_regular_and_spec_decode_uses_packed_decode_metadata(): batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3]) common_attn_metadata = create_common_attn_metadata( diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index bc49226c75b2..c7f7e525741d 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -556,10 +556,20 @@ def _prefill_conv( assert non_spec_state_indices_tensor is not None assert has_initial_state is not None - initial_state = gather_initial_states( - recurrent_state, - non_spec_state_indices_tensor, - has_initial_state, + # Chunk KDA initializes its recurrence to zero when no initial + # state is supplied. On ROCm it can also store final states + # directly into their cache rows. + direct_final_state = ( + current_platform.is_rocm() and m.all_initial_states_fresh + ) + initial_state = ( + None + if direct_final_state + else gather_initial_states( + recurrent_state, + non_spec_state_indices_tensor, + has_initial_state, + ) ) ( core_attn_out_non_spec, @@ -577,9 +587,21 @@ def _prefill_conv( output_final_state=True, use_qk_l2norm_in_kernel=True, cu_seqlens=non_spec_query_start_loc, + final_state_cache=( + recurrent_state if direct_final_state else None + ), + final_state_indices=( + non_spec_state_indices_tensor if direct_final_state else None + ), ) - # Init cache - recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state + if direct_final_state: + assert last_recurrent_state is None + else: + # Init cache + assert last_recurrent_state is not None + recurrent_state[non_spec_state_indices_tensor] = ( + last_recurrent_state + ) else: # pure-decode non-spec batch diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py index 77ed09d489ee..1fdf2750fc21 100644 --- a/vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/__init__.py @@ -16,9 +16,9 @@ # - OOB-mask correctness fix: present (all tl.load use mask=..., other=0). # Validated on gfx950: no core-dump, gsm8k 94.1%. # -# AMD-specific deltas vs the NVIDIA copy: NONE yet (byte-identical). Keep in sync -# with the NVIDIA copy on FLA updates; any divergence should be an intentional, -# documented gfx950-specific change (a #869-style AMD-only fix). +# AMD-specific deltas vs the NVIDIA copy: chunk KDA can write final recurrent +# states directly to indexed cache rows for all-fresh ROCm prefills. Keep other +# FLA updates in sync; divergence should remain intentional and documented. from .chunk import ( chunk_kda, diff --git a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py index 7d0ea0204ce6..cf4227db9b03 100644 --- a/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py +++ b/vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py @@ -596,6 +596,8 @@ def _chunk_kda_fwd_with_cumulative_g( chunk_indices: torch.Tensor | None = None, chunk_size: int = FLA_CHUNK_SIZE, safe_gate: bool = False, + final_state_cache: torch.Tensor | None = None, + final_state_indices: 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)). @@ -632,6 +634,8 @@ def _chunk_kda_fwd_with_cumulative_g( cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, use_exp2=True, + final_state_cache=final_state_cache, + final_state_indices=final_state_indices, ) del w, u, kg o = chunk_gla_fwd_o_gk( @@ -704,6 +708,8 @@ def chunk_kda_with_fused_gate_fwd( output_final_state: bool, lower_bound: float | None = None, cu_seqlens: torch.Tensor | None = None, + final_state_cache: torch.Tensor | None = None, + final_state_indices: torch.Tensor | None = None, ): chunk_size = FLA_CHUNK_SIZE chunk_indices = ( @@ -734,6 +740,8 @@ def chunk_kda_with_fused_gate_fwd( chunk_indices=chunk_indices, chunk_size=chunk_size, safe_gate=lower_bound is not None, + final_state_cache=final_state_cache, + final_state_indices=final_state_indices, ) @@ -785,6 +793,8 @@ def chunk_kda_with_fused_gate( lower_bound: float | None = None, use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.Tensor | None = None, + final_state_cache: torch.Tensor | None = None, + final_state_indices: torch.Tensor | None = None, **kwargs, ): """Run chunk KDA from raw gate and beta projections.""" @@ -808,6 +818,8 @@ def chunk_kda_with_fused_gate( output_final_state=output_final_state, lower_bound=lower_bound, cu_seqlens=cu_seqlens, + final_state_cache=final_state_cache, + final_state_indices=final_state_indices, ) return o, final_state diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 0bb4864d2918..3fdb402ccf84 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -403,11 +403,27 @@ def build( # type: ignore[override] # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares # its own chunk indices. Only causal-convolution metadata is needed here. nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None + all_initial_states_fresh = False if num_prefills > 0: has_initial_state = m.compute_num_computed_tokens() > 0 + has_initial_state_cpu = m._num_computed_tokens_cpu + non_spec_query_lens_for_fresh_cpu = query_start_loc_cpu.diff() if spec_sequence_masks_cpu is not None: has_initial_state = has_initial_state[active_non_spec_mask_cpu] + if has_initial_state_cpu is not None: + has_initial_state_cpu = has_initial_state_cpu[ + active_non_spec_mask_cpu + ] + non_spec_query_lens_for_fresh_cpu = ( + non_spec_query_lens_for_fresh_cpu[~spec_sequence_masks_cpu] + ) assert non_spec_query_start_loc_cpu is not None + if has_initial_state_cpu is not None: + all_initial_states_fresh = bool( + non_spec_query_lens_for_fresh_cpu.numel() > 0 + and (non_spec_query_lens_for_fresh_cpu > 0).all().item() + and not has_initial_state_cpu.any().item() + ) nums_dict, batch_ptr, token_chunk_offset_ptr = ( compute_causal_conv1d_metadata( non_spec_query_start_loc_cpu, @@ -472,6 +488,7 @@ def build( # type: ignore[override] num_spec_decode_tokens=num_spec_decode_tokens, num_actual_tokens=m.num_actual_tokens, has_initial_state=has_initial_state, + all_initial_states_fresh=all_initial_states_fresh, spec_query_start_loc=spec_query_start_loc, non_spec_query_start_loc=non_spec_query_start_loc, spec_state_indices_tensor=spec_state_indices_tensor, diff --git a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py index eb1c3af15297..b8b284117785 100644 --- a/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py +++ b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py @@ -27,6 +27,7 @@ "USE_GK": lambda args: args["gk"] is not None, "USE_INITIAL_STATE": lambda args: args["h0"] is not None, "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "USE_FINAL_STATE_INDICES": lambda args: args["ht_indices"] is not None, "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, } @@ -52,8 +53,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( h, h0, ht, + ht_indices, cu_seqlens, chunk_offsets, + ht_stride_n, T, H: tl.constexpr, Hg: tl.constexpr, @@ -65,6 +68,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( USE_GK: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, + USE_FINAL_STATE_INDICES: tl.constexpr, SAVE_NEW_VALUE: tl.constexpr, IS_VARLEN: tl.constexpr, USE_EXP2: tl.constexpr, @@ -107,7 +111,11 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( if USE_INITIAL_STATE: h0 = h0 + i_nh * V * K if STORE_FINAL_STATE: - ht = ht + i_nh * V * K + if USE_FINAL_STATE_INDICES: + i_state = tl.load(ht_indices + i_n).to(tl.int64) + ht = ht + i_state * ht_stride_n + i_h * V * K + else: + ht = ht + i_nh * V * K # load initial state if USE_INITIAL_STATE: @@ -331,7 +339,9 @@ def chunk_gated_delta_rule_fwd_h( chunk_indices: torch.Tensor | None = None, chunk_offsets: torch.Tensor | None = None, use_exp2: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: + final_state_cache: torch.Tensor | None = None, + final_state_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: # This kernel is slightly different from fla to support Q/K with different head numbers. # In fla, Q/K always have the same head number, so Hg is always equal to H. B, T, Hg, K, V = *k.shape, u.shape[-1] @@ -350,9 +360,24 @@ def chunk_gated_delta_rule_fwd_h( assert K <= 256, "current kernel does not support head dimension larger than 256." h = k.new_empty(B, NT, H, V, K) - final_state = ( - k.new_empty(N, H, V, K, dtype=torch.float32) if output_final_state else None - ) + if final_state_cache is not None: + assert output_final_state + assert final_state_indices is not None + assert final_state_indices.numel() == N + assert final_state_cache.shape[1:] == (H, V, K) + assert final_state_cache.dtype == torch.float32 + assert final_state_cache.device == k.device + assert final_state_cache.stride()[1:] == (V * K, K, 1) + final_state = final_state_cache + returned_final_state = None + else: + assert final_state_indices is None + final_state = ( + k.new_empty(N, H, V, K, dtype=torch.float32) + if output_final_state + else None + ) + returned_final_state = final_state v_new = torch.empty_like(u) if save_new_value else None @@ -369,8 +394,10 @@ def grid(meta): h=h, h0=initial_state, ht=final_state, + ht_indices=final_state_indices, cu_seqlens=cu_seqlens, chunk_offsets=chunk_offsets, + ht_stride_n=final_state.stride(0) if final_state is not None else 0, T=T, H=H, Hg=Hg, @@ -379,4 +406,4 @@ def grid(meta): BT=BT, USE_EXP2=use_exp2, ) - return h, v_new, final_state + return h, v_new, returned_final_state diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 0b32dfcf9444..6e31ee756fb5 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -49,6 +49,7 @@ class GDNAttentionMetadata: num_actual_tokens: int has_initial_state: torch.Tensor | None = None + all_initial_states_fresh: bool = False spec_query_start_loc: torch.Tensor | None = None # shape: [num_spec_decodes + 1,] non_spec_query_start_loc: torch.Tensor | None = ( @@ -330,6 +331,7 @@ def build( # type: ignore[override] prefill_query_start_loc: torch.Tensor | None = None prefill_state_indices: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None + all_initial_states_fresh = False if num_prefills > 0: from vllm.third_party.flash_linear_attention.ops.utils import ( FLA_CHUNK_SIZE, @@ -390,9 +392,24 @@ def build( # type: ignore[override] if num_prefills > 0: has_initial_state = context_lens_tensor > 0 + has_initial_state_cpu = m._num_computed_tokens_cpu + non_spec_query_lens_cpu = query_start_loc_cpu.diff() if spec_sequence_masks_cpu is not None: has_initial_state = has_initial_state[~spec_sequence_masks_cpu] + if has_initial_state_cpu is not None: + has_initial_state_cpu = has_initial_state_cpu[ + ~spec_sequence_masks_cpu + ] + non_spec_query_lens_cpu = non_spec_query_lens_cpu[ + ~spec_sequence_masks_cpu + ] assert non_spec_query_start_loc_cpu is not None + if has_initial_state_cpu is not None: + all_initial_states_fresh = bool( + non_spec_query_lens_cpu.numel() > 0 + and (non_spec_query_lens_cpu > 0).all().item() + and not has_initial_state_cpu.any().item() + ) nums_dict, batch_ptr, token_chunk_offset_ptr = ( compute_causal_conv1d_metadata( non_spec_query_start_loc_cpu, @@ -493,6 +510,7 @@ def build( # type: ignore[override] num_spec_decode_tokens=num_spec_decode_tokens, num_actual_tokens=m.num_actual_tokens, has_initial_state=has_initial_state, + all_initial_states_fresh=all_initial_states_fresh, chunk_indices=chunk_indices, chunk_offsets=chunk_offsets, prefill_query_start_loc=prefill_query_start_loc,