sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

gptoss: route lm_head thru asm_gemm (#17497)

Browse Source
This commit is contained in:
wozeparrot
2026-08-11 18:32:50 -07:00
committed by GitHub
parent 3686a1758f
commit f891f5ffd0
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/mlperf/models/gpt_oss.py
+6 -2
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
from extra.llama_kernels.rmsnorm import rmsnorm
from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
from extra.gemm.moe_gemm import grouped_mx_gemm
from extra.gemm.moe_routing import route, dispatch, combine
@@ -263,7 +263,11 @@ class GPTOSS:
w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
logits = self.norm(h) @ self.output.T
h_normed = self.norm(h)
pad = (-self.dim) % 256
h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
else: logits = h_normed @ self.output.T
return logits
def _get_pads(uop:UOp) -> list[UOp]: