[Kimi-K3][Bugfix] Fix low-latency GEMM fallback initialization (#54167)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Thien Tran committed
Aug 28, 2026 at 16:32 UTC
f956e1c343dc63720b50c01209bd5ada6d8589a2
2 files changed
+7
-13
tests/kernels/test_bf16_skinny_gemm.py
+6
-13
@@ -1031,19 +1031,12 @@ def test_residual_dispatch_falls_back_to_addmm(
1031
assert output is fallback
1032
1033
1034
-def test_fallback_preserves_default_method(monkeypatch: pytest.MonkeyPatch) -> None:
1035
- fallback = torch.empty(2, 8)
1036
- monkeypatch.setattr(
1037
- k3_gemm.UnquantizedLinearMethod,
1038
- "apply",
1039
- lambda *args: fallback,
1040
- )
1041
- # 1-D input fails the runtime check, forcing the base-method fallback.
1034
+def test_fallback_preserves_default_method() -> None:
1035
+ x = torch.randn(2, 4)
1036
+ weight = torch.randn(3, 4)
1037
+ # CPU tensors fail the runtime check, forcing the base-method fallback.
1038
method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {})
1039
1044
- output = method.apply(
1045
- SimpleNamespace(weight=torch.empty(0)),
1046
- torch.empty(0),
1047
- )
1040
+ output = method.apply(SimpleNamespace(weight=weight), x)
1041
1049
- assert output is fallback
1042
+ torch.testing.assert_close(output, torch.nn.functional.linear(x, weight))
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+1
@@ -725,6 +725,7 @@ class _KimiK3LowLatencyApply:
725
"""Mixin: try the precomputed plan, else defer to the base method."""
726
727
def __init__(self, plan: dict[int, ResolvedCall]) -> None:
728
+ super().__init__()
729
self._plan = plan
730
731
def apply(