fix(GDN): dtype guard on _ix_matmul/_ix_bmm — ixformer.matmul requires kHalf

Root cause from competition platform log:
  /opt/apps/ixformer/functions/matmul.cu:149 'Expected input.dtype() == kHalf'
  Repeats ~80 times — every GDN layer token pass calls _ix_matmul with float32

GDN chunked delta rule uses float32 accumulation (correct for precision).
_ix_matmul was calling ixformer.matmul on float32 tensors → stderr spam.
The try/except caught it and fell back to torch.matmul, but the stderr
output floods the log and may slow down inference.

Fix: check a.dtype == torch.float16 before calling ixformer.matmul.
     Non-half tensors go directly to torch.matmul — zero stderr noise.
This commit is contained in:
EX Engine
2026-08-10 04:45:24 +00:00
parent 41aec33955
commit d841c44e55

View File

@@ -163,8 +163,8 @@ except ImportError:
# ---------------------------------------------------------------------------
def _ix_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""BI-V100 accelerated matmul via ixformer, fallback to torch."""
if _ix_available:
"""BI-V100 accelerated matmul via ixformer. Only for half — ixformer rejects float32."""
if _ix_available and a.dtype == torch.float16:
try:
return _ix.matmul(a, b)
except Exception:
@@ -172,8 +172,8 @@ def _ix_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
return torch.matmul(a, b)
def _ix_bmm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""Batched matmul — ixformer.matmul handles batched inputs."""
if _ix_available:
"""Batched matmul — ixformer.matmul handles batched half inputs."""
if _ix_available and a.dtype == torch.float16:
try:
return _ix.matmul(a, b)
except Exception:
@@ -181,8 +181,8 @@ def _ix_bmm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
return torch.matmul(a, b)
def _ix_softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
"""BI-V100 accelerated softmax via ixformer."""
if _ix_available:
"""BI-V100 accelerated softmax via ixformer. Only for half."""
if _ix_available and x.dtype == torch.float16:
try:
return _ix.softmax(x, dim=dim)
except Exception: