diff --git a/tests/models/kimi_k3/test_amd_kda_checkpoint.py b/tests/models/kimi_k3/test_amd_kda_checkpoint.py index 6f3129a224..2f8bd3f0e4 100644 --- a/tests/models/kimi_k3/test_amd_kda_checkpoint.py +++ b/tests/models/kimi_k3/test_amd_kda_checkpoint.py @@ -482,7 +482,7 @@ def _spec_for(monkeypatch, **attrs) -> MambaSpec: layer = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention) for name, value in dict( use_prefill_checkpoint=True, - use_fused_chunk=True, + kda_prefill_backend="fused", num_spec=0, use_safe_gate=True, ).items(): @@ -506,7 +506,8 @@ def test_opt_in_publishes_chunk_walk_alignment(monkeypatch) -> None: "attr,value", [ ("use_prefill_checkpoint", False), - ("use_fused_chunk", False), + ("kda_prefill_backend", "triton"), + ("kda_prefill_backend", "gluon"), ("use_safe_gate", False), ], ) diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py index c53e9bacdb..2dcf886c91 100644 --- a/tests/models/kimi_k3/test_amd_kda_chunk.py +++ b/tests/models/kimi_k3/test_amd_kda_chunk.py @@ -77,7 +77,7 @@ def _run(inp: dict, use_fused: bool, out: torch.Tensor | None = None, **kw): output_final_state=True, use_qk_l2norm_in_kernel=True, cu_seqlens=inp["cu"], - use_fused_chunk=use_fused, + backend=kw.pop("backend", "fused" if use_fused else "triton"), out=out, **kw, ) @@ -112,6 +112,27 @@ def test_fused_chunk_matches_triton(seqlens: list[int]) -> None: torch.testing.assert_close(state_fused, state_ref, rtol=1e-3, atol=1e-3) +@pytest.mark.parametrize( + "seqlens", + [[1], [64], [1024], [4096], [513, 64, 1, 1200], [512] * 6], +) +def test_gluon_chunk_matches_triton(seqlens: list[int]) -> None: + """AITER's Gluon chunk_kda must match the Triton chunk path too.""" + from vllm._aiter_ops import rocm_aiter_ops + + if not rocm_aiter_ops.is_kda_gluon_prefill_available(): + pytest.skip("AITER Gluon chunk_kda is not available here") + + inp = _inputs(seqlens, seed=len(seqlens)) + o_ref, state_ref = _run(inp, use_fused=False) + o_gluon, state_gluon = _run(inp, use_fused=False, backend="gluon") + + torch.testing.assert_close(o_gluon.float(), o_ref.float(), rtol=2e-2, atol=2e-2) + torch.testing.assert_close(state_gluon, state_ref, rtol=5e-3, atol=5e-3) + state_rel = (state_gluon - state_ref).norm() / state_ref.norm() + assert state_rel < 5e-3, f"final state relative L2 error {state_rel:.2e}" + + @pytest.mark.parametrize("use_fused", [False, True]) def test_out_argument_receives_the_result(use_fused: bool) -> None: """``out`` is a promise about where the result lands. @@ -648,7 +669,7 @@ def _run_paged(inp: dict, cache, idx, warm, use_fused: bool, **kw): has_initial_state=warm, use_qk_l2norm_in_kernel=True, cu_seqlens=inp["cu"], - use_fused_chunk=use_fused, + backend="fused" if use_fused else "triton", **kw, ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 6a6c99840a..29364f204a 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2606,6 +2606,27 @@ class rocm_aiter_ops: return cls.is_linear_enabled() and on_gfx950() + @staticmethod + @functools.cache + def _kda_gluon_prefill_importable() -> bool: + try: + from aiter.ops.triton.attention.chunk_kda import chunk_kda # noqa: F401 + + return True + except ImportError: + return False + + @classmethod + @if_aiter_supported + def is_kda_gluon_prefill_available(cls) -> bool: + from vllm.platforms.rocm import get_cdna_version + + return ( + cls._AITER_ENABLED + and get_cdna_version() >= 4 + and cls._kda_gluon_prefill_importable() + ) + @classmethod def get_aiter_allreduce(cls): """Return the TP device communicator's AITER custom-allreduce if it has @@ -5111,6 +5132,50 @@ class rocm_aiter_ops: layer_input.view(*outer_shape, hidden_size), ) + @staticmethod + def kda_chunk_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, + cu_seqlens: torch.Tensor | None, + chunk_indices: torch.Tensor | None, + chunk_offsets: torch.Tensor | None, + scale: float, + out: torch.Tensor | None, + initial_state: torch.Tensor | None, + output_final_state: bool, + state_cache: torch.Tensor | None, + state_indices: torch.Tensor | None, + has_initial_state: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + from aiter.ops.triton.attention.chunk_kda import chunk_kda + + return chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + scale=scale, + out=out, + initial_state=initial_state, + output_final_state=output_final_state, + state_cache=state_cache, + state_indices=state_indices, + has_initial_state=has_initial_state, + ) + rocm_aiter_ops.register_ops_once() _sync_aiter_situv2_moe_env() diff --git a/vllm/model_executor/warmup/kimi_k3_triton_warmup.py b/vllm/model_executor/warmup/kimi_k3_triton_warmup.py index e7883fff44..68186ac085 100644 --- a/vllm/model_executor/warmup/kimi_k3_triton_warmup.py +++ b/vllm/model_executor/warmup/kimi_k3_triton_warmup.py @@ -19,9 +19,11 @@ if TYPE_CHECKING: logger = init_logger(__name__) -def _get_kda_layer(worker: Worker) -> KimiK3DeltaAttention | None: +def _get_kda_layer( + worker: Worker, module: str = "vllm.models.kimi_k3.nvidia.kda" +) -> KimiK3DeltaAttention | None: # Kimi model construction already imports kda. Avoid importing it here. - kda = sys.modules.get("vllm.models.kimi_k3.nvidia.kda") + kda = sys.modules.get(module) if kda is None: return None @@ -172,9 +174,59 @@ def _warm_recurrent_kda( ) +def _warm_gluon_chunk_kda(layer, input_dtype: torch.dtype) -> None: + from vllm.models.kimi_k3.amd.ops.kda_prefill import ( + chunk_kda_prefill, + get_gluon_chunk_kda_warmup_profiles, + ) + + if layer.kda_prefill_backend != "gluon": + return + kv_cache = layer.kv_cache + if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) < 2: + return + state = kv_cache[1] + if not isinstance(state, torch.Tensor) or not state.numel(): + return + + logger.info("Warming up Kimi-K3 Gluon chunk KDA prefill kernels.") + h, d = int(layer.local_num_heads), int(layer.head_dim) + for num_seqs, seq_len in get_gluon_chunk_kda_warmup_profiles(): + num_tokens = num_seqs * seq_len + shape = (1, num_tokens, h, d) + chunk_kda_prefill( + q=torch.zeros(shape, dtype=input_dtype, device=state.device), + k=torch.zeros(shape, dtype=input_dtype, device=state.device), + v=torch.zeros(shape, dtype=input_dtype, device=state.device), + raw_g=torch.zeros(shape, dtype=input_dtype, device=state.device), + raw_beta=torch.zeros( + (1, num_tokens, h), dtype=input_dtype, device=state.device + ), + A_log=layer.A_log, + g_bias=layer.dt_bias, + lower_bound=layer.gate_lower_bound, + state_cache=state, + state_indices=torch.zeros(num_seqs, dtype=torch.int32, device=state.device), + has_initial_state=torch.zeros( + num_seqs, dtype=torch.bool, device=state.device + ), + use_qk_l2norm_in_kernel=True, + cu_seqlens=torch.arange( + 0, num_tokens + 1, seq_len, dtype=torch.int32, device=state.device + ), + backend="gluon", + out=torch.empty(shape, dtype=input_dtype, device=state.device), + ) + + @torch.inference_mode() def kimi_k3_triton_warmup(worker: Worker) -> None: """Warm Kimi-K3 Triton kernels reachable by this server.""" + if current_platform.is_rocm(): + layer = _get_kda_layer(worker, "vllm.models.kimi_k3.amd.kda") + if layer is not None: + _warm_gluon_chunk_kda(layer, worker.model_config.dtype) + return if not current_platform.is_cuda(): return diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index 611a0c5d96..08ad0b650d 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -95,7 +95,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): assert isinstance(spec, MambaSpec) # Only the fused path exports the checkpoint snapshots can_checkpoint = ( - self.use_prefill_checkpoint and self.use_fused_chunk and self.use_safe_gate + self.use_prefill_checkpoint + and self.kda_prefill_backend == "fused" + and self.use_safe_gate ) if can_checkpoint and vllm_config.cache_config.mamba_cache_mode == "align": logger.info_once("Kimi-K3 KDA prefill checkpoint enabled.") @@ -242,13 +244,15 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): "The fused KDA chunk kernel requires gfx950 and a build that " "includes it." ) - self.use_fused_chunk = backend == "fused" or ( - backend == "auto" and is_fused_kda_chunk_supported() - ) - logger.info_once( - "Kimi-K3 KDA prefill backend: %s", - "fused" if self.use_fused_chunk else "triton", - ) + self.kda_prefill_backend = backend + if backend == "auto": + if self.use_safe_gate and rocm_aiter_ops.is_kda_gluon_prefill_available(): + self.kda_prefill_backend = "gluon" + elif is_fused_kda_chunk_supported(): + self.kda_prefill_backend = "fused" + else: + self.kda_prefill_backend = "triton" + logger.info_once("Kimi-K3 KDA prefill backend: %s", self.kda_prefill_backend) # TODO: Remove the extra config arg once tested self.use_prefill_checkpoint = ( additional_config.get("kda_prefill_checkpoint", True) @@ -694,7 +698,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): cu_seqlens=prefill_query_start_loc, chunk_indices=m.chunk_indices, chunk_offsets=m.chunk_offsets, - use_fused_chunk=self.use_fused_chunk, + backend=self.kda_prefill_backend, out=core_attn_out[:, nd_tok:num_actual_tokens] if direct else None, checkpoint_state=checkpoint_state, checkpoint_offsets=checkpoint_offsets, diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py index 5ba2fd6379..040f3afe66 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py +++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py @@ -9,6 +9,7 @@ path. import torch +from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( gather_initial_states, @@ -24,6 +25,11 @@ from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE logger = init_logger(__name__) +# Consumed by kimi_k3_triton_warmup.py during kernel_warmup(). +def get_gluon_chunk_kda_warmup_profiles() -> tuple[tuple[int, int], ...]: + return ((1, 64), (2, 64), (4, 64), (8, 64), (16, 64), (1, 4096), (2, 2048)) + + def chunk_kda_prefill( q: torch.Tensor, k: torch.Tensor, @@ -40,7 +46,7 @@ def chunk_kda_prefill( cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.Tensor | None = None, chunk_offsets: torch.Tensor | None = None, - use_fused_chunk: bool = False, + backend: str = "triton", out: torch.Tensor | None = None, checkpoint_state: torch.Tensor | None = None, checkpoint_offsets: torch.Tensor | None = None, @@ -68,9 +74,10 @@ def chunk_kda_prefill( cu_seqlens: int32 cumulative sequence lengths. chunk_indices: precomputed chunk indices for ``cu_seqlens``. chunk_offsets: int32 per-sequence first chunk index. - use_fused_chunk: request the two-kernel ROCm path. It is used only when - every one of its preconditions holds; otherwise the Triton path - runs unchanged. + backend: ``"triton"`` (the vendored chunk path), ``"fused"`` (the + two-kernel ROCm path; used only when every one of its + preconditions holds, otherwise the Triton path runs unchanged) or + ``"gluon"`` (AITER's ``chunk_kda``). out: buffer the result must land in. Honoured by both backends, so the caller can hand in a slice of its own output and skip a copy. checkpoint_state: destination for mid-prefill recurrent state @@ -89,6 +96,11 @@ def chunk_kda_prefill( The output and, when requested, the final recurrent state. """ + if backend not in ("triton", "fused", "gluon"): + raise ValueError( + "chunk_kda_prefill backend must be 'triton', 'fused' or 'gluon', " + f"got {backend!r}." + ) if scale is None: scale = k.shape[-1] ** -0.5 @@ -96,7 +108,7 @@ def chunk_kda_prefill( # needs the raw projections and a bounded gate rather than the pre-normalized # tensors the Triton path takes. fused = ( - use_fused_chunk + backend == "fused" and use_qk_l2norm_in_kernel and cu_seqlens is not None and lower_bound is not None @@ -116,6 +128,28 @@ def chunk_kda_prefill( if state_indices is None or has_initial_state is None: raise ValueError("state_cache needs state_indices and has_initial_state") + if backend == "gluon": + return rocm_aiter_ops.kda_chunk_prefill( + q=q, + k=k, + v=v, + g=raw_g, + beta=raw_beta, + A_log=A_log, + dt_bias=g_bias, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + scale=scale, + out=out, + initial_state=initial_state, + output_final_state=output_final_state, + state_cache=state_cache, + state_indices=state_indices, + has_initial_state=has_initial_state, + ) + # Only the fused walk addresses the paged rows directly. The Triton path # gathers the rows it needs and scatters the results back here, so callers # pass the cache the same way for either backend.