How to control computation precision

[999cube]

How to control computation precision in FP16, BF16, FP8 etc.
Do I have no choice but to build the cuda kernel?

[999cube]

I'm trying to evaluate the model's accuracy when all internal operations are strictly limited to FP16. However, I noticed that the selective_scan CUDA kernel seems to use FP32 accumulators by default.

When I simulated the FP16 truncation in Python, I saw a 0.04% accuracy drop. Now I want to replicate this at the CUDA kernel level, but I'm having trouble modifying the C++ source without breaking dependencies.

Does anyone know if there is a Triton-based implementation of Mamba? Or is there a standard way to control the internal precision of these fused kernels for research purposes?

Any advice would be appreciated. Thanks!