The llama journey

30 Dec, 2025

Day 1

12 December 2025 (~2 hr)

Today's Target: Build it myself

• First command gave error, Could Not find CURL
• found it on build.md, need to install libcurl-devel
• Now it works, running actual build. Very slow 🥲
• While that is working some good flags will be -j, lemme just stop and recompile for speed
• That's 8X fast, now I will add cuda flags next. No updates on dnf, all up to date
• Ok compiled the basic version, now cuda
• Building it was slow so I put up performance mode and flushed my battery, had to rm -rf build and start again
• DO NOT RUN PERFORMACE ON BATTERY
• everything compiled 🥳
• Lets test run it, using ./examples/simple, but wait how toooo,
• Why is everyone like download a model, no one is using the example folder or the models folder 😭
• All models in models folder fail to load 😭
• Ok now I'll try to download a model and then run it, given RTX 4050 I have quite some limitation as all the systems usually require 8GB VRAM.
• Downloading a small model DeepSeek-R1-Distill-Qwen-1.5B-Q4_K_M.gguf, and it runs 🥳🥳🥳. It is dumb though, but smart too, I knows how to solvemath so good!

Day 2

13 December 2025 (30 mins)

Today's Target: Try to understand maybe simple example

• looks like ./build/bin/llama-simple is the executable
• the print usage function is helpful let me just run it
• Just a simple run gave a lot of output, need to unravel slowly
• Ha this is funny, there was a little extension of my line, full on halucinating
• ahh so easy to understand near line 14-21 some simple initializations. It offloads layers to GPU itself. Ah that n_predict decides how much the line to extend.

  // path to the model gguf file
  std::string model_path;
  // prompt to generate text from
  std::string prompt = "Hello my name is";
  // number of layers to offload to the GPU
  int ngl = 99;
  // number of tokens to predict
  int n_predict = 32;
  

• Then a code block to parse CLI, lets see with just some error checking etc
  • If -m is found update model_path
  • If -n is found update n_predict
  • If -ngl is found update ngl
• Then everything else is just the prompt so no need of " in the prompt, lets try that ... and still as dumb
• Now there is load dynamic backend, need to see this later. Added to TODO
• Then there is model initialization, it gets default parameters, let's dig a bit deeper into what actually llama_model_params holds
  • It is a struct in llama.h
  • It has a list of devices the model will use I believe
  • llama_model_tensor_buft_override attribute seems way out of my understanding, let's just continue and come back
• Now we actually load the model using llama_model_load_from_file() that takes file loc and parameters
• after some error handling, a llama_vocab object is created, no idea what it is but maybe will understand after seeing its use
• Tokenizing begins !!
  • It finds num_tokens, using vocab object and prompt, wait the llama_tokenize function is fassed NULL for llama_token object and 0 for n_tokens_max?
  • My hypothesis is that it is previous data, like if we a conversation with AI the prompt includes previous stuff
  • Interesting that a vector of llama_tokens is intialized and fitted into llama_tokenize again
  • Ok I think I know what is happening, llama_tokenize returns the number of tokens it fails to tokenize, so all fail with zero so we get -ve of size and then we make a vector of that specific size.
• Next some context stuff
  • intitialize llama_context_params with default value and set it up.
  • So it looks like it will tell the system what size of matmuls will occur, like n_ctx attribute is total tokens(prompt+predict-1){why the -1?} and n_batch is number of tokens to process
• Then we actually create a context using our params and model
• Now some sampling that I believe is for next day

TODO:

• what is loading dynamic backends, need to dig deeper into ggml_backend_load_all()
• Get deeper into nearly all functions, their names are descriptive and so are their docs but need to read them properly
• Next day: complete simple.cpp and read the docs of every function

Day 3

14 December 2025 (1.5 hrs)

Today's Target: Finish understanding simple.cpp

• First I spent 30mins understanding how LLMs work from a 3Blue1Brown video.
• Now begins understanding the sampler
  • First default parameters are set for sampler and a sampler chain is initiailized(a llama_sambler)
  • then to the sampler chaing greedy is added
• Wait I cant understand where all the output is generate, lets try to tinker with printing of prompt token by token and see where is the change
• Oh ok it is quite at the bottom, first token is <｜begin▁of▁sentence｜>
• So llama_token_to_piece maps the token back to word in the vocab
• Now a batch is intialized and if encoding is required something else happens
• The main engine now
  • n_pos counts number of tokens done I think
  • First we evaluate current batch, and increament n_pos with obtained tokens
  • Then we sample next token, print it and get a new batch.
  • Note that these batches are size 1 so generate only 1 token.
• Finaly just some time printing and freeing allocating memory

TODO:

• Get deeper into nearly all functions, their names are descriptive and so are their docs but need to read them properly
• Maybe once look into simple-chat.cpp
• Maybe read a cuda file(found them in ./ggml/src/ggml-cuda/ using find . -regex ".*\.cu")
• Today I am reading reduction, I think argmax.cu will be very good for me. I have also read convolution chapter so rading conv2d.cu also should be helpful

Day 4

16 December 2025 (2.5 hrs)

Today's Target: Look into simple-chat.cpp and argmax.cu

• Just finished reduction, now let's dive into argmax.cu
• Just one function that takes in context and returns a tensor.
• So it calls argmax_f32 kernal, which instead of taking context as input reads it from a stream? or is it x?
• Need to look into how is this function called
• So looks like a ggml_tensor object holds the operation and operands and takes in the answer.
• So the only guy who uses this function is evaluating or capturing cuda graph and the tensor is a node in some computation graph
• Let me completely disect ggml_cuda_argmax once:
  • We get backend cuda context as ctx and a tensor dst probably meaning destination
  • src0 is another tensor that is dst->src[0], so it is a source tensor in dst tensor
  • As it is a f32 operation GGML asserts the type of src0 is GGML_TYPE_F32 and as output is integer(location) dst->type is GGML_TYPE_I32
  • Confusing on why a source is f32, My idea is a tensor is kept as operations and sources until its calculation is necessary/called
  • I think that should be it, then src0 must be contiguous.
  • Now ne00 is the number of elements in src0, 0th dim
  • nrows is number of rows in src0 which is tensor->ne[1]*tensor->ne[2]*tensor->ne[3] as product of other 4 dims
  • src0_d is the data in src0, in struct it is a void * and I think will be pointing to first element of array that actually stores data
  • similarily we get dst-data. Also it seems al these are pointers on device memory as no cudaMalloc is called
  • a stream in ctx is declared and given to kernal
  • Kernal has thread per block as min of 1024, ceil(ne00/WARP_SIZE)*WARP_SIZE. Brackets are somewhat confusing.
  • But wait num_blocks is nrows what if ne00 is too large, shouldn't there be more blocks, a possibility is coarsening but this is important
