flux-tinygrad-opt

Objective

Try optimizing Blackforest's Flux.1 with tinygrad's optimization, especially on H100.

Flux model class is based on PR #6334, which can seamlessly load weights of official one, downloaded from huggingface.

Optim. methods include

• RAW: The raw code.
• HC: Hand-coded; usual optim. such as dead code elimination, loop unrolling, etc.
• TC: Seems like NVIDIA's tensor cores optim. Only applied when applicable, else raw code.

---

• BEAM: Beam search the kernel.
• MCTS: MCTS the kernel.
• [TODO-RL]: Utilize RL like AlphDev did for CPU.

[TODO-1] Reverse-engineer optimized IR back to high-level tinygrad code and benchmark with "torch compiled" original Flux, measured by torch profiler or torch.cuda.Event. Also, refer to this.

Verification status

There is currently no reproducible end-to-end Flux performance result in this repository. The historical run did not preserve raw per-kernel records, an executable combined schedule, numerical-equivalence checks, or exact dependency revisions. Current tinygrad no longer exposes the imported APIs.

The closest reconstructable upstream snapshot is tinygrad commit 8074c0ec8f1e3ffc5a9459294570d778b5d1fa4f, the last upstream commit before this repository's final 2024 code change. With that full source checkout, scripts/check_environment.py imports the experiment and creates an 18-item dummy schedule. That validates the code's historical API environment only; it does not validate the performance numbers below.

Historical result — do not cite as a model benchmark

Benchmark

Exploratory oracle result: 592.40 ms and 60,225 weighted GFLOPS. This is the sum of independently timed per-kernel minima, not executable end-to-end Flux inference latency.

For each kernel, the best time across all optimization methods is selected. No single optimized schedule is assembled and executed, and numerical equivalence of the selected programs is not checked. Treat the total as an optimistic lower bound for follow-up work, not a model benchmark. For example, in the table below, BEAM (594.451 ms) is chosen for kernel_X.

Flux's flow is broken down into 1,423 tinygrad kernels.

---

Additional results Performance when using a single optimization method exclusively.

Optimization Method	Time (ms)
RAW	201654.093
HC	7989.641
TC	6317.517
BEAM	594.451
MCTS	825.108

(The raw result significantly differs from the values from Peter and Antonio, check.) => the time seems to be correct. tinygrad/tinygrad#6759

Caveat Couldn't solve TODO-1, therefore, the times(ms) and GFLOPS are measured within tinygrad's method. The cumulative table's work column is GFLOP (operation count), while the final weighted rate is GFLOPS. [TODO-2] Does tinygrad lookup hardware intrinsics? or actually measure the time? The runtime dispatcher returns execution time when called?

The original environment did not record exact dependency revisions. The code uses tinygrad internals from the PR #6334 era, so current unpinned dependencies may require adaptation before the experiment can be reproduced.

Reconstruct the historical code environment

git clone https://github.com/tinygrad/tinygrad.git .deps/tinygrad
git -C .deps/tinygrad checkout 8074c0ec8f1e3ffc5a9459294570d778b5d1fa4f
python -m pip install -r requirements.txt
PYTHONPATH="$PWD/.deps/tinygrad" python scripts/check_environment.py

main.py now constructs the Flux topology without downloading model weights by default. This is suitable for schedule-development smoke tests, but random unrealized parameters can add initialization work and must not be compared with the historical timing table. A new publishable benchmark needs to persist raw kernel records, assemble one executable schedule, verify outputs against a reference implementation, and record hardware/software provenance.

Prerequisites

pip install -r requirements.txt

apt install graphviz git-lfs

Example outputs

An optimized kernel

Block size-tuned kernel

#define INFINITY (__int_as_float(0x7f800000))
#define NAN (__int_as_float(0x7fffffff))
extern "C" __global__ void __launch_bounds__(64) r_2048_4_16_16(float* data0, const float* data1) {
  __shared__ float temp1[64];
  int gidx0 = blockIdx.x; /* 2048 */
  int lidx0 = threadIdx.x; /* 4 */
  int lidx1 = threadIdx.y; /* 16 */
  int alu0 = (lidx0*16);
  float acc0 = 0.0f;
  for (int ridx0 = 0; ridx0 < 16; ridx0++) {
    float val0 = data1[(gidx0*1024)+(lidx0*256)+lidx1+(ridx0*16)];
    acc0 = (acc0+val0);
  }
  temp1[alu0+lidx1] = acc0;
  __syncthreads();
  if (((bool)(lidx1)!=1)) {
    float acc1 = 0.0f;
    for (int ridx1 = 0; ridx1 < 16; ridx1++) {
      float val1 = temp1[alu0+ridx1];
      acc1 = (acc1+val1);
    }
    data0[(gidx0*4)+lidx0] = acc1;
  }
}

