nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

gptoss: 3d mx block scale (#17062)

Browse Source
This commit is contained in:
wozeparrot
2026-07-16 23:30:24 -07:00
committed by GitHub
parent 6f1176ea90
commit a836c3822a
2 changed files with 11 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/mlperf/models/gpt_oss.py
+6 -3
View File
@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
from tinygrad.uop.ops import Ops, UOp
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
from extra.llama_kernels.rmsnorm import rmsnorm
from extra.gemm.cdna_asm_gemm import _mx_block_scale, quantize_mxfp8
from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
FP8_DTYPE = dtypes.fp8e4m3
FP8_MAX = 448.0
@@ -39,8 +39,11 @@ def quant_dequant_mx(x:Tensor) -> Tensor:
fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
def _mx_scale(e8:Tensor) -> Tensor:
return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
return w_q.cast(dtypes.bfloat16) * _mx_block_scale(w_scale)
return w_q.cast(dtypes.bfloat16) * _mx_scale(w_scale)
@functools.cache
def _dequant_fwd_fxn(wq_p, ws_p, device):
@@ -48,7 +51,7 @@ def _dequant_fwd_fxn(wq_p, ws_p, device):
def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
w_scale = Tensor(call.src[2])
return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_block_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
extra/gemm/cdna_asm_gemm.py
+5
View File
@@ -128,6 +128,11 @@ def _mx_block_scale(e8:Tensor) -> Tensor:
rows, scale_K = e8.shape
return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(rows, scale_K, 1).expand(rows, scale_K, 32).reshape(rows, scale_K*32)
def _mx_block_scale_3d(e8:Tensor) -> Tensor:
# batched (E, rows, scale_K) dequant scale 2^(e8-127) broadcast to (E, rows, scale_K*32)
E, rows, scale_K = e8.shape
return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(E, rows, scale_K, 1).expand(E, rows, scale_K, 32).reshape(E, rows, scale_K*32)
counters = {"used":0, "todos":[]}
def todo(msg:str) -> bool: counters["todos"].append(msg); return False
def _asm_gemm_report():