@@ -542,32 +542,6 @@ def _fused_gemm_rmsbwd_ybwd_kernel(
542542 p2 = tl .load (pre_ptr + pid * 4 + 2 )
543543 p3 = tl .load (pre_ptr + pid * 4 + 3 )
544544
545- # Load grad_mixes (24 scalars) — reused across all D-blocks
546- gm00 = tl .load (grad_mixes_ptr + gm_base + 0 )
547- gm01 = tl .load (grad_mixes_ptr + gm_base + 1 )
548- gm02 = tl .load (grad_mixes_ptr + gm_base + 2 )
549- gm03 = tl .load (grad_mixes_ptr + gm_base + 3 )
550- gm04 = tl .load (grad_mixes_ptr + gm_base + 4 )
551- gm05 = tl .load (grad_mixes_ptr + gm_base + 5 )
552- gm06 = tl .load (grad_mixes_ptr + gm_base + 6 )
553- gm07 = tl .load (grad_mixes_ptr + gm_base + 7 )
554- gm08 = tl .load (grad_mixes_ptr + gm_base + 8 )
555- gm09 = tl .load (grad_mixes_ptr + gm_base + 9 )
556- gm10 = tl .load (grad_mixes_ptr + gm_base + 10 )
557- gm11 = tl .load (grad_mixes_ptr + gm_base + 11 )
558- gm12 = tl .load (grad_mixes_ptr + gm_base + 12 )
559- gm13 = tl .load (grad_mixes_ptr + gm_base + 13 )
560- gm14 = tl .load (grad_mixes_ptr + gm_base + 14 )
561- gm15 = tl .load (grad_mixes_ptr + gm_base + 15 )
562- gm16 = tl .load (grad_mixes_ptr + gm_base + 16 )
563- gm17 = tl .load (grad_mixes_ptr + gm_base + 17 )
564- gm18 = tl .load (grad_mixes_ptr + gm_base + 18 )
565- gm19 = tl .load (grad_mixes_ptr + gm_base + 19 )
566- gm20 = tl .load (grad_mixes_ptr + gm_base + 20 )
567- gm21 = tl .load (grad_mixes_ptr + gm_base + 21 )
568- gm22 = tl .load (grad_mixes_ptr + gm_base + 22 )
569- gm23 = tl .load (grad_mixes_ptr + gm_base + 23 )
570-
571545 # Pass 1: accumulate dot(x, grad_x_scaled) and grad_pre
572546 dot = 0.0
573547 gp0 = tl .zeros ([], dtype = tl .float32 )
@@ -698,7 +672,6 @@ def _small_stage_bwd_kernel_hc4(
698672 sb = pid * ITERS * 16
699673
700674 # ---- FORWARD RECOMPUTE (pre / post / logits / sinkhorn) ----
701- m_pre = tl .arange (0 , 4 )
702675 # Load per-token mixes/base slices
703676 mix_pre0 = tl .load (mixes_ptr + mb + 0 )
704677 mix_pre1 = tl .load (mixes_ptr + mb + 1 )
0 commit comments