#define INFINITY (__int_as_float(0x7f800000))
#define NAN (__int_as_float(0x7fffffff))
extern "C" __global__ void __launch_bounds__(128) r_16_256_2_64_4(float* data0, const float* data1, const float* data2) {
  __shared__ float temp1[128];
  int gidx0 = blockIdx.x; /* 256 */
  int gidx1 = blockIdx.y; /* 16 */
  int lidx0 = threadIdx.x; /* 2 */
  int lidx1 = threadIdx.y; /* 64 */
  int alu0 = (lidx0*64);
  float val0 = data2[(gidx1*2)+lidx0];
  float acc0 = 0.0f;
  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
    float val1 = data1[(gidx1*131072)+(gidx0*256)+(lidx0*65536)+(lidx1*4)+ridx0];
    float alu1 = (val1+(val0*(-1.0f)));
    acc0 = (acc0+(alu1*alu1));
  }
  temp1[alu0+lidx1] = acc0;
  __syncthreads();
  if (((bool)(lidx1)!=1)) {
    float acc1 = 0.0f;
    for (int ridx1 = 0; ridx1 < 64; ridx1++) {
      float val2 = temp1[alu0+ridx1];
      acc1 = (acc1+val2);
    }
    data0[(gidx1*512)+gidx0+(lidx0*256)] = acc1;
  }
}

Metrics by kernel

             kernel 818 r_2560_21504_3072                     [21504, 2560, 1]   [1, 1, 1]    takes 1862.52 ms,    182 GFLOPS -- RAW                                                                                                    
             kernel 818 r_80_448_8_16_768_4_3_4               [448, 80, 1]       [8, 16, 1]   takes   76.35 ms,   4431 GFLOPS -- HC                                                                                                     
             kernel 818 r_32_168_2_2_2_2_2_4_192_8_2_2_2_5_4  [168, 32, 1]       [16, 2, 4]   takes    8.21 ms,  41219 GFLOPS -- TC                                                                                                     
             kernel 818 r_192_40_2_2_2_2_2_192_8_2_2_2_4_7_2  [40, 192, 1]       [8, 2, 2]    takes    5.16 ms,  65530 GFLOPS -- BEAM                                                                                                   
             kernel 818 r_32_384_2_2_2_2_2_192_8_2_2_2_7_5    [384, 32, 1]       [8, 2, 2]    takes    5.77 ms,  58602 GFLOPS -- MCTS 

e.g. r_80_448_8_16_768_4_3_4, r: reduce, 80_448: global size 8_16: local size. [TODO-3] Clarify the output. Is this OpenCL context?

r is the reduction operations like sum, dot product.

Related code

@functools.cached_property
def name(self) -> str:
  # kernel name (before late upcast)
  name = ("r" if self.reduceop else ("C" if all(x.op in BUFFER_UOPS for x in self.ast.parents) else "E")) + \
               (f"{len(self.ast.src)}_" if len(self.ast.src) > 1 else "_") + \
               colored('_', 'BLACK').join([colored(str(x), c) for x,c in zip(self.full_shape, self.colors())])

  # name the function something unique
  Kernel.kernel_cnt[(function_name := to_function_name(name))] += 1
  suffix = f"{'n'+str(Kernel.kernel_cnt[function_name]-1)}" if Kernel.kernel_cnt[function_name] > 1 else ""
return name+colored(suffix, 'BLACK')

print(f"                 kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS -- {colored(nm, 'green') if lin is sorted_choices[0][2] else nm}")

 @property
  def global_dims(self) -> int: return self.first_reduce-self.local_dims

  # there's eight chunks of the shape
  # blue   -- global dims
  # cyan   -- local dims (warp ones first)
  #  *** self.first_reduce
  # green  -- reduce-local dims
  # white  -- reduce-late upcasted dim (self.upcast_in_mid_reduce_axes)
  # red    -- reduce loops
  #  *** self.upcasted
  # purple -- reduce upcasted
  # yellow -- normal upcasted dimensions
  def colors(self) -> List[str]:
    # first non local non reduce dims are global (blue)
    colors = ["blue"] * self.global_dims if not self.dont_use_locals else ["BLUE"] * self.global_dims
    # after global are local_dims; warp ones used in tensor cores must be closest to first_reduce (cyan)
    colors += ["cyan"] * self.local_dims
    # between first_reduce and first_reduce + group_for_reduces, they are either upcast mid reduce (white), or late upcasted (green)
    colors += ["white" if i in self.upcast_in_mid_reduce_axes else "green" for i in range(self.first_reduce, self.first_reduce + self.group_for_reduces)]  # noqa: E501
    # between first_reduce + group_for_reduces and upcasted, they are reduce (red)
    colors += ["red"] * (self.first_upcast - (self.first_reduce + self.group_for_reduces))
    # upcasted dimensions are reduce (magenta) or normal (yellow)
    colors += ["magenta" if self.full_shape[i] != self.sts[0].shape[i] else "yellow" for i in range(self.first_upcast, self.shape_len)]
    assert len(colors) == self.shape_len, "colors size mismatch"
    return colors

Background

Supported ops

https://github.com/tinygrad/tinygrad/blob/d1bae42d3527e265b2772e39e563cdbaea34592e/tinygrad/ops.py