ViperEkura
7e98a419a7
- re-register the linear family with the operator dispatcher (ASTR_OPS / op_backend / resolve) - fix bf16 gemv misaligned-address faults and element mispairing for offset weights - reject misaligned bf16_swiglu inputs with a clear error and fall back in the backend gate - make the rollout reuse decision, validation, and return atomic under one policy snapshot - add the documented post-scoring rollout version check - derive live+1 under the scheduler lock in optimizer_step via apply_weight_update(None, ...) - reject rollout_max_policy_lag below rollout_interval - 1 at config time - sync gemv stream-test inputs before switching streams; drop dead loader imports
68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
"""Inference-only fused SwiGLU selection for dense MLP layers."""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import Tensor
|
|
|
|
from astrai.extension.backend.linear import linear
|
|
from astrai.extension.dispatch import env_mode
|
|
from astrai.extension.loader import is_available
|
|
from astrai.extension.ops.swiglu import bf16_swiglu
|
|
|
|
|
|
def _unfused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
|
|
# Keep the existing linear backend in the fallback chain. This preserves
|
|
# any independently qualified GEMV batches instead of making the fusion
|
|
# decision suppress linear-level optimizations.
|
|
return linear(x, up_weight) * F.silu(linear(x, gate_weight))
|
|
|
|
|
|
def _fused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
|
|
return bf16_swiglu(x.detach(), up_weight.detach(), gate_weight.detach())
|
|
|
|
|
|
def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
|
|
return not (
|
|
torch.is_grad_enabled()
|
|
or not x.is_cuda
|
|
or x.dtype != torch.bfloat16
|
|
or up_weight.dtype != torch.bfloat16
|
|
or gate_weight.dtype != torch.bfloat16
|
|
or x.ndim not in (1, 2)
|
|
or up_weight.ndim != 2
|
|
or gate_weight.ndim != 2
|
|
or (x.ndim == 2 and not 1 <= x.shape[0] <= 8)
|
|
or up_weight.shape != gate_weight.shape
|
|
or x.shape[-1] != up_weight.shape[1]
|
|
or x.shape[-1] % 8 != 0
|
|
or x.device != up_weight.device
|
|
or x.device != gate_weight.device
|
|
or not x.is_contiguous()
|
|
or not up_weight.is_contiguous()
|
|
or not gate_weight.is_contiguous()
|
|
# The fused kernel reads all streams as uint4; contiguous-but-offset
|
|
# views are routed to the unfused chain instead of failing.
|
|
or (x.data_ptr() & 15) != 0
|
|
or (up_weight.data_ptr() & 15) != 0
|
|
or (gate_weight.data_ptr() & 15) != 0
|
|
or not is_available("bf16_swiglu")
|
|
)
|
|
|
|
|
|
def swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
|
|
"""Apply the dense-MLP SwiGLU projection with a safe torch fallback.
|
|
|
|
``ASTRAI_SWIGLU=0`` and ``auto`` keep the unfused linear-backend chain;
|
|
``1`` forces the fused primitive for supported inputs. Auto will adopt
|
|
an M-banded rule mirroring the linear backend once end-to-end evidence
|
|
qualifies one.
|
|
"""
|
|
if env_mode("ASTRAI_SWIGLU") != "1" or not _swiglu_capable(
|
|
x, up_weight, gate_weight
|
|
):
|
|
return _unfused_swiglu(x, up_weight, gate_weight)
|
|
return _fused_swiglu(x, up_weight, gate_weight)
|
|
|
|
|
|
__all__ = ["swiglu"]
|