b4e055e9a9
EngineX replaces the missing corex_gdn/corex_moe/corex_fa2 operator chain
that Sub168 has but our BI-V100 image lacks.
Architecture (mirrors CCCL dispatch/tuning/kernel three-layer system):
Registry (policy_selector) → three-tier dispatch:
Tier 1: Native .so via dlopen (libcorex_gdn.so, libixattn.so)
Tier 2: ixformer Python ops (vendor-provided)
Tier 3: PyTorch fallback (always available)
Critical fixes vs comp 168 docker log:
- moe_topk_softmax: replacement for missing ixformer op
- gdn_prefill: NaN-stable chunked impl (chunk_size=16)
- gdn_decode: state clamp prevents NaN accumulation
18 operators, all tests pass.
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)
|