Where the last talk of the day to wrap up, and this is a fantastic thing, because more
it will tell about the actual application of AI and what he's doing with Ben Storin hardware
on drug discovery, and yeah, I've just learned like it was like a hobby project or it was
a hobby project last year, and now he's working at Ben Storin to continue his work,
so like hobby project turned into a job, and a good one with a passion, this is awesome,
so let's listen to what you want to say about the progress of your project.
Awesome, thanks.
So yeah, I'm going to give a little update on my project T-Dibolt.
I actually gave already a talk at the second edition of the AI,
I plan the conference about T-Dibolt, but that was the like Bolt's one on Tense on Warmhol,
and really a lot has happened since then, and so I thought it would be like cool to give
a little update, so for people who are not already familiar with the project, T-Dibolt
is basically a part of Bolt's two Tense-Dorin hardware, and Bolt's two is a biomelacular structure
prediction model, and basically just improved successor of Alpha Fold 3, and I think like
more people are familiar with Alpha Fold than Bolt's, but actually like Bolt's become industry
standard, and so all the biotech labs use Bolt's, and not Alpha Fold anymore.
And yeah, the project is completely open source, another MIT license, and it works out of
the box on Tense-Dorin hardware.
And I started the project in November, 24, just as my pet project as a side project during
the university, and then in April, 25 at the first working prototype, and since then I was just
optimizing the model, and after I finished my studies, I joined Tense-Dorin, and I get a lot of
time to work on the project and more resources, which is awesome, and I also wrote my thesis about
the topics, so if anyone is interested, you can also check out the report on the thesis.
Yeah, but first I want to talk a little bit about Tense-Dorin hardware,
I want to keep it really brief, because like there were so many talks about Tense-Dorin hardware.
So here we have the Tense-Dorin black hole P150, which is currently the most performant PCI
express card that is available in the top, and yeah, in the middle we have the black hole
processor, then 8, GTDR6, memory banks with 4G over each, so 32 gigabytes, but for me like
the coolest thing is that it has a large S-RAM, 200, 10 megabytes of S-RAM, and it's a scratch
bet, so we can control it explicitly, and stay for a long time in S-RAM for certain operations,
which can be a major advantage against like, or in comparison with GPUs who have like transparent
catches that you can only control implicitly on the right is my acquired box, and yeah,
the boss project right now is only running on a single card, but very soon we want to
polarize it, of course, multi-baccades, there's also a Tense-Dorin Galaxy server with like
32 processors, and yeah, I want to leverage everything for for drug discovery on the
boss project, yeah, here's the architecture of the black hole processor and the Tense-Dorin
Core, so black hole processor, the architecture is basically just a grid of pulse, we have
even at course DRAM cores system controller, but we're actually the compute heavens,
that's the Tense-Dorin Core, and on the right we have the architecture of Tense-Dorin Core, and we have
five small risk five chips, and they basically control the routers, and the data movement and
compute engines, and each Tense-Dorin Core has also, I think, on black hole 1.5 megabits of S-RAM,
yeah, but I think not everyone here is familiar with structure prediction,
so proteins are basically the bidding box of life, they end up in all processes,
in life, and they are defined by sequence of amino acids, and the sequence of amino acids
falls into a 3D structure in the human body, and the 3D structure determines the function
of what the protein does in the human body, so it is really important to predict the structure,
but it was a huge problem for 50 years and biology, and I thought for two basically solved
this problem and therefore was awarded the Nobel Prize in 2024, and it can predict those structures
in minutes and not years. So, both two builds on top of alpha-fold, but it adds also affinity
prediction there in the top right corner, and affinity prediction predicts basically the probability
and how strong a small molecule binds to a protein, and for example, a protein could be a harmful
protein, and then you want to design a drug so a small molecule, the ligand here, that attaches
to the protein, and maybe disables this function, and that's why binding affinity is so
important for drug discovery, and yeah, I'm contrast to alpha-fold, the bottom model is fully open
source and not only for academic usage, and the bottom way of the whole architecture of the
builds, two models it can see pretty complex, but the most important modules are really just
the performer in the middle, the diffusion model, and maybe also the MSA module, so formally
it is basically just diffusion module, but it is a pretty unique diffusion module because most
of the computational complexities in this pairformer, which is a transformer, and yeah, that's what
makes a different from pretty different from other diffusion modules. This is another way of
looking at the architecture, and there we can see that both two basically consists of a few
number of big modules, which are composed of smaller modules, and the main big modules are the
pairformer diffusion module, and the MSA module, and for example, see that the template module,
the confidence module, and the affinity module are just thin repas around the pairformer, and so
the task of plotting the builds model to transform hardware was just plotting those big modules
to the tensed-on architecture. And yeah, I would say the most interesting small modules are
trying a multiplication and triangle attention, and they have cubic complexity, not just
quadratic, complexity, and context, for example, to regular self-attention, and the triangle
comes from the fact that between those atoms, the triangle inequality holds, and that is just
inductive bias, and that's basically what the researchers did with alpha-fold and bolts. It feels
like engineering inductive biases, so they take regular machine learning operations, but they have
some requirements, for example, physics, and they have to engineer those inductive biases
in the model. And another interesting thing about the bolts to model are the activations. So
we have a single representation, this is a usual representation, for example, in transformers,
in transformers, we have a sequence of tokens, and therefore sequence of embeddings, and therefore
this matrix on the left, but additionally, in the bolts to model, we have a pair representation,
and this pair representation has additional dimension, which is as large as the sequence length,
and this is the root cause for this increased complexity. This means, for example, that the activations,
the memory complexity of the activations are quadratic and sequence length, and so for larger
molecules, most of the memories consume by the activations and not the weights.
The bolts to model is actually a pretty small model, but activations can get really large,
and this is also this additional dimension, which is as large as the sequence length,
is also the root cause for the increased for the cubic time complexity of triangular tension
and triangle multiplication. And yeah, that's basically this skeleton, how I
ported our bolts to 10 stone hardware. Those are just two of the big modules, but I just
rewrote them in TTNN, then put them in a PyTorch wrapper, so from the outside they look like
PyTorch modules, and then I could use it in the existing code base. And then yeah,
most of my time was actually spent on performance optimizations, and the key performance optimizations
that I did across all those modules was for triangle tension, I integrated flash attention to,
I will definitely go into that later. Then of course, like always, fusing operations, for example,
if you compute queries, keys, and values, you have three linear transformation, but you can
fuse them into one, then you have a bigger matrix multiplication, that means higher arithmetic
intensity, and so you can utilize the matrix units more, but in general, just fusing operations,
having less kernel overhead staying long and SRAM, there was important. And since the
activations have high memory complexity, I use chunking a lot, so processing those activations,
those tensors, and chunks that we also can keep in SRAM, that was really important. And of course,
mixed precision, so by default, we use BFLO16, but whenever we can, whenever it doesn't affect
the accuracy too much, we use block f, p8, and for example, for triangle attention, we can
basically do all of it and block f8. And because, yeah, because triangle attention is one of the
most interesting small modules, I just wanted to give briefly the idea how we used regular
flash attention to for that. And we have here the standard attention formula, and in the standard
attention formula, we add attention mass. And in triangle attention, we just have to replace the
attention mass with the triangle bias. And there's like, there's a small difference between the
attention mask and the triangle bias. The attention mask is the same for all attention heads.
So we can broadcast it across the head dimension. But the triangle bias is actually a different
per head, but we have to broadcast this across the bedstimension. And so the first trick that we did
was just permuting the head dimension and the bedstimension because they are both independent.
But even those permutations took too long. And so in the end, we just added support for
batch board, brush casking, to the flash attention color. And this way, we can just use
flash attention to for triangle attention. And no one can talk a little bit about the evaluation
about the performance. This is, of course, just the baseline. So both, both, two on a CPU
and both two on black hole. And yeah, we can see that it is significantly faster than a CPU
for all sequence length. And with sequence length, I always mean the number of amino acids
and are protein. More interesting is, of course, the comparison with a GPU. And so I wanted to
compare the best, the best PCR specouts from NVIDIA and from 10-strand, but I available for
consumers. But of course, just comparing the inference time would be a little unfair because
like the average over the last year of RTX 5090 or were $3,000, right now, it's not even available
for this price. It turns out significantly faster. And if we look at the predicted structures
per hour per dollar, we have better results for the 10-strand cards for all proteins except
the one with more than 1,000 amino acids. But I still think that there's a buck for large proteins.
And I hope we can solve it and also get better for larger proteins. So I think like 10-strand
cards already have very good price performance and for small proteins, 10-strand cards actually,
the 10-strand black hole cards actually faster than RTX 5090. And when optimizing the performance,
poor filings, of course, are very important and looking at how the runtime is distributed across
all sequence length. And I think the across all the big modules and I think the most interesting
thing is here that as the sequence length increases, the performer more and more dominates the
runtime. And that's because only the performer contains trying a lot of tension and trying a
multiplication, which are the only small modules that have cubic complexity. And so for large
proteins, we mostly care about the performer, but for small proteins, it's also very important
to optimize, for example, the diffusion module. And yeah and that's basically how a predicted
structure looks like. This protein structure, the blue protein structure was predicted on a
10-strand black hole card. In green, we have the experimental ground truth. And so yeah,
it is pretty accurate and that's also confirmed by all the important metrics. And I guess we
still have a little bit of time so I can also hopefully show a little demo.
Oh, maybe I have to start it again.
Yeah, so I think it works now. So the first part was generating the MSA. Then it is loading the
module into the memory. Then the performer runs and right now that is the diffusion module.
And you can see how the protein is folded in return. There was pretty fast. And then you can also
see, for example, for each amino acid here, how confidence the model is. And the prediction
all the over all confidence and some metrics that are important for biology.
So I think we're that we can open up for a Q&A maybe.
Yeah. On the comparison is the NVIDIA running on FP32 and TN-10-strand on BF16 or base.
Yeah, the NVIDIA cards also running on BF16, but they don't use block FPA anywhere.
Okay, good. Thank you very much. This is a wrap. I'll give it to you.
Thank you.
Okay, thank you everybody for joining us today. So if you want to continue this discussion,
you can join us also on Monday on the fringe event. You can also just continue to the breweries
and continue the evening. Thank you everyone.