diff --git a/dsv4/kernels/cuda/mhc_sinkhorn.cu b/dsv4/kernels/cuda/mhc_sinkhorn.cu index f688d4b1..834f7326 100644 --- a/dsv4/kernels/cuda/mhc_sinkhorn.cu +++ b/dsv4/kernels/cuda/mhc_sinkhorn.cu @@ -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 @@ -20,9 +25,12 @@ #include #include -// 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<<>>( logits.data_ptr(), out.data_ptr(), 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)"); } diff --git a/dsv4/layers/mhc.py b/dsv4/layers/mhc.py index eb95a2a9..38133693 100644 --- a/dsv4/layers/mhc.py +++ b/dsv4/layers/mhc.py @@ -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) # ---------------------------------------------------------------------------