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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user