Follow torch.sign in the engram gate

[truong-v]

Description

engram_gate_value_rms_kernel gates with tl.where(gate_raw >= 0, 1.0, -1.0), which has two outcomes, while EngramPt.forward in the same file gates with gate.sign(), which returns 0 for a zero input. With the clamp_min(1e-6) in front of the square root, the tie is not degenerate: at gate_raw == 0 the module gives sigmoid(0) = 0.5 and the kernel gives sigmoid(sqrt(1e-6)) = 0.50025. A zero gate_raw is what a masked or padded position produces, since its query row is all zeros.

Changes:

• Use the three-way sign in the kernel, so a zero stays zero.

Measured on an L40S with both modules at the same weights and one masked position (kernel_size=1, conv weight 1.0, so the fused path's half buffers are exact and only the gate can differ): masked rows had max abs error 6.85e-4 against ordinary rows at 6.20e-6 before this change, and 0.0 against 6.20e-6 after. Performance is unchanged — the sign is computed once per program, outside the channel loop: gate kernel at B=8, T=256, G=4, C=1024, do_bench interleaved three times, 177.5–177.7 us before and 177.1–177.9 us after.

Tests

The op had no test file. Added tests/kernels/test_engram.py, which compares EngramTri against EngramPt on masked query rows in the shape EngramPt.forward supports. It fails on the unfixed kernel (4838 of 5120 masked elements mismatched, max 6.85e-4 against a 1e-5 tolerance) and passes with the fix. The existing in-file test_engram() passes unchanged.

Fixes #176

Environment

• based on DLBlas main @ 26271d1
• NVIDIA L40S, driver CUDA 12.4
• torch 2.6.0+cu124, triton 3.2.0, Python 3.12