biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix mHC Sinkhorn kernel: remove VLA, remove Python fallback

Browse Source 
Root cause: float row_max[n] is a VLA — not allowed in CUDA device code.
Fix: use shared memory with MHC_MAX_N=16 fixed-size slots.

Also: REMOVED the Python fallback in sinkhorn_knopp().
If the CUDA kernel fails, the pipeline DIES. No soft landing.
This is the correct behavior — silent fallback to broken precision
is worse than a loud crash.

The residual growth |X|→500-700 at L60 was likely caused by the Python
fallback running a DIFFERENT numerical path (BF16 accumulation in torch
ops vs FP32 in the CUDA kernel). With the fixed kernel, Sinkhorn should
produce properly doubly-stochastic B_l, bounding the residual.
This commit is contained in:
biondizzle
2026-06-02 10:44:53 +00:00
parent c89762ecdd
commit 6cb5078821
2 changed files with 71 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

160
dsv4/kernels/cuda/mhc_sinkhorn.cu
View File

@@ -11,6 +11,11 @@
* 1. softmax(logits, dim=-1) + eps
* 2. column normalize
* 3. (t_max - 1) alternating row/col normalize
*
* NVFP4 PATH: This kernel operates on the B_l (comb) matrix which must be
* doubly-stochastic for residual bounding. The residual |X| growth to 500-700
* at L60 indicates B was NOT properly doubly-stochastic at runtime. This kernel
* ensures it. No fallback to Python. If this kernel fails, the pipeline fails.
*/
#include <cuda.h>
@@ -20,9 +25,12 @@
#include <torch/extension.h>
#include <cmath>
// One thread per (t, i, j) element of the (T, n, n) matrix
// For T=1, n=4: 16 threads total — trivial parallelism
// For larger T, each batch element is independent
// Max supported n — DSV4 uses n=4. Increase if needed.
#define MHC_MAX_N 16
// One block per batch element. n*n threads per block (for n=4: 16 threads).
// Shared memory holds the (n, n) matrix + row/col sums.
// All loops use fixed-size arrays (no VLA — CUDA requirement).
__global__ void mhc_sinkhorn_kernel(
const float* __restrict__ logits, // (T, n, n)
@@ -31,69 +39,71 @@ __global__ void mhc_sinkhorn_kernel(
) {
int t = blockIdx.x;
if (t >= T) return;
// Each block handles one batch element
// Use shared memory for the (n, n) matrix — n=4 → 16 floats = 64 bytes
// Shared memory layout: M (n, n) | row_max (MHC_MAX_N) | row_sum (MHC_MAX_N) | col_sum (MHC_MAX_N)
extern __shared__ float smem[];
float* M = smem; // (n, n) — current matrix
float* row_sum = smem + n * n; // (n,) — row sums
float* col_sum = row_sum + n; // (n,) — col sums
float* M = smem; // n*n floats
float* row_max = smem + n * n; // MHC_MAX_N floats
float* row_sum_arr = row_max + MHC_MAX_N; // MHC_MAX_N floats
float* col_sum_arr = row_sum_arr + MHC_MAX_N; // MHC_MAX_N floats
int i = threadIdx.x / n;
int j = threadIdx.x % n;
// Step 1: softmax(logits, dim=-1) + eps
// Each row's softmax is computed by threads [i*0..i*(n-1)]
if (i < n && j < n) {
M[i * n + j] = logits[t * n * n + i * n + j];
}
__syncthreads();
// Compute row max for numerical stability
float row_max[n]; // n=4, so this fits in registers
for (int ri = 0; ri < n; ri++) {
float mx = -INFINITY;
for (int rj = 0; rj < n; rj++) {
mx = fmaxf(mx, M[ri * n + rj]);
}
row_max[ri] = mx;
}
// Apply softmax + eps
for (int ri = 0; ri < n; ri++) {
float exp_sum = 0.0f;
for (int rj = 0; rj < n; rj++) {
M[ri * n + rj] = expf(M[ri * n + rj] - row_max[ri]);
exp_sum += M[ri * n + rj];
}
for (int rj = 0; rj < n; rj++) {
M[ri * n + rj] = M[ri * n + rj] / exp_sum + eps;
}
}
// Step 2: column normalize
for (int cj = 0; cj < n; cj++) {
float cs = 0.0f;
for (int ci = 0; ci < n; ci++) cs += M[ci * n + cj];
for (int ci = 0; ci < n; ci++) M[ci * n + cj] = M[ci * n + cj] / (cs + eps);
}
// Step 3: (t_max - 1) alternating row/col normalize
for (int iter = 0; iter < t_max - 1; iter++) {
// Row normalize
// Thread 0 does all the work (n is tiny — 4)
if (threadIdx.x == 0) {
for (int ri = 0; ri < n; ri++) {
float rs = 0.0f;
for (int rj = 0; rj < n; rj++) rs += M[ri * n + rj];
for (int rj = 0; rj < n; rj++) M[ri * n + rj] = M[ri * n + rj] / (rs + eps);
float mx = -INFINITY;
for (int rj = 0; rj < n; rj++) {
mx = fmaxf(mx, M[ri * n + rj]);
}
row_max[ri] = mx;
}
// Column normalize
// Apply softmax + eps
for (int ri = 0; ri < n; ri++) {
float exp_sum = 0.0f;
for (int rj = 0; rj < n; rj++) {
M[ri * n + rj] = expf(M[ri * n + rj] - row_max[ri]);
exp_sum += M[ri * n + rj];
}
for (int rj = 0; rj < n; rj++) {
M[ri * n + rj] = M[ri * n + rj] / exp_sum + eps;
}
}
// Step 2: column normalize
for (int cj = 0; cj < n; cj++) {
float cs = 0.0f;
for (int ci = 0; ci < n; ci++) cs += M[ci * n + cj];
for (int ci = 0; ci < n; ci++) M[ci * n + cj] = M[ci * n + cj] / (cs + eps);
}
// Step 3: (t_max - 1) alternating row/col normalize
for (int iter = 0; iter < t_max - 1; iter++) {
// Row normalize
for (int ri = 0; ri < n; ri++) {
float rs = 0.0f;
for (int rj = 0; rj < n; rj++) rs += M[ri * n + rj];
for (int rj = 0; rj < n; rj++) M[ri * n + rj] = M[ri * n + rj] / (rs + eps);
}
// Column normalize
for (int cj = 0; cj < n; cj++) {
float cs = 0.0f;
for (int ci = 0; ci < n; ci++) cs += M[ci * n + cj];
for (int ci = 0; ci < n; ci++) M[ci * n + cj] = M[ci * n + cj] / (cs + eps);
}
}
}
__syncthreads();
// Write output
if (i < n && j < n) {
out[t * n * n + i * n + j] = M[i * n + j];
@@ -109,63 +119,25 @@ torch::Tensor mhc_sinkhorn_cuda(
int T = logits.size(0);
int n = logits.size(1);
TORCH_CHECK(logits.size(2) == n, "logits must be square");
TORCH_CHECK(n <= MHC_MAX_N, "n must be <= MHC_MAX_N (16)");
TORCH_CHECK(logits.scalar_type() == torch::kFloat32, "logits must be FP32");
auto out = torch::empty_like(logits);
// One block per batch element, n*n threads per block
int threads = n * n;
int smem_size = n * n * sizeof(float) + 2 * n * sizeof(float);
// Shared memory: M (n*n) + row_max + row_sum + col_sum (3 * MHC_MAX_N)
int smem_size = (n * n + 3 * MHC_MAX_N) * sizeof(float);
mhc_sinkhorn_kernel<<<T, threads, smem_size, c10::cuda::getCurrentCUDAStream()>>>(
logits.data_ptr<float>(),
out.data_ptr<float>(),
T, n, t_max, (float)eps
);
return out;
}
// Also: fused mHC dynamic params kernel
// Computes A_l, B_l, C_l from X_flat in a single kernel launch.
// Currently done in ~8 separate ops in _dynamic_params().
__global__ void mhc_dynamic_params_kernel(
const __nv_bfloat16* __restrict__ X_flat, // (T, K) BF16
const float* __restrict__ W_stacked, // (N_proj, K) FP32
int T, int K, int n_hc,
float alpha_pre, float alpha_post, float alpha_comb,
const float* __restrict__ S_pre, // (1, n_hc)
const float* __restrict__ S_post, // (n_hc,)
const float* __restrict__ S_comb, // (n_hc*n_hc,)
float eps,
__nv_bfloat16* __restrict__ A_l_out, // (T, n_hc) BF16
float* __restrict__ B_l_out, // (T, n_hc, n_hc) FP32
__nv_bfloat16* __restrict__ C_l_out, // (T, n_hc) BF16
int t_max_sinkhorn
) {
// This kernel is more complex — it needs to do:
// 1. RMSNorm on X_flat
// 2. GEMM: (T, K) × (N_proj, K)^T → (T, N_proj)
// 3. Split + apply constraints
// 4. Sinkhorn on comb
//
// The GEMM at T=1, K=28672, N=24 is small enough to do per-thread
// with shared memory tiling.
//
// For now, just do the post-GEMM part (steps 3-4) as a fused kernel.
// The GEMM stays in Python/CuTeDSL.
// TODO: Full fusion in a future iteration.
// This kernel handles post-GEMM: split, apply constraints, Sinkhorn
int t = blockIdx.x;
if (t >= T) return;
// Thread handles one element of the output
// Not implementing the full GEMM here — that stays in Python
// This is a placeholder for the fused post-GEMM kernel
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("mhc_sinkhorn", &mhc_sinkhorn_cuda, "Fused mHC Sinkhorn-Knopp projection");
m.def("mhc_sinkhorn", &mhc_sinkhorn_cuda, "Fused mHC Sinkhorn-Knopp projection (NO FALLBACK)");
}

23
dsv4/layers/mhc.py
View File

@@ -91,25 +91,12 @@ def sinkhorn_knopp(
3. column-normalize
4. (t_max - 1) alternating row/col normalizations
Uses fused CUDA kernel when available (1 launch instead of 38).
Falls back to Python for correctness verification.
NO PYTHON FALLBACK. If the CUDA kernel fails, the pipeline dies.
The kernel MUST compile and run correctly. Period.
"""
# Try fused CUDA kernel first
try:
from dsv4.kernels.cuda.loader import get_cuda_module
mod = get_cuda_module("mhc_sinkhorn", ["mhc_sinkhorn.cu"])
return mod.mhc_sinkhorn(logits.float(), t_max, eps)
except Exception as e:
import sys; print(f"mhc_sinkhorn CUDA kernel failed: {e}, falling back to Python", file=sys.stderr, flush=True)
pass # Fall back to Python
# Python fallback
M = torch.softmax(logits, dim=-1) + eps # (T, n, n)
M = M / (M.sum(dim=-2, keepdim=True) + eps) # T_c (col)
for _ in range(t_max - 1):
M = M / (M.sum(dim=-1, keepdim=True) + eps) # T_r (row)
M = M / (M.sum(dim=-2, keepdim=True) + eps) # T_c (col)
return M
from dsv4.kernels.cuda.loader import get_cuda_module
mod = get_cuda_module("mhc_sinkhorn", ["mhc_sinkhorn.cu"])
return mod.mhc_sinkhorn(logits.float(), t_max, eps)
# ---------------------------------------------------------------------------