auto-round
f96ec416 - fix: BF16 sparse native kernel element-aware MMA and softmax scale

Commit
9 days ago
fix: BF16 sparse native kernel element-aware MMA and softmax scale Two device-level bugs made the native BF16 sparse kernel unusable: - SparseSageConfig hardcoded the Q*K MMA to an int8 DPAS (XE_DPAS_TT<..., int32_t, int8_t>) regardless of ElementQ. Feeding bf16 fragments into the int8 DPAS hung the device; INT8 worked only because its fragments genuinely are int8. The MMA is now ElementQ-aware: int8 keeps the int32/int8 DPAS, while bf16/fp16 use a float/<ElementQ> DPAS. - The mainloop applied the softmax scale only inside the scale_block_size != 0 branches. The native BF16 path uses scale_block_size 0 (no quantization), so the 1/sqrt(d) scale was never applied and the softmax ran on unscaled scores. Added the else branch that multiplies by params.scale.
Author
Parents
Loading