[INTERPRETER] Compute bf16 ops in a value dtype#10945
Conversation
|
These are two different problems. int1 is handled in #10923 bfloat16 I'm not sure why we need an upcast. How does it match the GPU side? |
The interpreter stores bf16 as a uint16 bit pattern, but binary_op, unary_op, fma and dot ran the numpy op directly on that storage, so bf16 arithmetic and comparisons computed on the bit pattern instead of the value, diverging from the GPU. Convert to float32 before the op (as the GPU does) and round back to bf16 after. Adds interpreter tests for bf16 binary_op, comparison, neg and fma. (int1, the other half of triton-lang#10919, was fixed separately in triton-lang#10923.)
545ed87 to
9faa170
Compare
|
Thanks @Jokeren. Dropped the int1 changes (handled by #10923) and rebased; this is now bf16-only. The bug: bf16 is stored as a On the upcast to fp32: numpy has no native bf16 (smallest float is fp16), so the op has to run in some numpy float type. fp32 holds bf16 losslessly, and a single RTNE round back to bf16 is bit-identical to a native bf16 op (verified for add/sub/mul), so this changes representation only, not numerics. Same as the int1 branch in #10923 (compute in a wider domain, narrow back). Tests fail-before/pass-after on Same-type fp8 is out of scope here. It is also uint8-stored and semantic.py#L106 keeps the dtype rather than promoting, so the interpreter has the same bug ( |
Refs #10919 (bf16 half; int1 half fixed in #10923).
The interpreter stores bf16 as a
uint16bit pattern, butbinary_op,unary_op,create_fmaandcreate_dotran the numpy op directly on that storage, so bf16 arithmetic and comparisons computed on the bit pattern instead of the value, diverging from the GPU.2.5 + 2.0gave-2.9e-39(from16416 + 16384as uint16) instead of4.5;-5.0 < -2.0gave0instead of1(bf16 is sign-magnitude, so uint16 order != value order).This converts operands to float32 before the op (as the GPU does) and rounds the result back to bf16 (RTNE) after, covering
binary_op,unary_opandcreate_fma;create_dotnow converts any sub-32-bit float operand, not just 8-bit fp8. Adds interpreter tests for bf16binary_op, comparison,negandfma.