• Now let's look into the kernal itself

  • row processing is blockIdx.x and maxval is initilized to max float with argmax=-1

  • Why is argmax an int, a uint32_t or size_t seems better suited

  • Oh I forgot to map, data of src 0th dim is now x and dst_d is dst. I need to learn what __restrict__ is?

  • rowx is the current row to process as x+row*ncols, note ncols is ne00

  • As I thought following for loop coarsens the threads for large ncols

  • What it does is gets the maxval and argmax for every thread

  • Now another for loop with unroll pragma for performance?

  • I did this yesterday, this looks like convergent reduction

  • Why is offset >>= 1 and not offset /= 2 and offset is int, can save space by using short?

  • Thing is does it really matter cause the pragama will open the loop for constant values of offset.

  • Whoah __shfl_xor_sync is a pro thing I need to learn, looks like it gets the maxval and argmax of some other thread and evaluates the better one, with no conditional statements like thread id < something etc

  • I am thinking of a potential optimization to prevent control divergence of

    if (val > maxval) {
      maxval = val;
      argmax = col;
    }
    

    by replacaing this with something like this

    bool is_greater= (val > maxval);
    maxval = val * is_greater + maxval * !is_greater;
    argmax += (col-argmax) * is_greater;
    

    It occurs at many lines so can be better, but might as well have larger arithmatic overhead and if not might be optimized by compiler itself

  • So we can assume that every thread in a warp now has argmax of that respective warp.

  • some new id's for warps like lane_id, warp_id are declared

  • If n_warps > 1 then we need more processing so, now we define stuff in shared memory as the argmax and maxval of respective warps

  • after syncing to make sure everyone has filled respective filed in shared memory, 0th warp activates and does the same thing.

  • Clever thing done here is that max_warps cannot exceed 32 ie the WARP_SIZE.

  • So we can just load everything into their respective regs and process again.

  • Why do we need the shared memory then,...Oh to share data between threads. Looks good

  • Finally 0th thread gives output.

• What if future sizes of WARP increases the offset used in argmax will be affected, it can be made more robust by using offset = WARP_SIZE/2
• After some discussion with AI I posted a PR with two lines changed

TODO:

• Why is argmax an int, a uint32_t or size_t seems better suited
• Read about CUDA Warp-Level Primitives from nvidia dev technical blogs
• Read simple-chat.cpp tomorrow didnt have time today
• ggml_cuda_compute_forward function has a lot of such functions being called, look for some other function that looks good to anaylse

Day 5

19 December 2025 (1 hrs)

• My last PR was merged 🥳🥳🥳
• Today I comleted scan chapter in PMPP, I think acc.cu implements accumulate ie sum scan
• Well I was wrong it is something else, so let;s just try looking at matmul today, for acedamic curiosity. I am sure it would have been deeply analyzed by a lot of people before me so no change there
• I should go for mm*.cu files
• Interesting that they use #pragma once in place of headerguards using #ifndef ...
• Ok it is too big to analyse with a lot of namespaces and all and so many matmuls
• It is better I try something small like mean and sum today. When I am ready back to matmul
• If cub is active then something happens but I believe I can skip that for now.
• First lets see mean of a tensor
• As last takes in ctx and dst pointers to context and tensor
• src0 is the 0th source whose sum it will calculate
• src0_d is data whose sum is required, dst_d is where data is to keep
• then stream I think pmpp mentions is about clusters and stuff so keep myself away for now
• input output are both float32 and source is contiguous mem array
• ncols is the size of first dimension of tensor and nrows is how many of them are there
• I wonder why are GGML_MAX_DIMS just 4 whereas pytorch tensors can be n dimentional, also the ggml_nrows function hardcodes the number of rows, wondering if a loop will be more robust with a #pragma unroll(similar for ggml_nelements)
• If we skip CUB used for graph reduction we go to that block dims is just number of rows(nrows) and not min(nrows, 1024). That is suspicious
• Now number of streaming multiprocessors is requested and if they are more than twice the rows, block dimention is 512, else it is lesser
• If nsm is larger we have a lot of resource, and as we have a block per row we can devote a much larger block but these numbers are still magical, need to understand the intent begind them
• Now the kernal is called, reduce_rows_f32, need to find it
• It is in reduce_rows.cuh
• x is source and dst gets the sum/mean
• Interesting... manual unrolling!
• we have a block per row, so the col, row variable make sense as thread and block idx
• for every col we take first j values in temp and then add them to sum_temp[j], why not directly add? Why use temp as they are private to threads?
• Finally sum has sum of all the points that the thread was required to maintain
• now warp_reduce_sum was called that use warp level primitives to add all the sum variables in warp and put the answer at thread 0?
• Now if block dimensions are greater than WARP_SIZE we calculate the warp_id and lane_id of all and add them to s_sum
• Now we set sum to zero except to process lanes where we again do the warp_reduce_sum trick and finally return the mean or sum based on norm
• Unroll is defined using const int num_unroll = 8; which is very fishy
• If it was a compile time constant a pragma unroll would do a better job. Maybe the pragma works here too
• Also now I need to understand why so complications while calling this kernal
• If we have (nrows / nsm) < 2 we have 2 rows per nsm? I still dont get it
• Oh yeah lets use gitlens to see when it was commited, who did it and why, maybe add a comment
• Ok this was a big PR that solved a lot of problems and put a lot for new comers, lemme force my head harder maybe these numbers just work?
• So an important thing is a block can remain only on one SM
• I think I should ask some one what these means
• Oh okay in the PR the person mentions that pragma unrolls did not work
• I think there is a lot more to learn from this PR

TODO:

• Read this PR properly and completely analyse

Day 6

20 December 2025 (1 hrs)

