perf(flydsl): select validated Kimi-K3 KDA schedule

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."""

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"):