[Bug] Non-FP8 `general_gemm` silently misreads a bias whose dtype differs from the output dtype and reads out of bounds

[MarkXCloud]

Describe the bug

For a non-FP8/FP4 GEMM, cublaslt_gemm.cu never sets CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE; it is only set inside the if (use_fp8 || use_fp4) branch (main, around L476–L646). cuBLASLt then assumes the bias has the same type as D. Nothing checks this, so a call like

general_gemm(weight_bf16, inp_bf16, torch.float32, layout="TN", bias=bias_bf16)

returns wrong results without any error:

• the BF16 bias bytes are reinterpreted as FP32, so the first half of the output columns get garbage bias values;
• the epilogue reads 2 * bias.nbytes, i.e. past the end of the bias tensor, so the second half of the columns pick up whatever memory follows it.

The Python side already computes bias_dtype = TE_DType[bias.dtype] and passes it down, which makes the mixed-dtype call look supported.

Impact in practice

Megatron-Core's MoE router (RouterGatingLinearFunction) called te_general_gemm(weight, inp, router_dtype, layout="TN", bias=bias) with router_dtype=torch.float32 and a BF16 router.bias. This is the default setup for gpt-oss (router with bias, moe_router_dtype=fp32, BF16 params). With parameters held as views in the DDP param buffer, the overflow read the neighbouring parameter (a layernorm weight ≈ 1.0), adding a fake +2..+8 bias to half of the experts. In RL training, this showed up as a 1.3–1.5 nat train/inference logprob mismatch on gpt-oss-120b (0.002 after the fix). Megatron-LM has since worked around it on main by casting the bias (#6000, commit 883533e049). Any other caller passing a bias with a different dtype than the output is still exposed.

Steps/Code to reproduce bug

import torch
from transformer_engine.pytorch.cpp_extensions import general_gemm

torch.manual_seed(0)
E, H, T = 128, 2880, 256
buf = torch.randn(4 * E, dtype=torch.bfloat16, device="cuda")
buf[E:] = 1.0                      # data that happens to follow the bias in memory
bias = buf[:E]
weight = torch.randn(E, H, dtype=torch.bfloat16, device="cuda")
inp = torch.randn(T, H, dtype=torch.bfloat16, device="cuda")

ref = torch.addmm(bias.float(), inp.float(), weight.float().t())
gemm = lambda b: general_gemm(weight, inp, torch.float32, layout="TN", bias=b)[0]
no_bias, out_bf16, out_fp32 = gemm(None), gemm(bias), gemm(bias.float())
effective = (out_bf16 - no_bias).mean(0)

print((out_bf16 - ref).abs().max())                                   # 4.1272  <- wrong
print((out_fp32 - ref).abs().max())                                   # 0.0004  <- correct
print((effective - buf.view(torch.float32)[:E]).abs().max())          # 9.8e-07: bf16 bytes read as fp32
print(effective[E // 2:].mean())                                      # 1.002: bytes beyond the bias tensor

Output (torch 2.11.0+cu130, TE 2.15.0+42b84005, GB200):

max|out - ref|, bf16 bias        : 4.1272
max|out - ref|, bias cast to fp32: 0.0004
effective bias vs bf16 bytes reinterpreted as fp32: max abs diff 9.83e-07
effective bias, columns 64..127 (bytes beyond the bias tensor): mean 1.002

Expected behavior

Either

1. set CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE in the non-FP8 path as well (if cuBLASLt supports the combination), or
2. reject the call with an NVTE_CHECK (bias dtype must equal D dtype when not using FP8/FP4), or cast the bias in the PyTorch wrapper.

Silently producing wrong values and reading out of bounds is the problem; any of these would prevent it.

Environment overview

• Environment location: Slurm cluster, container
• Method of Transformer Engine install: prebuilt wheel in the NeMo-RL container
• The code path is unchanged on TE main (checked at f5f4e83, 2026-09-22).

Environment details

• OS version: Ubuntu (container), aarch64
• PyTorch version: 2.11.0+cu130
• Python version: 3.13
• Transformer Engine version: 2.15.0+42b84005
• CUDA version: 13.0

Device details

• GPU model: NVIDIA GB200

[dundysm]

i'd like to work on this.

root cause: non-fp8/fp4 cublaslt_gemm sets bias pointer/epilogue but not CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE (only set in the fp8/fp4 branch). cublas then assumes bias dtype == D, so a bf16 bias with fp32 out is reinterpreted and read past the end of the tensor. te modules usually cast bias to activation dtype first; general_gemm does not.

proposed fix: NVTE_CHECK that bias dtype equals D dtype on the non-fp8/fp4 path (matches cublas support tables and the megatron #6000 workaround). alternatively cast in the python wrapper. setting BIAS_DATA_TYPE alone does not make bf16 bias + fp32 D valid for non-fp8 per cublas 12.8 tables.

not related to #3248.