mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-06-24 02:14:17 +00:00
fp8 gemm cleanup (#16607)
This commit is contained in:
parent
bcdb988df0
commit
2e77bd01db
1 changed files with 3 additions and 8 deletions
|
|
@ -346,20 +346,15 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_pt
|
|||
}
|
||||
|
||||
// apply x_scale * w_scale before bf16 store to prevent overflow
|
||||
#if SCALE_MODE != 0
|
||||
#if SCALE_MODE == 1
|
||||
float scale = *x_scale_ptr;
|
||||
mul(cA, cA, scale);
|
||||
mul(cB, cB, scale);
|
||||
mul(cC, cC, scale);
|
||||
mul(cD, cD, scale);
|
||||
#elif SCALE_MODE == 2
|
||||
float scale = *w_scale_ptr;
|
||||
mul(cA, cA, scale);
|
||||
mul(cB, cB, scale);
|
||||
mul(cC, cC, scale);
|
||||
mul(cD, cD, scale);
|
||||
#elif SCALE_MODE == 3
|
||||
float scale = *x_scale_ptr * *w_scale_ptr;
|
||||
#endif
|
||||
|
||||
mul(cA, cA, scale);
|
||||
mul(cB, cB, scale);
|
||||
mul(cC, cC, scale);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue