Repository: databricks/megablocks (language: python)

Issue:
Wrong outputs for hidden dim 14336
Was trying to run this: magnet:?xt=urn:btih:5546272da9065eddeb6fcd7ffddeef5b75be79a7&dn=mixtral-8x7b-32kseqlen&tr=udp%3A%2F%2Fopentracker.i2p.rocks%3A6969%2Fannounce&tr=udp%3A%2F%[2Ftracker.opentrackr.org](http://2ftracker.opentrackr.org/)%3A1337%2Fannounce&tr=udp%3A%2F%[2Ftracker.openbittorrent.com](http://2ftracker.openbittorrent.com/)%3A6969%2Fannounce&tr=http%3A%2F%[2Ftracker.openbittorrent.com](http://2ftracker.openbittorrent.com/)%3A80%2Fannounce&tr=https%3A%2F%[2Ftracker.gbitt.info](http://2ftracker.gbitt.info/)%3A443%2Fannounce

but the current code generates random outputs for hidden dim 14336
Support new model



Discussion / hints:

Solves issue https://github.com/stanford-futuredata/megablocks/issues/46
Hey Pierre! Nice fix!

Thanks for the contribution :) 
Added one comment - this should work fine as a quick fix to unblock you. 

I've been meaning to port this to Triton and I can also clean up the grid dimensions for different problem sizes as part of that :)

Fix patch:
diff --git a/csrc/indices.h b/csrc/indices.h
index 3787d831..df227e87 100644
--- a/csrc/indices.h
+++ b/csrc/indices.h
@@ -13,11 +13,9 @@
 namespace megablocks {
 namespace construct_indices {
 
-// We expect the number of outputs per block to be
-// small. With ffn_hidden_size=4096, we only need
-// to write 32 elements per block per iteration.
-// This is the largest we're every likely to use
-// so we keep the blocks small.
+// We expect the number of outputs per block to be small. For
+// example, with ffn_hidden_size=4096, we only need to write
+// 32 elements per block per iteration.
 const int kThreadsPerBlock = 32;
 
 __global__ void __launch_bounds__(kThreadsPerBlock)
@@ -39,13 +37,11 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
 
   // Write the indices to the output.
   int bin_offset = blockIdx.y;
-  int tid = threadIdx.x;
   int num_rows = end - start;
   for (; bin_offset < num_rows; num_rows -= gridDim.y) {
-    int elements = num_columns;
     short *out = indices;
-    for (; tid < elements; elements -= kThreadsPerBlock) {
-      *out = threadIdx.x + (blockIdx.x * num_columns);
+    for (int bid = threadIdx.x; bid < num_columns; bid += kThreadsPerBlock) {
+      *out = bid + (blockIdx.x * num_columns);
       out += kThreadsPerBlock;
     }
     indices += gridDim.y * num_columns;


Test patch:
diff --git a/megablocks/ops/topology_test.py b/megablocks/ops/topology_test.py
index a17a3813..0ff3491e 100644
--- a/megablocks/ops/topology_test.py
+++ b/megablocks/ops/topology_test.py
@@ -26,7 +26,9 @@
     (16384, 768, 128),
     (16384, 768, 256),
     (16384, 768, 512),
-    (16384, 768, 1024))
+    (16384, 768, 1024),
+    (8, 14336, 8),
+)
 
 
 class TopologyTest(parameterized.TestCase):