call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cublas_v2.h>\n#include <cstdio>\n\n// Constants for thread/block organization\n#define THREADS_PER_BLOCK 256\n\n// Simple top-2 kernel: finds top 2 values and indices for each row\n__global__ void top2_kernel(const float* __restrict__ scores,\n float* __restrict__ top_vals,\n int* __restrict__ top_idx,\n int num_tokens,\n int num_experts) {\n int token_idx = blockIdx.x * blockDim.y + threadIdx.y;\n \n if (token_idx >= num_tokens) return;\n \n // Each thread in y dimension handles one token\n const float* row = scores + token_idx * num_experts;\n \n float max1 = -INFINITY, max2 = -INFINITY;\n int idx1 = 0, idx2 = 0;\n \n for (int e = 0; e < num_experts; e++) {\n float val = row[e];\n if (val > max1) {\n max2 = max1;\n idx2 = idx1;\n max1 = val;\n idx1 = e;\n } else if (val > max2) {\n max2 = val;\n idx2 = e;\n }\n }\n \n top_vals[token_idx * 2] = max1;\n top_vals[token_idx * 2 + 1] = max2;\n top_idx[token_idx * 2] = idx1;\n top_idx[token_idx * 2 + 1] = idx2;\n}\n\n// Softmax on 2-element vectors and produce output\n__global__ void softmax_and_output_kernel(const float* __restrict__ top_vals,\n const int* __restrict__ top_idx,\n float* __restrict__ output,\n int num_tokens) {\n int token_idx = blockIdx.x * blockDim.x + threadIdx.x;\n \n if (token_idx >= num_tokens) return;\n \n float v0 = top_vals[token_idx * 2];\n float v1 = top_vals[token_idx * 2 + 1];\n \n // Softmax\n float max_val = fmaxf(v0, v1);\n float exp0 = expf(v0 - max_val);\n float exp1 = expf(v1 - max_val);\n float sum = exp0 + exp1;\n \n float w0 = exp0 / sum;\n float w1 = exp1 / sum;\n \n // Output: [num_tokens, 2, 2] - dim=-1 is last\n // output[token, 0, :] = [idx0_float, idx1_float] \n // output[token, 1, :] = [w0, w1]\n output[token_idx * 4 + 0] = (float)top_idx[token_idx * 2];\n output[token_idx * 4 + 1] = (float)top_idx[token_idx * 2 + 1];\n output[token_idx * 4 + 2] = w0;\n output[token_idx * 4 + 3] = w1;\n}\n\n// Initialize cuBLAS handle - use a global handle\nstatic cublasHandle_t get_cublas_handle() {\n static cublasHandle_t handle = nullptr;\n if (handle == nullptr) {\n cublasCreate(&handle);\n }\n return handle;\n}\n\n// Main kernel function\nvoid grounded_top2_moe_router(\n const float* token_hidden,\n const float* router_logits,\n const float* expert_ground,\n float* output,\n float alpha,\n int num_tokens,\n int hidden_dim,\n int num_experts,\n cudaStream_t stream) {\n \n cublasHandle_t handle = get_cublas_handle();\n cublasSetStream(handle, stream);\n \n // Allocate workspace for similarity scores\n float* similarity;\n cudaMalloc(&similarity, num_tokens * num_experts * sizeof(float));\n \n // Allocate workspace for top-2 results\n float* top_vals;\n int* top_idx;\n cudaMalloc(&top_vals, num_tokens * 2 * sizeof(float));\n cudaMalloc(&top_idx, num_tokens * 2 * sizeof(int));\n \n // Step 1: similarity = token_hidden @ expert_ground.T\n // token_hidden: (num_tokens, hidden_dim)\n // expert_ground: (num_experts, hidden_dim)\n // similarity: (num_tokens, num_experts)\n \n // cuBLAS Sgemm: C = alpha*A*B + beta*C\n // For C = A @ B.T, we use C = A @ B^T\n // In column-major: C(m,n) = A(m,k) @ B(k,n) with transB\n // A = token_hidden (row-major: num_tokens x hidden_dim) -> column-major: hidden_dim x num_tokens\n // B = expert_ground (row-major: num_experts x hidden_dim) -> column-major: hidden_dim x num_experts\n // C = similarity (row-major: num_tokens x num_experts) -> column-major: num_experts x num_tokens\n \n // We want: C = A @ B.T\n // In row-major: C[i,j] = sum_k A[i,k] * B[j,k]\n // In column-major: C[j,i] = sum_k A[k,i] * B[k,j]\n \n // For cuBLAS, we compute C^T = B @ A^T, then we have C in correct layout\n // Actually let's think again:\n // Row-major A: num_tokens x hidden_dim\n // Row-major B: num_experts x hidden_dim\n // Row-major C: num_tokens x num_experts\n // C = A @ B^T\n \n // cuBLAS is column-major\n // If we pass A as (num_tokens x hidden_dim) with transA=N, it's read as hidden_dim x num_tokens\n // If we pass B as (num_experts x hidden_dim) with transB=T, it's read as hidden_dim x num_experts (transposed)\n // Then C = A @ B gives hidden_dim x hidden_dim - wrong\n \n // Correct approach:\n // A is stored row-major as num_tokens x hidden_dim\n // B is stored row-major as num_experts x hidden_dim\n // C should be num_tokens x num_experts\n \n // In column-major terms:\n // A_col = A^T, shape hidden_dim x num_tokens\n // B_col = B^T, shape hidden_dim x num_experts\n // C_col should be num_experts x num_tokens\n \n // C_col = B_col^T @ A_col = B @ B^T ... no\n \n // Let's use: C^T_row = B_row @ A_row^T\n // B_row: num_experts x hidden_dim\n // A_row^T: hidden_dim x num_tokens\n // C^T_row: num_experts x num_tokens\n // C_row: num_tokens x num_experts\n \n // In column-major:\n // C_col (num_experts x num_tokens) = B_col (hidden_dim x num_experts)^T @ A_col (hidden_dim x num_tokens)\n // = op(B_col) @ A_col where op(B) = B^T\n // So: cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, num_tokens, num_experts, hidden_dim,\n // 1.0f, A, hidden_dim, B, hidden_dim, 0.0f, C, num_tokens)\n \n // Wait, let me be more careful.\n // For C_col(m_col, n_col) = op(A_col) @ op(B_col)\n // cublasSgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc)\n // m = rows of op(A), n = cols of op(B), k = cols of op(A) = rows of op(B)\n \n // We want C_col(num_experts, num_tokens) = B_col.T @ A_col\n // B_col is hidden_dim x num_experts, B_col.T is num_experts x hidden_dim\n // A_col is hidden_dim x num_tokens\n // Result: num_experts x num_tokens\n \n // op(B_col) = B_col.T, so transb = CUBLAS_OP_T\n // op(A_col) = A_col, so transa = CUBLAS_OP_N\n // m = num_experts, n = num_tokens, k = hidden_dim\n // ldb = hidden_dim (leading dim of B_col)\n // lda = hidden_dim (leading dim of A_col)\n // ldc = num_experts (leading dim of C_col)\n \n // But our arrays are in row-major. The pointer to row-major array can be interpreted as column-major.\n // Row-major A (num_tokens x hidden_dim) -> pointer starts at A[0,0], stride=hidden_dim\n // Column-major interpretation: first column has num_tokens elements with stride=1? No.\n // In column-major, elements in same column are adjacent. But in our row-major data, adjacent elements are in same row.\n \n // So if we interpret row-major A (num_tokens x hidden_dim) as column-major:\n // Column-major shape: hidden_dim x num_tokens\n // First column contains A[0,0], A[0,1], ..., A[0, hidden_dim-1] -> this is wrong\n \n // Actually, the stride tells the story:\n // Row-major A: A[i,j] is at position i*hidden_dim + j\n // Column-major interpretation: this is an array with leading dimension = hidden_dim\n // Shape: hidden_dim x num_tokens\n // Element at (row, col) = A_col[row, col] = data[col * hidden_dim + row]\n // But row-major A[i,j] = data[i * hidden_dim + j]\n // So A_col[j, i] = A_row[i, j]\n // Hence A_col = A_row.T\n \n // So:\n // token_hidden row-major (num_tokens, hidden_dim) -> column-major A_col = (hidden_dim, num_tokens) = token_hidden.T\n // expert_ground row-major (num_experts, hidden_dim) -> column-major B_col = (hidden_dim, num_experts) = expert_ground.T\n \n // We want similarity_row (num_tokens, num_experts) = token_hidden @ expert_ground.T\n // similarity_col (num_experts, num_tokens) = similarity_row.T\n // = (token_hidden @ expert_ground.T).T\n // = expert_ground @ token_hidden.T\n \n // expert_ground row-major (num_experts, hidden_dim)\n // token_hidden.T row-major (hidden_dim, num_tokens)\n // expert_ground @ token_hidden.T: (num_experts, num_tokens) \u2713\n \n // In column-major:\n // expert_ground row-major -> expert_ground_col (hidden_dim, num_experts) = expert_ground.T\n // token_hidden row-major -> token_hidden_col (hidden_dim, num_tokens) = token_hidden.T\n \n // We want: expert_ground @ token_hidden.T\n // In col-major: expert_ground_col.T @ token_hidden_col\n // = (hidden_dim x num_experts).T @ (hidden_dim x num_tokens)\n // = (num_experts x hidden_dim) @ (hidden_dim x num_tokens)\n // = (num_experts x num_tokens) \u2713\n \n // cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, num_tokens, num_experts, hidden_dim,\n // 1.0f, token_hidden, hidden_dim, expert_ground, hidden_dim, 0.0f, similarity, num_tokens)\n // Hmm, let me verify:\n // op(A) with transa=CUBLAS_OP_T: A is token_hidden, leading dim hidden_dim\n // op(A) has rows = num_tokens, cols = hidden_dim\n // So A_col has shape hidden_dim x num_tokens, and A_col.T has shape num_tokens x hidden_dim\n // op(B) with transb=CUBLAS_OP_N: B is expert_ground, leading dim hidden_dim\n // op(B) has rows = hidden_dim, cols = num_experts\n // So B_col has shape hidden_dim x num_experts\n // C = op(A) @ op(B) = (num_tokens x hidden_dim) @ (hidden_dim x num_experts) = num_tokens x num_experts\n // ldc = num_tokens\n \n // But we need ldc = rows of C = num_tokens? Let me check.\n // In column-major C, C[i,j] is at j * ldc + i\n // We're computing C with m rows (num_tokens) and n columns (num_experts)\n // So ldc >= m = num_tokens\n // And the result C_col has shape num_tokens x num_experts\n // But we want similarity_col to have shape num_experts x num_tokens...\n \n // OK I think the convention I'm using is:\n // Row-major similarity[num_tokens][num_experts]\n // The pointer to this data, when interpreted as column-major, gives:\n // similarity_col with shape num_experts x num_tokens\n \n // If I compute with cublasSgemm and store result in similarity with ldc = num_tokens:\n // The resulting column-major matrix has rows = m = num_tokens and cols = n = num_experts\n // This means similarity_col[num_experts * i + j] = ... wait\n \n // I think I need to step back. Let me use a simpler formula.\n // For row-major matrices A(m,k) and B(n,k), we want C(m,n) = A @ B.T\n // This is equivalent to: C_col.T = B_col @ A_col.T\n // Where A_col = A.T has shape (k,m), B_col = B.T has shape (k,n)\n // C_col.T = B_col.T @ A_col = B @ A.T... hmm\n \n // Alternative: Use cublasSgemm with proper arguments for row-major interpretation\n // For row-major A(m,k), B(n,k), we want C(m,n) = A @ B.T\n // Formula: cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, m, n, k, alpha, A, k, B, k, beta, C, n)\n // transa=N means op(A) = A (not transposed)\n // transb=T means op(B) = B.T\n // But A and B are column-major in cublas, so:\n // A_col(k,m) -> op(A) = A_col has shape k x m\n // B_col(k,n) -> op(B) = B_col.T has shape n x k\n // op(A) @ op(B) = (k x m) @ (n x k) doesn't work\n \n // Ugh. Let me just use the correct formula:\n // For row-major A(m,k), B(n,k), compute C(m,n) = A @ B.T\n // \n // The answer is: cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_T, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc)\n // where lda = k, ldb = k, ldc = n\n // \n // Reasoning:\n // A is stored row-major as (m,k), so in column-major it's (k,m) with leading dim k\n // B is stored row-major as (n,k), so in column-major it's (k,n) with leading dim k\n // \n // cublasSgemm computes C = op(A) @ op(B)\n // op(A) with transa=N: A_col(k,m), result is k x m\n // op(B) with transb=T: B_col(k,n).T, result is n x k\n // This gives k x m @ n x k which doesn't work...\n \n // Actually the convention is:\n // cublasSgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc)\n // computes C(m,n) = alpha * op(A) @ op(B) + beta * C\n // where m = rows of op(A), n = cols of op(B), k = cols of op(A) = rows of op(B)\n \n // So if transa=N, op(A) has same shape as A_col, which is k x m (since A_row(m,k) -> A_col(k,m))\n // Wait no. A_col has leading dimension lda. \n // A_col[j*lda + i] is the element at row i, col j\n // If A_row is (m,k) row-major, then A_row[i,j] = A_col[j,i] \n // So A_col has k rows and m columns, shape (k,m)\n // lda = k\n \n // For transa=N: op(A) = A_col, shape (k,m)\n // For transa=T: op(A) = A_col.T, shape (m,k)\n \n // So m = rows of op(A) means:\n // - If transa=N: m = k (which doesn't match our desired m)\n // - If transa=T: m = m (correct!)\n \n // For row-major A(m,k), B(n,k), we want C(m,n) = A @ B.T\n // op(A) should have m rows -> transa = T (so op(A) = A_col.T has shape m x k)\n // op(B) should have n cols -> transb = N (so op(B) = B_col has shape k x n)\n // op(A) @ op(B) = (m x k) @ (k x n) = (m x n) \u2713\n // m = m (first arg), n = n (second arg), k = k (third arg)\n // lda = k, ldb = k, ldc = n\n \n // So: cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, alpha, A, k, B, k, beta, C, n)\n // Let's double-check with our dimensions:\n // A = token_hidden (num_tokens, hidden_dim), so m = num_tokens, k = hidden_dim\n // B = expert_ground (num_experts, hidden_dim), so n = num_experts, k = hidden_dim\n // C = similarity (num_tokens, num_experts), so m = num_tokens, n = num_experts\n \n // cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, num_tokens, num_experts, hidden_dim,\n // 1.0f, token_hidden, hidden_dim, expert_ground, hidden_dim, 0.0f, similarity, num_experts)\n \n // But wait, ldc should be >= m = num_tokens, not num_experts!\n // If C_col has m rows and n cols, and ldc is the leading dimension,\n // then ldc >= m = num_tokens\n \n // Hmm, but I said C_row(m,n) -> C_col(n,m). Let me reconsider.\n // If C is stored row-major as (m,n), then C[i,j] = data[i*n + j]\n // Interpreted as column-major: C_col[j*ldc + i]\n // For these to match: C_col[j*ldc + i] should be C[i,j] = data[i*n + j]\n // So j*ldc + i = i*n + j\n // This means ldc = n? Then j*n + i = i*n + j... no, that's not right for all i,j.\n \n // Let me think differently. \n // Row-major data layout: memory offset = row * num_cols + col\n // Column-major data layout: memory offset = col * num_rows + row\n // \n // For the same memory, if we store as row-major (m,n), the array has m*n elements\n // Interpreted as column-major, we'd have:\n // memory[offset] interpreted as column-major element at (offset % num_rows_cm, offset / num_rows_cm)\n // where num_rows_cm is the leading dimension of the column-major matrix\n \n // If the leading dimension is ldc = m (same as row-major rows), then:\n // memory[i*n + j] -> column-major element at (j*m + i) % m, (j*m + i) / m = (i, j)\n // This means C_col[i,j] = C_row[i,j]\n // So the matrices are identical! (both m x n)\n \n // So if we allocate C as row-major (m,n) and use ldc = m in cublas:\n // cublasSgemm writes C as column-major with leading dimension m\n // This means C_col[i,j] is written to memory position j*m + i\n // And when we read as row-major, C_row[i,j] = memory[i*n + j] \n // These don't match in general!\n \n // For them to match, we need j*m + i = i*n + j for all i,j, which requires m=n.\n \n // So the conclusion is: we can't directly use row-major storage with cublas without being careful.\n // The simplest approach: let cublas write column-major, then transpose if needed.\n \n // Actually, a better approach: compute C_col = B_col.T @ A_col (as I was trying before)\n // C_col = (hidden_dim x num_experts).T @ (hidden_dim x num_tokens)\n // = (num_experts x hidden_dim) @ (hidden_dim x num_tokens)\n // = (num_experts x num_tokens)\n // This is similarity.T (as row-major would be num_tokens x num_experts)\n \n // So if we use ldc = num_experts, and cublas writes to C_col with shape num_experts x num_tokens,\n // then when we read as row-major, we get:\n // memory[j * num_experts + i] (column-major C_col[i,j])\n // = memory[j * num_experts + i]\n // interpreted as row-major C_row[?][?]\n // row-major offset = row * num_experts + col\n // So row = j, col = i -> C_row[j,i] = C_col[i,j]\n // Hence C_row = C_col.T = similarity.T.T = similarity \u2713\n \n // So the formula is:\n // cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, num_tokens, num_experts, hidden_dim,\n // 1.0f, expert_ground, hidden_dim, token_hidden, hidden_dim, 0.0f, similarity, num_experts)\n \n // Wait, let me redo this. We want:\n // C_col(num_experts, num_tokens) = expert_ground.T @ token_hidden\n // where expert_ground is (num_experts, hidden_dim) row-major\n // and token_hidden is (num_tokens, hidden_dim) row-major\n \n // In column-major:\n // expert_ground_col = expert_ground.T has shape (hidden_dim, num_experts)\n // token_hidden_col = token_hidden.T has shape (hidden_dim, num_tokens)\n \n // C_col = expert_ground.T @ token_hidden = expert_ground_col.T @ token_hidden_col\n // = (hidden_dim x num_experts).T @ (hidden_dim x num_tokens)\n // = (num_experts x hidden_dim) @ (hidden_dim x num_tokens)\n // = (num_experts x num_tokens)\n \n // cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, num_tokens, num_experts, hidden_dim, ...)\n // Wait, m = rows of op(A) where A = first matrix in cublas call\n // Let's say A = expert_ground (pointer to row-major data)\n // A_col has shape (hidden_dim, num_experts)\n // transa = T -> op(A) = A_col.T has shape (num_experts, hidden_dim)\n // So m = num_experts\n \n // B = token_hidden (pointer to row-major data)\n // B_col has shape (hidden_dim, num_tokens)\n // transb = N -> op(B) = B_col has shape (hidden_dim, num_tokens)\n // So n = num_tokens, k = hidden_dim\n \n // C = op(A) @ op(B) = (num_experts x hidden_dim) @ (hidden_dim x num_tokens) = num_experts x num_tokens \u2713\n \n // cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, num_tokens, num_experts, hidden_dim,\n // 1.0f, expert_ground, hidden_dim, token_hidden, hidden_dim, 0.0f, similarity, num_experts)\n // \n // Hmm wait, the arguments are (m, n, k) where m = rows of op(A), n = cols of op(B)\n // So m = num_experts, n = num_tokens, k = hidden_dim\n // \n // cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, num_experts, num_tokens, hidden_dim,\n // 1.0f, expert_ground, hidden_dim, token_hidden, hidden_dim, 0.0f, similarity, num_experts)\n //\n // ldc should be >= m = num_experts\n // The result is stored in column-major C with shape (num_experts, num_tokens)\n // When read as row-major (num_tokens, num_experts) with row stride = num_experts:\n // C_row[i,j] = memory[i * num_experts + j]\n // C_col[j,i] = memory[i * num_experts + j]\n // So C_row[i,j] = C_col[j,i]\n // Hence C_row = C_col.T\n // C_col = similarity.T (from our derivation)\n // C_row = similarity \u2713\n \n // Great! So the correct call is:\n // cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, num_experts, num_tokens, hidden_dim,\n // alpha, expert_ground, hidden_dim, token_hidden, hidden_dim, 0.0f, similarity, num_experts)\n \n // Compute similarity = alpha * token_hidden @ expert_ground.T\n const float gemm_alpha = alpha;\n const float gemm_beta = 0.0f;\n \n // similarity(num_tokens, num_experts) = alpha * token_hidden @ expert_ground.T\n // Using: C_col = expert_ground.T @ token_hidden (scaled by alpha)\n cublasStatus_t status = cublasSgemm(handle, \n CUBLAS_OP_T, CUBLAS_OP_N,\n num_experts, num_tokens, hidden_dim,\n &gemm_alpha,\n expert_ground, hidden_dim,\n token_hidden, hidden_dim,\n &gemm_beta,\n similarity, num_experts);\n \n if (status != CUBLAS_STATUS_SUCCESS) {\n printf(\"cuBLAS Sgemm failed with status %d\\\\n\", status);\n }\n \n // Step 2: Add router_logits to similarity\n // We need: grounded = similarity + router_logits (element-wise)\n // This is done in the top2 kernel by reading both arrays\n \n // Actually we need to add them first. Let me create a fused kernel.\n // For simplicity, let's use a separate kernel for the addition + top2\n \n // Step 3: Top-2 selection and softmax\n int blocks = (num_tokens + 255) / 256; // Each block handles 256 tokens\n dim3 threads(1, 256); // 256 threads in y dimension, each handling 1 token\n \n // Fused kernel: add router_logits, find top-2, and store results\n // We'll modify the top2 kernel to read both arrays\n // Actually, let me create a unified kernel for efficiency\n \n // For now, let's add router_logits to similarity in the top2 kernel\n // No, let's keep it simple with multiple kernels first\n \n // ...\n \n // Free workspace\n cudaFree(similarity);\n cudaFree(top_vals);\n cudaFree(top_idx);\n}\n\n// Fused kernel: compute grounded scores (similarity + logits), find top-2, softmax, output\n__global__ void fused_grounded_top2_kernel(\n const float* __restrict__ similarity,\n const float* __restrict__ router_logits,\n float* __restrict__ output,\n int num_tokens,\n int num_experts) {\n \n int token_idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (token_idx >= num_tokens) return;\n \n const float* sim_row = similarity + token_idx * num_experts;\n const float* logit_row = router_logits + token_idx * num_experts;\n \n // Find top-2 among grounded scores\n float max1 = -INFINITY, max2 = -INFINITY;\n int idx1 = 0, idx2 = 0;\n \n for (int e = 0; e < num_experts; e++) {\n float val = sim_row[e] + logit_row[e];\n if (val > max1) {\n max2 = max1;\n idx2 = idx1;\n max1 = val;\n idx1 = e;\n } else if (val > max2) {\n max2 = val;\n idx2 = e;\n }\n }\n \n // Softmax\n float max_val = fmaxf(max1, max2);\n float exp0 = expf(max1 - max_val);\n float exp1 = expf(max2 - max_val);\n float sum = exp0 + exp1;\n \n float w0 = exp0 / sum;\n float w1 = exp1 / sum;\n \n // Output\n output[token_idx * 4 + 0] = (float)idx1;\n output[token_idx * 4 + 1] = (float)idx2;\n output[token_idx * 4 + 2] = w0;\n output[token_idx * 4 + 3] = w1;\n}\n\n// Main entry point - called from Python\nvoid grounded_top2_moe_router_entry(\n const float* token_hidden,\n const float* router_logits,\n const float* expert_ground,\n float* output,\n float alpha,\n int num_tokens,\n int hidden_dim,\n int num_experts) {\n \n cudaStream_t stream = 0; // Default stream\n \n cublasHandle_t handle = get_cublas_handle();\n cublasSetStream(handle, stream);\n \n // Allocate workspace for similarity scores\n float* similarity;\n cudaMalloc(&similarity, num_tokens * num_experts * sizeof(float));\n \n // Compute similarity = alpha * token_hidden @ expert_ground.T\n const float gemm_alpha = alpha;\n const float gemm_beta = 0.0f;\n \n cublasSgemm(handle, \n CUBLAS_OP_T, CUBLAS_OP_N,\n num_experts, num_tokens, hidden_dim,\n &gemm_alpha,\n expert_ground, hidden_dim,\n token_hidden, hidden_dim,\n &gemm_beta,\n similarity, num_experts);\n \n // Fused kernel\n int threads_per_block = 256;\n int blocks = (num_tokens + threads_per_block - 1) / threads_per_block;\n \n fused_grounded_top2_kernel<<<blocks, threads_per_block>>>(\n similarity, router_logits, output, num_tokens, num_experts);\n \n // Free workspace\n cudaFree(similarity);\n}\n\"\"\"\n\ncpp_source = \"\"\"\nvoid grounded_top2_moe_router_entry(\n const float* token_hidden,\n const float* router_logits,\n const float* expert_ground,\n float* output,\n float alpha,\n int num_tokens,\n int hidden_dim,\n int num_experts);\n\"\"\"\n\nmodule = load_inline(\n name=\"grounded_top2_moe_cuda\",\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=[\"grounded_top2_moe_router_entry\"],\n extra_cuda_cflags=[\"-O3\", \"-arch=sm_90\"],\n extra_ldflags=[\"-lcublas\"],\n verbose=True\n)\n\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n self.cuda_module = module\n\n def forward(\n self,\n token_hidden: torch.Tensor,\n router_logits: torch.Tensor,\n expert_ground: torch.Tensor,\n alpha: float,\n ) -> torch.Tensor:\n num_tokens, hidden_dim = token_hidden.shape\n num_experts = router_logits.shape[1]\n \n # Ensure contiguous tensors\n token_hidden = token_hidden.contiguous()\n router_logits = router_logits.contiguous()\n expert_ground = expert_ground.contiguous()\n \n # Output tensor: [num_tokens, 2, 2]\n output = torch.empty(num_tokens, 2, 2, dtype=torch.float32, device=token_hidden.device)\n \n self.cuda_module.grounded_top2_moe_router_entry(\n token_hidden.data_ptr(),\n router_logits.data_ptr(),\n expert_ground.data_ptr(),\n output.data_ptr(),\n alpha,\n num_tokens,\n hidden_dim,\n num_experts\n )\n \n return output"
}