diff --git a/aiter/ops/flydsl/kimi_k3_kda_input_group64.py b/aiter/ops/flydsl/kimi_k3_kda_input_group64.py index a2ca838915..f4bf0ec2f3 100644 --- a/aiter/ops/flydsl/kimi_k3_kda_input_group64.py +++ b/aiter/ops/flydsl/kimi_k3_kda_input_group64.py @@ -105,8 +105,8 @@ def kimi_k3_kda_input_group64( scale: torch.Tensor, output: torch.Tensor | None = None, *, - rows_per_wave: int = 1, - cu_count: int = 240, + rows_per_wave: int = 2, + cu_count: int = 256, weight_cache_modifier: int = 2, ) -> torch.Tensor: """Launch only after the caller has passed the typed support predicate.""" diff --git a/op_tests/flydsl_tests/test_kimi_k3_kda_input_group64.py b/op_tests/flydsl_tests/test_kimi_k3_kda_input_group64.py index f20f5b715a..7ed29d366e 100644 --- a/op_tests/flydsl_tests/test_kimi_k3_kda_input_group64.py +++ b/op_tests/flydsl_tests/test_kimi_k3_kda_input_group64.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +import importlib + import pytest import torch @@ -14,6 +16,8 @@ supports_kimi_k3_kda_input_group64, ) +group64_module = importlib.import_module("aiter.ops.flydsl.kimi_k3_kda_input_group64") + def test_support_predicate_fails_closed_off_gpu() -> None: tensor = torch.empty(1) @@ -25,6 +29,34 @@ def test_quantizer_rejects_non_cuda_input() -> None: quantize_kimi_k3_kda_input_group64(torch.empty(1, dtype=torch.bfloat16)) +def test_wrapper_uses_validated_gfx950_schedule_by_default(monkeypatch) -> None: + hidden = torch.empty((1, 7168), dtype=torch.bfloat16) + weight = torch.empty((6284, 7168), dtype=torch.float8_e4m3fn) + scale = torch.empty((6284, 112), dtype=torch.float32) + output = torch.empty((1, 6288), dtype=torch.bfloat16) + schedules = [] + launches = [] + + monkeypatch.setattr( + group64_module, + "supports_kimi_k3_kda_input_group64", + lambda *args: True, + ) + monkeypatch.setattr( + group64_module, + "_launcher", + lambda *args: ( + schedules.append(args) or (lambda *args, **kwargs: launches.append(args)) + ), + ) + monkeypatch.setattr(group64_module, "ptr_arg", lambda tensor: tensor) + monkeypatch.setattr(torch.cuda, "current_stream", lambda device: None) + + assert kimi_k3_kda_input_group64(hidden, weight, scale, output) is output + assert schedules == [(2, 256, 2)] + assert len(launches) == 1 + + @pytest.mark.parametrize("rows_per_wave", [0, 5]) def test_builder_rejects_unsupported_rows_per_wave(rows_per_wave: int) -> None: with pytest.raises(ValueError, match="rows_per_wave"):