• I feel this is not very robust as num_unroll is just hardcoded

  float     sum        = 0.0f;
  const int num_unroll = 8;
  float     temp[num_unroll];
  float     sum_temp[num_unroll] = { 0.0f };
  for (int i = col; i < ncols;) {
    for (int j = 0; j < num_unroll; ++j) {
      if (i < ncols) {
        temp[j] = x[row * ncols + i];
      } else {
        temp[j] = 0;
      }
      i += blockDim.x;
    }
    for (int j = 0; j < num_unroll; ++j) {
      sum_temp[j] += temp[j];
    }
  }
  for (int j = 0; j < num_unroll; ++j) {
    sum += sum_temp[j];
  }
  

  I am currently thinking if this will be better, and what if float4 is used?

  #define NUM_UNROLL 8
  
  .
  .
  .
  
  for (int i = col; i < ncols;) {
    #pragma unroll
    for (int j = 0; j < NUM_UNROLL; ++j) {
      if (i < ncols) {
        temp[j] = x[row * ncols + i];
      } else {
        temp[j] = 0;
      }
      i += blockDim.x;
    }
    #pragma unroll
    for (int j = 0; j < NUM_UNROLL; ++j) {
      sum_temp[j] += temp[j];
    }
  }
  #pragma unroll
  for (int j = 0; j < NUM_UNROLL; ++j) {
    sum += sum_temp[j];
  }
  

  • Using float4 wont be very useful as we are accessing locations spaced by blockDim.x but as in same run i is continuous for warp memory access is coalesced
  • The above changes gave a definate improvement over older ones though not very large

    Device 0: NVIDIA GeForce RTX 4050 Laptop GPU, compute capability 8.9, VMM: yes
  | model                          |       size |     params | backend    | ngl |            test |                  t/s |
  | ------------------------------ | ---------: | ---------: | ---------- | --: | --------------: | -------------------: |
  | qwen2 1.5B Q4_K - Medium       |   1.04 GiB |     1.78 B | CUDA       |  99 |           pp512 |       1864.15 ± 1.39 |
  | qwen2 1.5B Q4_K - Medium       |   1.04 GiB |     1.78 B | CUDA       |  99 |           tg128 |         43.79 ± 0.20 |
  

  after the changes:

    Device 0: NVIDIA GeForce RTX 4050 Laptop GPU, compute capability 8.9, VMM: yes
  | model                          |       size |     params | backend    | ngl |            test |                  t/s |
  | ------------------------------ | ---------: | ---------: | ---------- | --: | --------------: | -------------------: |
  | qwen2 1.5B Q4_K - Medium       |   1.04 GiB |     1.78 B | CUDA       |  99 |           pp512 |       1874.16 ± 0.60 |
  | qwen2 1.5B Q4_K - Medium       |   1.04 GiB |     1.78 B | CUDA       |  99 |           tg128 |         44.04 ± 0.16 |
  

  This would mean a 0.5 % speedup in pp512(what ever that is) and around same in other one on RTX 4050 Ok running again it is kind of slower so this change holds nearly no value Also running it on godbolt proves there is absolutely no change on compiled assembly

• After thinking put up a PR to add this comment so that others dont get confused by magic numbers

  /*
   * Occupancy heuristic:
   * 1) If number of rows is small (low parallelization), devote more threads (512) per row.
   *    This gives the SM 16 full warps to cycle through, hiding memory latency via pipelining.
   * 2) If nrows is large but cols are medium-large, use 128 threads (4 warps).
   *    Since there are many blocks (one per row), the GPU scheduler can hide latency by switching blocks.
   *    128 threads is enough to saturate the concurrent execution of 4 warps typical in an SM.
   * 3) If columns are small (< 1024), use 32 threads (1 warp).
   *    With 128 threads and 8x unroll, <1024 cols results in only 1 loop iteration,
   *    making the synchronization overhead of larger blocks inefficient.
   */
  

TODO:

• Next time matmul I am coming for you!

Day 7

21 December 2025 (1 hrs)

Today's Target: Start fighting with matmul

• Last day pr was suggested to just add link to prevous pr with discussion and was then merged 🥳
• Now I think mmf.cuh is best for mat mul float, it must be simplest one
• This is much heavier than other .cuh files I analyzed with 835 lines!
• It is using ggml_cuda_mma which is in mma.cuh another very heavy file
• I think tackeling it would be better through the ggml_cuda_mul_mat_f function in mmf.cu

  • It takes in ctx like all but being a binary operator takes in two tensors src0 and src1, also some other tensor ids with abigous purpose and dst for destination

  • As per the asserts src1 and dst are of type f32 and ids is either null or i32 type

  • This define seems weird GGML_TENSOR_BINARY_OP_LOCALS

  • So shouldnt we use GGML_TENSOR_UNARY_OP_LOCALS in argmax.cu and mean.cu and sum.cu?

  • Wow I learned about the ## operator

  • Let me unpack GGML_TENSOR_BINARY_OP_LOCALS once:

    GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \
    GGML_TENSOR_LOCALS(size_t,  nb0, src0, nb) \
    GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \
    GGML_TENSOR_LOCALS(size_t,  nb1, src1, nb) \
    GGML_TENSOR_LOCALS(int64_t, ne,  dst,  ne) \
    GGML_TENSOR_LOCALS(size_t,  nb,  dst,  nb)
    

    • so we have const int64_t ne03 = (src0) ? (src0)->ne[3] : 0
    • so ne0i would mean the number of elements in 0th source ith dimension
    • similarily ne1i would mean the number of elements in 1st source ith dimension
    • for dst we eliminate the 1 and 0 to nei
    • Wait what is nb then? Ok in ggml.h it is the stride in bytes so for transversing u units in ith dim we use A + u * nb[i]
    • I am thinking of doing all this in argmax.cu, but in that the dunction just takes dst and uses its src.
    • Now we get source vector type sizes and start comparing their sizes
    • But if we are sure that they are floats as the function is why do all this? Maybe leave this as having asserts is good, no harm in it or maybe they take up some cycles while running
    • Need to test if removing redundant ones improves performance and it doesn't so nvm
    • Important line is GGML_ASSERT(ne13 == ne3), why this happens?
    • Oh yeah simple, it means the third dim of src1 is same as that of dst but why?
    • Lets leave that once and see, now the transversal bytes for 0th dim must be same as type of tensor
    • Other dims can have different ratios represented by s[01]?[123] regexp
    • so for whatever ids are
    • I dont know what MUL_MAT_ID is so I assume ids is null now and analyse normal matmul
    • bassed on type of src0 different functions are further called all based on template of mul_mat_f_switch_cols_per_block
    • That is a very large function in mmf.cuh only because of a large switch case

      switch (ncols_case) {
          case  1: {
              mul_mat_f_cuda<T,  1>(x, y, ids, dst, ncols_x, nrows_x, ncols_dst, stride_row, stride_col_y, stride_col_dst,
                  stride_col_id, stride_row_id, nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
                  nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, stream, ids_data);
          } break;
          case  2: {
              .
              .
              .
          case 16: {
              mul_mat_f_cuda<T, 16>(x, y, ids, dst, ncols_x, nrows_x, ncols_dst, stride_row, stride_col_y, stride_col_dst,
                  stride_col_id, stride_row_id, nchannels_x, nchannels_y,  nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
                  nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, stream, ids_data);
          } break;
          default: {
              GGML_ABORT("fatal error");
          } break;
      }
    

    After asking AI if who recommended function pointer table, indexing or macros I think this would be best

    #define MUL_MAT_F_CUDA_CASE(N) case N: {
            mul_mat_f_cuda<T, N>(x, y, ids, dst, ncols_x, nrows_x, ncols_dst, stride_row, stride_col_y, stride_col_dst,
                stride_col_id, stride_row_id, nchannels_x, nchannels_y,  nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
                nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, stream, ids_data);
        } break;
    .
    .
    .
    
    switch (ncols_case): {
      MUL_MAT_F_CUDA_CASE(1)  MUL_MAT_F_CUDA_CASE(2)  MUL_MAT_F_CUDA_CASE(3)  MUL_MAT_F_CUDA_CASE(4)
      MUL_MAT_F_CUDA_CASE(5)  MUL_MAT_F_CUDA_CASE(6)  MUL_MAT_F_CUDA_CASE(7)  MUL_MAT_F_CUDA_CASE(8)
      MUL_MAT_F_CUDA_CASE(9)  MUL_MAT_F_CUDA_CASE(10) MUL_MAT_F_CUDA_CASE(11) MUL_MAT_F_CUDA_CASE(12)
      MUL_MAT_F_CUDA_CASE(13) MUL_MAT_F_CUDA_CASE(14) MUL_MAT_F_CUDA_CASE(15) MUL_MAT_F_CUDA_CASE(16)
      default: {
        GGML_ABORT("fatal error");
      } break;
    }
    

  • will continue from this function tomorrow

  • maybe add this PR to improve readability

