engram_gate_value_rms_kernel gates a zero as positive, where the eager module uses torch.sign

[truong-v]

Summary

engram_gate_value_rms_kernel gates with a two-way sign (dlblas/kernels/engram.py:203):

gate = gate_sqrt * tl.where(gate_raw >= 0, 1.0, -1.0)

The eager module in the same file, EngramPt.forward (engram.py:71), gates with torch.sign, which has three outcomes:

gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign()

torch.sign returns 0 for a zero input, and the clamp_min(1e-6) is what keeps that case from being degenerate: at gate_raw == 0 the module gives sigmoid(0) = 0.5, while the kernel takes the positive branch and gives sigmoid(sqrt(1e-6)) = 0.50025. Every value on that row is then scaled by a gate that is 5.0e-4 too large.

gate_raw is exactly zero whenever the query row is all zeros — which is what a masked or padded position looks like — or whenever key and query are orthogonal. The in-file test compares the two modules on random data at rtol=atol=1e-3, and random data never lands on the tie, which is why this has not shown up.

The same expression is copied into dlblas/kernels/engram_demo_v1.py by the open PR #106, so a fix probably wants to cover that file too.

Reproduction

Both modules at the same weights, with one masked position. kernel_size=1 and a conv weight of 1.0 keep the fused path's half-precision buffers exact, so the gate is the only thing that can differ:

import torch
from dlblas.kernels.engram import EngramPt, EngramTri

def build(cls):
    torch.manual_seed(7)
    m = cls(1024, 1024, kernel_size=1, dilation=1, hc_mult=4, activation=False)
    with torch.no_grad():
        m.conv.weight.fill_(1.0)
    if hasattr(m, "_precompute_buffers"):
        m._precompute_buffers()
    return m

tri, pt = build(EngramTri), build(EngramPt)
torch.manual_seed(0)
emb = torch.randn(1, 6, 1024, device="cuda")
hs = torch.randn(1, 6, 4, 1024, device="cuda")
hs[0, 3] = 0.0                        # a masked / padded position: zero query row
d = (tri(embeddings=emb, hidden_states=hs) - pt(embeddings=emb, hidden_states=hs)).abs()
masked = hs.abs().sum(-1) == 0
print(f"masked rows: max |diff| {d[masked].max():.2e}")
print(f"other rows : max |diff| {d[~masked].max():.2e}")

masked rows: max |diff| 6.85e-04
other rows : max |diff| 6.20e-06

The masked rows carry about 110 times the error of the ordinary rows, and it is systematic rather than round-off: reading the gate the kernel applied gives 0.50024998 against the module's 0.50000000 on every masked row.

If the two-way sign is the intended behaviour, then the eager module is the side that should change — but the two paths should agree.

Environment

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