| 1 | # ruff: noqa: ANN001, ANN201, ERA001, N803, N806 |
| 2 | import triton |
| 3 | import triton.language as tl |
| 4 | |
| 5 | |
| 6 | @triton.jit |
| 7 | def fused_add_round_kernel( |
| 8 | x_ptr, |
| 9 | output_ptr, # contents will be added to the output |
| 10 | seed, |
| 11 | n_elements, |
| 12 | EXPONENT_BIAS, |
| 13 | MANTISSA_BITS, |
| 14 | BLOCK_SIZE: tl.constexpr, |
| 15 | ): |
| 16 | """ |
| 17 | A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding |
| 18 | and add them to bfloat16 output weights. Might be used to upcast original model weights |
| 19 | and to further add them to precalculated deltas coming from LoRAs. |
| 20 | """ |
| 21 | # Get program ID and compute offsets |
| 22 | pid = tl.program_id(axis=0) |
| 23 | block_start = pid * BLOCK_SIZE |
| 24 | offsets = block_start + tl.arange(0, BLOCK_SIZE) |
| 25 | mask = offsets < n_elements |
| 26 | |
| 27 | # Load data |
| 28 | x = tl.load(x_ptr + offsets, mask=mask) |
| 29 | rand_vals = tl.rand(seed, offsets) - 0.5 |
| 30 | |
| 31 | x = tl.cast(x, tl.float16) |
| 32 | delta = tl.load(output_ptr + offsets, mask=mask) |
| 33 | delta = tl.cast(delta, tl.float16) |
| 34 | x = x + delta |
| 35 | |
| 36 | x_bits = tl.cast(x, tl.int16, bitcast=True) |
| 37 | |
| 38 | # Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for |
| 39 | # normal numbers and -14 for subnormals. |
| 40 | fp16_exponent_bits = (x_bits & 0x7C00) >> 10 |
| 41 | fp16_normals = fp16_exponent_bits > 0 |
| 42 | fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14) |
| 43 | |
| 44 | # Add the target dtype's exponent bias and clamp to the target dtype's exponent range. |
| 45 | exponent = fp16_exponent + EXPONENT_BIAS |
| 46 | MAX_EXPONENT = 2 * EXPONENT_BIAS + 1 |
| 47 | exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent) |
| 48 | exponent = tl.where(exponent < 0, 0, exponent) |
| 49 | |
| 50 | # Normal ULP exponent, expressed as an fp16 exponent field: |
| 51 | # (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15 |
| 52 | # Simplifies to: fp16_exponent - MANTISSA_BITS + 15 |
| 53 | # See https://en.wikipedia.org/wiki/Unit_in_the_last_place |
| 54 | eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15)) |
| 55 | |
| 56 | # Calculate epsilon in the target dtype |
| 57 | eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True) |
| 58 | |
| 59 | # Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) -> |
| 60 | # fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 = |
| 61 | # 16 - EXPONENT_BIAS - MANTISSA_BITS |
| 62 | eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True) |
| 63 | eps = tl.where(exponent > 0, eps_normal, eps_subnormal) |
| 64 | |
| 65 | # Apply zero mask to epsilon |
| 66 | eps = tl.where(x == 0, 0.0, eps) |
| 67 | |
| 68 | # Apply stochastic rounding |
| 69 | output = tl.cast(x + rand_vals * eps, tl.bfloat16) |
| 70 | |
| 71 | # Store the result |
| 72 | tl.store(output_ptr + offsets, output, mask=mask) |
| 73 |