TODO: mul_mat_f_cuda function

Day 8

22 December 2025 (30 mins)

• My last pr was rejected as they preferred readability and grepability over macros, it just reduced LoC but not code maintainability and so I closed it.

• Did not do a lot, just opened Issues and saw this one tagged ggml cuda with no comments:

  Problem description & steps to reproduce:
  When running llama-fit-params on a model that the tool believes fits (or also if the model cannot be loaded which is #18084) the suggested command line of "-c X -ngl 999" is printed in a random location during the output, varies from run to run

• So I asked AI if this could be done by me, though and it recommended finding and looking into llama-fit-params file and told there is a race condition between stderr and stdout. Thing is both these streams work at different pace and sometimes the final result output using stdout is output before stderr(used for logs) or in between

• There was a line with something like std::this_thread::sleep_for(10ms); // to avoid a race between stderr and stdout

• AI told to use std::flush(stderr) but it did not work

• So then I started understanding how logging happened, apparently all log commands add a log line to a logger that priints them when the thread is idle such as when sleeping.

• So I put a line common_log_pause(common_log_main()); in place of the wait and put a pr.

Day 9

23 December 2025(1 hr)

• My last PR got a comment to use common_log_flush instead and so I ammended and force pushed that commit
• Today I continue with the matmul
• I was on end of ggml_cuda_mul_mat_f function, which calls mul_mat_f_switch_cols_per_block ungodly huge switch case function that only checks the number of cols and calls mul_mat_f_cuda with proper template args
• Now we look into mul_mat_f_cuda
• First different tiles are defined of type ggml_cuda_mma::tile, need to see what it is first
• WHAT it is an empty struct with bunch of template args??
• Will see how tile is used in final function then
• Now we need to note that
  • ncols_x is even, so is stride_row and so is stride_col_y
  • Also either ids should be a proper array and not null otherwise n_channels_ds must be devisible by nchaneels_x
  • Also nsamples_dst is divisible by nsamples_x
• We get both the previous ratios, the device, cc and warpsize.
• So we got compute capabiliity and warp_size(point to note about my first PR in other places they do get warpsize as run time constant, potential for refactoring that #define WARP_SIZE 32)
• Wondering why MMF_ROWS_PER_BLOCK is 32?
• I think I should build up from bones: going from here -> mul_mat_f_switch_ids -> mul_mat_f which I think does the calc for id less case
• a bunch of tiles are defined first based on what MMA is available.
• This is unlike most other kernals I have seen that do not use data about the device
• say I am neither using cuda hip or musa, so tile_A is 16X8 and tile_B is 8X8 with output tile_C as 16X8 float type
• Again using ggml_cuda_get_physical_warp_size, maybe I should refactor argmax and mean
• then there is a tile_k_padded = warp_size + 4 for unknown reason
• and ntA is rows_per_block/tile_A::I;
• and ntB is (cols_per_block + tileB::I - 1)/tile_B::I;
• Was it mentioned that rows_per_block is multiple of tile_A::I
• Btw now I get why we template arg tile but why not just add them as attributes of the struct
• row0 was blockIdx.x * rows_per_block ie 0th row that this block processes
• there is some expert_idx=0 and some col_base=0
• if it has ids then channel_dst is 0 other wise blockIdx.y, I think let us assume it has no ids
• A lot of things happening ... I am very confused need a break will continue tomorrow, did nothing important today 😞

Day 10

24 December 2025 (5hrs)

Target: I am too saturated to go behind matmul right now and saw potential of improvement in cumsum hence I shall be going full on to it to improve performance on llm arch like mamba/qwen3

• First I need to recon on how much potential is there if there is a fall back from CUB

• Putting a printf anywhere in cumsum doesnt work, and apparantly using gdb also it never breaks in cumsum file, now adding a GGML_ASSERT(false) also does not break the code, hence it never arrives to cumsum!

• Wait asking AI it says that mamaba relies on ssm-scan, what is that now!

• Ok so GGML_ASSERT(false) here breaks the code so it passes through this file 🤦

• No I must go back to cumsum as mamba does not use it but qwen3 does so need to download a qwen3 model

• I saw this

  #ifdef GGML_CUDA_USE_CUB
    // Check if we can use CUB (data must be contiguous along innermost dimension)
    const bool is_contiguous = (nb00 == type_size) && (nb0 == type_size);
  
    if (is_contiguous) {
      use_cub = true;
    }
  #endif // GGML_CUDA_USE_CUB
  

  and I think this would be better if ggml_is_contiguos_0 is used but that is on tensor and this function takes in alll data

• Say we do not use cub then we see the call to kernal cumsum_kernal

• why is this a 3d thing!?

• I cant do that, I checked up if I could download a qwen3 model but it too did not use cumsum, hence I got to know it is qwen3-next that uses it and that model is 50GB in size much more than my VRAM

• Ok so I am going to do it! and I am going to build a unit tester for it

• Lemme vibe code for some time 😎

• Ok so N=152000 cumsum takes around 1.0075 ms

• Interesting that even after turning off cub cmake shows no compilation hence it was never active

• now the assert breaks it 🥳 but code says it runs cub, more analysis results in that there is no way to force it to fallback on cumsum but I can change source code in cumsum

• Hmm the fallback kernal takes 1.0175 ms quite close, maybe it scales with N, let me do some checks

• now cub with 15_200_000 takes 109.927 ms and no cub takes 112.8467 ms which is a very small improvement around 2.5%, small but significant.

• Lets read gpu gems now to research the best algo Chapter 39. Parallel Prefix Sum (Scan) with CUDA

• Reading thoroughly i see that everywhere multi pass is used for (int64_t start = 0; start < ne00; start += BLOCK_SIZE)

• idk why but this seems to be the most efficient implementation

• so after thinking and planning and doing this is what I get, need to filter it once and we are done

  template<typename T>
  static __global__ void cumsum_kernel(
          const T * src, T * dst,
          const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
          const int64_t  s00, const int64_t  s01, const int64_t  s02, const int64_t  s03,
          const int64_t   s0, const int64_t   s1, const int64_t   s2, const int64_t   s3) {
  
      GGML_UNUSED_VARS(s00, s0);
  
      const int tid = threadIdx.x;
      constexpr int warp_size = ggml_cuda_get_physical_warp_size();
      const int lane = tid % warp_size;
      const int warp = tid / warp_size;
      const int warps_per_block = blockDim.x / warp_size;
  
      extern __shared__ float smem[];
      float * s_vals          = smem;
      float * s_warp_sums     = smem + blockDim.x;
      float * s_carry         = smem + blockDim.x + warps_per_block;
      float * s_chunk_total   = s_carry + 1;
  
      // Initialize carry
      if (tid == 0) {
          *s_carry = 0.0f;
      }
      __syncthreads();
  
      const int64_t i3 = blockIdx.z;
      const int64_t i2 = blockIdx.y;
      const int64_t i1 = blockIdx.x;
      if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
          return;
      }
  
      const T * src_row = src + i1 * s01 + i2 * s02 + i3 * s03;
      T       * dst_row = dst + i1 * s1  + i2 * s2  + i3 * s3;
  
      // every thread is assignmed 8 nums
      const int num_unroll = 4;
      // this stores all the nums
      T temp[num_unroll];
      // i represents where we start from
      for (uint32_t i = 0; i < ne00; i += num_unroll*blockDim.x) {
          // idx is where this thread starts from
          uint32_t idx = i + tid * num_unroll;
          // we set the first element of temp manually and rest is 
          temp[0] = (idx<ne00 ? src_row[i + num_unroll*tid] : 0);
          for (uint32_t j = 1; j < num_unroll; j++) {
              temp[j] = temp[j-1];
              if (idx + j < ne00) {
                  temp[j] += src_row[idx + j];
              } else {
                  temp[j] += 0;
              }
          }
          // so currently we have temp[i] = temp[0]+temp[1]...temp[i]
          // we assign the net sum to val, this will be sum that is added to next
          float val = (idx < ne00) ? ggml_cuda_cast<float, T>(temp[num_unroll-1]) : 0.0f;
  
          // 1. Warp inclusive scan adds up the vals in the warp
          // hence we just need to add val-temp.back() to all temp to get result for entire warp
          val = warp_prefix_inclusive_sum<T, warp_size>(val);
          s_vals[tid] = val;
  
          // Store warp total of all into s_warp_sum and as warp per block is strictly less than 32 we fit them all in a warp 0 to process next
          if (lane == warp_size - 1) {
              s_warp_sums[warp] = val;
          }
          __syncthreads();
  
          // 2. Exclusive scan of warp sums (warp 0 only)
          if (warp == 0) {
              float w = (tid < warps_per_block) ? s_warp_sums[tid] : 0.0f;
              float inc = warp_prefix_inclusive_sum<T, warp_size>(w);
              if (tid < warps_per_block) {
                  s_warp_sums[tid] = inc - w;   // exclusive sum
              }
              if (tid == warps_per_block - 1) {
                  *s_chunk_total = inc;          // total sum of this chunk
              }
          }
          __syncthreads();
  
          // this is carry of previous i block iteration
          float carry = *s_carry;
          // s_warp_sums[warp] is the sum added from previous warps
          // s_vals[tid] is net sum of all vals in the present warp upto this inclusive
          float final_val_offset = s_vals[tid] + s_warp_sums[warp] + carry - temp[num_unroll-1];
          for (uint32_t j=0; j<num_unroll; j++) {
              if (idx+j < ne00) {
                  dst_row[idx+j] = temp[j] + ggml_cuda_cast<T, float>(final_val_offset);
              }
          }
          // this syncthread usage is ambiguous as there is no correlation
          // need to check requirement
          __syncthreads();
  
          // Update carry for next chunk
          if (tid == 0) {
              *s_carry += *s_chunk_total;
          }
          __syncthreads();
      }
  }
  

  Now filtering and stuff, final PR comment: https://github.com/ggml-org/llama.cpp/pull/18343#issue-3760074245

• data is great

  |Implementation| Time (ms)| Relative Speed| |:---|:---|:---| |Old Fallback (Naive) |~1.018 ms |1.0x| |Current CUB Wrapper |~1.00 ms |1.018x| |New Fallback (This PR)| 0.40 ms |2.54x|

• Finally submitted the PR, this was kind of exhausting

• signing off for today, let's see what a new day brings 😆

Day 11

28 December 2025 (2.5 hrs)

Target: I am thinking of learning more about warp level primitives from nvidia dev blog and then potentially getting a good optimization on the ssm kernals as they seem to be quite naive implementations

• Starting with ssm-scan, let me test if mamba-2.8b-q4_k_m.gguf uses it using the good old assert false
• Got it ggml_cuda_op_ssm_scan will be target to understand today
• Maybe should read about mamba ssm first how it works, should skim through this https://arxiv.org/pdf/2312.00752
• It looks like this model is defined by four params ($\Delta$,A,B,C), I can see A,B,C being unpacked from the dst tensor but not $\Delta$, I would assume that it is actually src(dt)
• Another assumption that would be safe to make is that A,B,C here is discretized
• Yeah it is a part of coputation ($\Delta$,A,B,C) |-> ($\overline{A},$\bar{B}$,C), maybe it is done autocatically or that is what is stored
• Oh so Linear Time Invariance says that they are constant, good thing all A,B and their bars are const
• Everything then on seems distant so lets stop reading and observe what the code does
• Lets step into ssm_scan_f32_cuda as the wrapper is just unpacking tensors
• First lemme get AI to generate a temporary test script
• Wow mamba-2 has a lot going on, lets just read mamba 1 code first
• interesting that number of threads is fixed before hand and n_head must be its multiple
• also head_dim is 1 and so is n_group
• still the block defined is using n_seq, ceil(n_head/threads), 1
• we declare shared memory size as (threads * (d_state + 1) * 2) * sizeof(float)
• Now if d_state is 16 which seems to be the size of state like the N used in the paper with A and B were NxN matrices and C was Nx1 I think
• Now based on number of tokens the kernal is called
• Let's step into the ssc_scan_f32 kernal
• Some value L is set as L_param if L_template is zero otherwise it is L_template
• given the calls splitD is fixed as threads = 128. N is 16 and L_template is highly variable
• An insanly weird method is used to define offsets
  • They type cast origin to const char * so addition moves by 1 byte then add some number multiplied with sizeof(float) or srci_nb3 which is sizeof(float) because of earlier asserts
• Let me play with it to simplify the system
• Ok I get the system nb0 are asserted as sizeof(float) not nb3 so need to leave it here
• Ok shit they use CUB that is like very fast, lets see if we can optimize the fallback
• No wait they use CUB to load data from global mem to shared mem, now thats lazy!
• So I think N is unrolls? Idk there is this regA[N] and regs0[N] where data is copied and so is done in shared for B and C
• First we load data from A_block to regA and from s0_block to regs0
• Now for every step L if threadIdx.x < N so first into the shared memory we store B and C values
• get a dt_soft_plus and if it is <= 20.0f we change it to log1pf(expf(dt_soft_plus)) ie if dt <= 20 we do dt = $\ln (1+e^{dt})$, I wonder why
• Then sumf is 0.0f and we do

  for (size_t n = 0; n < N; n++) {
      float state = regs0[n] * expf(dt_soft_plus * regA[n]) + smemB[n] * x_dt;
      sumf += state * smemC[n];
      regs0[n] = state;
  }
  

• output is in y_block and s_block
• for given block y_block[threadIdx.x + i*stride_y] is the sumf of state*C_i
• consider a vector C of size N. So let s0, A, B
• we need output s = s0 * e^{A dt} + B * x dt
• ... I need a place where all this properly defined there is not even a test for this ...
• I am asking ai for resources
• Wow this is such a great way of understanding what SSM models do! https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mamba-and-state
• Reading on very interesting idea is to use constant memroy for A,B,C as they never ever change
• Ok now I seem to understand better delta is the discrete step so dt, hence bar(A) or the discretized A is e^{Adt}
• Also bar(B) or discretized B is (delta A)^-1(bar(A) - 1) * delta B
• What we are given: initial state s, input signal x, delta dt, params A, B and C (and ids)
• It does look like A, B, and C are small enough for const mem operation, or maybe not
• A is like just 48 elements, B is 128*512 and so is C, but this is test, lets try mamba model: that has A as 16 *5120 and B and C both 16?
• I think this wont work here but holds potential for ssm-conv
• I am done with this today learnt a lot about ssms and all
• Just a little look w in conv can be loaded into const mem, will see about that later

Day 12

29 Decmeber 2025 (3 hrs)

• Now I have some idea what ssm mathematically is lets try agiain
• Once lets see what values n_tok takes using some log instructions
• My idea is to kind of precompute ${\overline{A}}^2$ and then do it in multiple steps if n_tok is large enough which it should as it is number of tokens?
• While mamba1 model is downloading lets do something else
• I think I also should creat a cheat sheet lable for sizes of matrices(filled along with analysis) | matrix | src | dims | |:-------|:-----|:-----| | state| 0 | state dim X 1(or head dim?) X heads size X ?(id related) | x | 1 | ? X heads size X tokens per seq X num of seq in batch | dt | 2 | | A | 3 | | B | 4 | ? X number of groups | C | 5 |
• blocks used are (number of sequences in the batch, number of heads / thread count, 1)
• so each block line processess a sequence and there are as many lines as there are sequence in a batch
• It looks like each thread is assigned a head dim
• Now each block is simply hardcoded 128 threads
• for mamaba 1 it seems head_dim is one, so head is one dimentional
• as n_tok is token per sequence each block line processes one sequence of tokens
• Value of N is fixed as 16 and L is L_template == 0 ? L_param : L_template
• Point to note is this is useless as l_param is n_tok which is set using the switch case
• Wait I realized my mamba model is 1, but with 2.8b params hence mamba-2.8 and it uses the Mamba2 code? I need to unit test this then
• Now we see that s0_block is the state and is moved by first something id related along 3rd dim and blockIdx.y * splitD * src0_nb2 means which head to start at
• So I believe head represent the state size etc and we have a thread per state val
• next x_block should represent the input waveform displaced form src1 by threads before it and different seq stuff, same for dt
• interesting that A is invariate to sequence and B and C are invariate on block
• So A being a property of state is same for sequences but B and C working for how input goes to next state and how output is generate from state they are same for all threads
• there is some tiling around here we need to take care of
• we load N initial state into regs0 and N of A into regA, same for B and C in shared mem
• So we have N initial states, N A,B,C vals ready, but why only N and not how many threads are there?
• For how many ever tokens are there, we loop now
• First load some values into shared memory for B and C(first N), B_block[i * stride_B + threadIdx.x] means every i value loads in a new B val, so a different value of B is used for every time step!
• now we calculate dt_soft_plus as $\ln (1+e^{d{t}_i})$, where dt is also different for time steps!
• x_dt is just value of x at that time times the dt_soft_plus
• now we define a sumf as 0
• for n from 0 to N we now get the state as stored state in regs0[n] * e^(A[n]*dt_soft_plus) + B[n]*x_dt, this seems to involve discretization.
• Then in sumf we add state*C[n] and set regs0[n] to state for next time step
• finally after the loop y(output waveform) at that timestep is set as sumf
• and also we store final regs0 back to dst
• Hey wait, if there is a cuda implementation ssm_scan must have a cpu one too, I can read that!!
  • We read this in kind of a subsection
  • this seems to be it ggml_compute_forward_ssm_scan_f32
  • Wow the comments are so beautiful 🥲
  • Its easy to tabulate it for understanding: | src | what it is | ne0 | ne1 | ne2 | ne3 | |-------|:------------|----------:|----------:|----------:|------:| |0 | state(s) |d_state |dim |n_head |n_seq| |1 |input wave(x)|dim |n_head |n_seq_tok|n_seq| |2 |delta T (dt) |n_head |n_seq_tok|n_seq |1 | |3 | A |d_state/1|n_head |1 |1 | |4 | B |d_state |n_group |n_seq_tok|n_seq| |5 | C |d_state |n_group |n_seq_tok|n_seq| |6 | ids |n_seq |1 |1 |1 |
  • So we have state size as d_state and for mamba 1 we assume 16 states and for mamba 2 we assume 128 or 256 state size on cuda!!
  • Head are channels and in mamba 1 it is equal to 1
  • on cpu we loop for every sequence
    • for every token in that sequence
      • for every head
        • for every dimension of state(dim)
          • for every element of state in that dim update state
          • calculate sum of them y = rowwise_dotprod(state, C)
  • On gpu we launch as many threads as size of channel that must size of input/output waveform
• Wait lets recollect all we have
  • We launch a thread for every input in channel representing the input and hance gives one output(and potentially final state)
  • For the number of tokens given we iterate throughout the number of tokens and find the output for all tokens individually, this being the bottleneck as the number of tokens can be large for a hefty model?
  • So what exactly happens is we have a given input wave x with n_heads size and note each token has different wave and needs different output y. We parallelize over size of n_heads and for every token using the given starting state iteratively calculate the new state for that wave position and get an output for y.
  • Hence, every thread has its own state, each thread works on a position on the waveform and consequently gives outputs for that waveform on that position for different time
  • Time travels with tokens, so it is kind of like at t=0 we have a waveform x, we generate y update state and at t=1 we have different waveform x and generate new y and update state and so on
• Potential of speedup do seem scarce
• will look into the possibility to unfold recurrence tomorrow

Day 13

30 December 2025 (4 hrs)

• Using the given command I tested how ssm_scan works and it turns out the original SSM scan is quite good and is memory bound(190 GB/s for small token size and 90 for 8000 tokens) and so improving it will require better memory access that will take time

  ./build/bin/test-backend-ops perf -o SSM_SCAN -b CUDA0 
  

• On the other hand the mamba 2 version is hgihly compute bound running at 9 GB/s memory hence a lot of optimization possible
• Lets read into the mamba 2 code ssm_scan_f32_group
• call is simple and depends on the d_state that is 128 or 256 in mamba 2 case, lets focus on the 128 case
• there is some splitH constant as 16 maybe the states maintianed by a thread or how many ways we split the head(the waveform)
• number of blocks is (n_head * head_dim/splitH, n_seq) hence splitH is the states per thread(coarsening I believe)
• Testing various splitH 16 gave best performance
• Lets step into the kernal now
  • head_idx is which dim we start from, head_off is the access point on that dim
  • so each block processess splitH? thats weird no blockDim.x is used
  • we have a seq_idx as the sequence index, great
  • group_off is some group related offset?
  • Oh yeah on top it writes assumes as many threads as d_state that we might change for potential speedup
  • What exactly is splitH then? Looks like this is quite a bit different
  • Lets do some calc, threads and d_state is 128, for a sequence if we consider a line of blocks there are n_head * head_dim * 128 / splitH total threads, lets wait and see
  • we get the state of block, x, A, B, C, y, final state, dt
  • we get their strides and all
  • Now a reg array of splitH size is defined and it is the state!
  • A shared array stateC is defined of splitH * d_state for parallel accumulation
  • step 1: load splitH items from state array, so each thread has is responsible for one point in the state and at once loads splitH state points in the array along different dimension of state
  • So if the state is say a x b then there are a threads and each handles b?
  • then for every token we do
    • get a dt_soft_plus that is same throughout the threads
    • get discretized dA and wait!! B and C are called from global memory😲😲😲
    • ah wait the problem is every thread has a different B and C here so large space in shared mem is required
    • Oh got it each thread does not need what other thread does so it happens this way, in shared mem we would have loaded into shared mem only to be used once by the thread who loaded it
    • There are some TODOs here that might be good to work with

      // TODO: only calculate dA and dt_soft_plus once per head instead of every splitH head elements
      // TODO: only calculate B and C once per head group
      // NOTE: dt_soft_plus, dA and x_dt have the same value across threads here.
      

    • Looking into that we have options
      • Let it be, all threads calculate same thing at their pace(seems bad)
      • Let threadIdx.x == 0 calculate it and others wait on __syncthreads()(seems better but stop seems quite big)
      • Let 0th thread in warp calculate in every warp and use __syncwarps()(seems best)
    • Another good option is, in shared memory for following many i we calculate dt and dA, like say in shared mem we store 2 128 sized arrays and after 128 turns all threads process dt_soft and dA for current time + thread index and then keep using it for 128 tunrs
    • // TODO: only calculate B and C once per head group is confusing, or in better words questionable given the design choice
    • Lets just try to implement the pre calc dt_soft and dA idea
      • First we need I believe 2 arrays, one for elements of dt_soft, one for dA and I believe one space for A_block[0] itself though it might just reside in the L1/L2 cache, lets do that later

      if (i%d_state == 0) {
          // we must refresh precomputed values
          if (i+threadIdx.x < n_tok) {
              float dt_soft_plus = dt_block[(i+threadIdx.x) * stride_dt];
              if (dt_soft_plus <= 20.0f) {
                  dt_soft_plus = log1pf(expf(dt_soft_plus));
              }
              dt_precompute[threadIdx.x] = dt_soft_plus;
              dA_precompute[threadIdx.x] = expf(dt_soft_plus * A_block[0]);
          }
      }
      __syncthreads();
      

      • No visible speedup in perf with same memory access rate, I wonder why? 😭
      • I need to go deeper into why this did absolutely no help, it reduces floating point calcluations by a factor of d_state(128) in this part of code, it should go to be memory bound?
      • Ok so new plan, __syncthreads() makes warps delay by 100s of cycles, buttttt __syncwarp() wont. Lets preprocess 32 of them then
      • Even that does not work so lets remove this and get back later, we need to work on the other todo about B and C preloading
      • So we sort of implement tiling for B and C, using not shared but I beleive register tiling

      if (i%tile == 0) {
          // load B and C into tiles
          for (int j=0; j<tile; j++) {
              if (i+j < n_tok) {
                  regsB[j] = B_block[(i+j) * stride_B + threadIdx.x];
                  regsC[j] = C_block[(i+j) * stride_C + threadIdx.x];
              }
          }
      }
      

      • even this seems to not be very helpful 😢, similar speeds only
      • Have to move to the real kernal then
      • Ok now I profiled it and bottleneck is stateC[k] += stateC[k + (w >> 1)];
      • Apparantly all the stateC load stores create bottlenecks
      • Master plan:
        • remove state C from static memory, and put a register to store the result privately
        • then conduct the cumsum kernal ideas to do the cumulative sums
    • Fingers crossed lets implement this
      • opening actual reduce_rows_f32 for reference
      • I cant understand this indexing and what exactly is committed, lets see the gpu implementation in ggml-cpu.c
      • After long time and trying to make it better still I get no speedup

            /*
            plan
            we use a big algo for reducing the sum however we can do better IMO using warp level primitives first
            first let us see how stateC was used -> it was stored in following format
                       d_state
              [ [                ]
                [                ]
                [                ]
                [                ] ]
            lets see what exactly must be stored etc
            only y_block is the commited part
            my hypothesis is that we store
            y_block[0] = sum of row 0 of stateC
            y_block[1] = sum of row 1 of stateC
            and so on
            now we privitize each column to each thread
            using float stateC[splitH];
      
            */
      
            // the +1 padding is to prevent bank conflicts
            const int lane_id = threadIdx.x % WARP_SIZE;
        #pragma unroll
            for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) {
            #pragma unroll
                for (int j=0; j<splitH; j++) {
                    stateC[j] += __shfl_xor_sync(0xffffffff, stateC[j], offset, WARP_SIZE);
                }
            }
            if (lane_id == 0) {
            #pragma unroll
                for (int j=0; j<splitH; j++) {
                    atomicAdd(y_block + (i * stride_y + threadIdx.x), stateC[j]);
                }
            }
        }
      

