40 lines
1.1 KiB
Python
40 lines
1.1 KiB
Python
|
"""
|
|||
|
|
EngineX activation operators.
|
|||
|
|
|
|||
|
|
These map to CCCL's dispatch_transform pattern — element-wise kernels
|
|||
|
|
that fuse activation + multiply in a single pass.
|
|||
|
|
|
|||
|
|
ixformer provides these natively (confirmed working in hardware probe).
|
|||
|
|
PyTorch fallbacks here for completeness.
|
|||
|
|
|
|||
|
|
CCCL tuning: tuning_transform.cuh bytes_in_flight = 64KB on BI-V100
|
|||
|
|
(56 GB/s per-SM × 1100ns latency, 16 SMs)
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
|
|||
|
|
|
|||
|
|
def silu_and_mul_pytorch(x: torch.Tensor, out: torch.Tensor) -> None:
|
|||
|
|
"""Fused SiLU(x[..., :d]) * x[..., d:]"""
|
|||
|
|
d = x.shape[-1] // 2
|
|||
|
|
gate = x[..., :d]
|
|||
|
|
up = x[..., d:]
|
|||
|
|
out.copy_(F.silu(gate) * up)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def gelu_and_mul_pytorch(x: torch.Tensor, out: torch.Tensor) -> None:
|
|||
|
|
"""Fused GELU(x[..., :d]) * x[..., d:]"""
|
|||
|
|
d = x.shape[-1] // 2
|
|||
|
|
gate = x[..., :d]
|
|||
|
|
up = x[..., d:]
|
|||
|
|
out.copy_(F.gelu(gate) * up)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def gelu_tanh_and_mul_pytorch(x: torch.Tensor, out: torch.Tensor) -> None:
|
|||
|
|
"""Fused GELU_tanh(x[..., :d]) * x[..., d:]"""
|
|||
|
|
d = x.shape[-1] // 2
|
|||
|
|
gate = x[..., :d]
|
|||
|
|
up = x[..., d:]
|
|||
|
|
out.copy_(F.gelu(gate, approximate='tanh') * up)
|