Day 14

31 December 2025(5 hrs)

• New plan instead of the given corsening along heads I propose a coarsening of each warp maintaining a state step
• So each thread handles 4 data points on state
• Lets test it out!
  • Important coarsening factor here will be d_state/WARP_SIZE so each warps handle one time step hence a thread handles lets call it c_factor for now
  • So there are as many warps as heads, hence ceil(n_head*head_dim / n_warps) and n_warps is threads/WARP_SIZE
  • num_warps == c_factor as threads == d_state
  • head_idx must be blockIdx.x * c_factor / d_head
  • Instead of anything_block we will use anything_warp however now I must completely understand the indexing procedure
  • After a long debugging session(and some hitting head in the AI wall) I got a 28% speedup!! and also the code is much much more simplified than it was initially!!

  // assumes as many threads as d_state
  template <int c_factor, int d_state>
  __global__ void __launch_bounds__(d_state, 1)
      ssm_scan_f32_group(
          const float * __restrict__ src0, const float * __restrict__ src1, const float * __restrict__ src2,
          const float * __restrict__ src3, const float * __restrict__ src4, const float * __restrict__ src5,
          const int32_t * __restrict__ src6, float * __restrict__ dst,
          const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3,
          const int src2_nb1, const int src2_nb2, const int src3_nb1,
          const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3,
          const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok) {
  
      const int warp     = threadIdx.x / WARP_SIZE;
      const int lane     = threadIdx.x % WARP_SIZE;
      const int warpIdx  = blockIdx.x  * c_factor + warp;
  
      const int head_idx =  warpIdx / d_head;
      const int head_off = (warpIdx % d_head) * sizeof(float);
      const int seq_idx  = blockIdx.y;
  
      const int group_off = (head_idx / (n_head / n_group)) * d_state * sizeof(float);
  
      const float * s0_warp = (const float *) ((const char *) src0 + src6[seq_idx] * src0_nb3 + head_idx * src0_nb2 + head_off * d_state);
      const float * x_warp  = (const float *) ((const char *) src1 + (seq_idx * src1_nb3) + (warpIdx * sizeof(float)));
      const float * dt_warp = (const float *) ((const char *) src2 + (seq_idx * src2_nb2) + head_idx * sizeof(float));
      const float * A_warp  = (const float *) ((const char *) src3 + head_idx * src3_nb1);
      const float * B_warp  = (const float *) ((const char *) src4 + (seq_idx * src4_nb3) + (group_off));
      const float * C_warp  = (const float *) ((const char *) src5 + (seq_idx * src5_nb3) + (group_off));
      float *       y_warp  = dst + (seq_idx * n_tok * n_head * d_head) + warpIdx;
      float *       s_warp  = (float *) ((char *) dst + s_off + seq_idx * src0_nb3 + head_idx * src0_nb2 + head_off * d_state);
  
      // strides across n_seq_tokens
      const int stride_x  = src1_nb2 / sizeof(float);
      const int stride_dt = src2_nb1 / sizeof(float);
      const int stride_B  = src4_nb2 / sizeof(float);
      const int stride_C  = src5_nb2 / sizeof(float);
      const int stride_y  = n_head * d_head;
  
      float state[c_factor];
      float state_sum = 0.0f;
  
  #pragma unroll
      for (int j = 0; j < c_factor; j++) {
          state[j] = s0_warp[WARP_SIZE * j + lane];
      }
  
      for (int64_t i = 0; i < n_tok; i++) {
          // TODO: only calculate dA and dt_soft_plus once per warp instead of every warp thread
          // NOTE: dt_soft_plus, dA and x_dt have the same value for a warp here.
          float dt_soft_plus = dt_warp[i * stride_dt];
          if (dt_soft_plus <= 20.0f) {
              dt_soft_plus = log1pf(expf(dt_soft_plus));
          }
  
          state_sum = 0.0f;
          const float dA   = expf(dt_soft_plus * A_warp[0]);
          const float x_dt = x_warp[i * stride_x] * dt_soft_plus;
      #pragma unroll
          for (int j=0; j<c_factor; j++) {
              float B_val = B_warp[i * stride_B + WARP_SIZE * j + lane];
              float C_val = C_warp[i * stride_C + WARP_SIZE * j + lane];
              state[j] = (state[j] * dA) + (B_val * x_dt);
              state_sum += state[j] * C_val;
          }
  
          // parallel accumulation for output
          state_sum = warp_reduce_sum(state_sum);
  
          if (lane == 0) {
              y_warp[i * stride_y] = state_sum;
          }
      }
  
      // write back the state
  #pragma unroll
      for (int j = 0; j < c_factor; j++) {
          s_warp[WARP_SIZE * j + lane] = state[j];
      }
  }
  

• WOW!! apparantly it is on my RTX4050 I am getting a 28% speedup, some other contributers with better GPUs like RTX4090, reported a 2-3x speedup in specific function and 10% speedup in inference of a mamba based model.
• This was painful but I think this was a rather good experience and I learnt a lot in 1 month(I started reading PMPP from start of month)
• Hence I finish with a PR I am quite proud of.