cuda: RMS norm rows and token rows for every family, Nemotron's embed the first user

zig/build/cuda.zig

@@ -36,6 +36,8 @@ const kernels = [_]Kernel{
     .{ .name = "nemotron_mamba", .flags = glue },
     .{ .name = "nemotron_attention", .flags = glue },
     .{ .name = "nemotron_keyed", .flags = glue },
+    .{ .name = "norms", .flags = &.{"-O3"} }, // ours: RMS norm rows for every family (cuda/norms.zig)
+    .{ .name = "embed_rows", .flags = &.{"-O3"} }, // ours: 4-bit affine and bf16 token rows (cuda/embed_rows.zig)
     .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3
     .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1
     .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts
@@ -125,6 +127,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu
     b.installArtifact(b.addExecutable(.{ .name = "tensorfold", .root_module = cli }));
     const runner = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip });
     runner.addImport("cuda", cuda);
+    runner.addImport("core", mods.core);
     runner.addImport("lanes", mods.lanes);
     runner.addImport("nemotron", mods.nemotron);
     const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core);

zig/kernels/cuda/embed_rows.cu

@@ -0,0 +1,44 @@
+// Token rows for every CUDA family (float64 reference: core/act_ref.zig's affine4): 4-bit affine or bf16 tables.
+
+#include <cuda_bf16.h>
+#include <stdint.h>
+
+namespace tf_embed {
+
+constexpr int NT = 256;
+
+}  // namespace tf_embed
+
+using namespace tf_embed;
+
+// out[r, i] = bf16(q * s + b) of token ids[r]: q 4 bits (8 a word, low first), s and b bf16 a group of 64; d % 64 == 0.
+extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_affine4(const int* __restrict__ ids,
+                                                                       const uint32_t* __restrict__ w,
+                                                                       const __nv_bfloat16* __restrict__ s,
+                                                                       const __nv_bfloat16* __restrict__ b,
+                                                                       __nv_bfloat16* __restrict__ out, int d) {
+    const int64_t r = blockIdx.x;
+    const int c = blockIdx.y * NT + threadIdx.x; // 8 values, one word
+    if (8 * c >= d) return;
+    const int64_t tok = ids[r];
+    const uint32_t word = w[tok * (d / 8) + c];
+    const int64_t g = tok * (d / 64) + c / 8;
+    const float sc = __bfloat162float(s[g]), bi = __bfloat162float(b[g]);
+    uint4 o;
+    __nv_bfloat162* h = reinterpret_cast<__nv_bfloat162*>(&o);
+#pragma unroll
+    for (int i = 0; i < 4; ++i)
+        h[i] = __floats2bfloat162_rn(__fmaf_rn(static_cast<float>((word >> (8 * i)) & 0xFu), sc, bi),
+                                     __fmaf_rn(static_cast<float>((word >> (8 * i + 4)) & 0xFu), sc, bi));
+    *reinterpret_cast<uint4*>(out + r * d + 8 * c) = o;
+}
+
+// out[r] = w[ids[r]] for a bf16 table, 16 bytes a thread (d % 8 == 0).
+extern "C" __global__ void __launch_bounds__(NT) tf_embed_rows_bf16(const int* __restrict__ ids,
+                                                                    const __nv_bfloat16* __restrict__ w,
+                                                                    __nv_bfloat16* __restrict__ out, int d) {
+    const int64_t r = blockIdx.x;
+    const int c = blockIdx.y * NT + threadIdx.x;
+    if (8 * c >= d) return;
+    *reinterpret_cast<uint4*>(out + r * d + 8 * c) = *reinterpret_cast<const uint4*>(w + static_cast<int64_t>(ids[r]) * d + 8 * c);
+}

zig/kernels/cuda/nemotron_norms.cu

@@ -52,19 +52,6 @@ __device__ __forceinline__ void scale_store(const float (&x)[MAXE], int D, float
 
 using namespace tf_glue;
 
-// MLX affine 4-bit token rows: out = bf16(fma(q, scale, bias)), D / 8 words and D / 64 scales a token.
-extern "C" __global__ void __launch_bounds__(NT) tf_nemo_embed(const int* __restrict__ ids, const uint32_t* __restrict__ W,
-                                                                 const __nv_bfloat16* __restrict__ S, const __nv_bfloat16* __restrict__ B,
-                                                                 __nv_bfloat16* __restrict__ out, int D) {
-    const int64_t tok = ids[blockIdx.x];
-    for (int i = threadIdx.x; i < D; i += NT) {
-        const uint32_t word = W[tok * (D / 8) + i / 8];
-        const float q = static_cast<float>((word >> (4 * (i % 8))) & 0xFu);
-        const int64_t g = tok * (D / 64) + i / 64;
-        out[static_cast<int64_t>(blockIdx.x) * D + i] = __float2bfloat16_rn(__fmaf_rn(q, bf(S[g]), bf(B[g])));
-    }
-}
-
 // A bf16 token table's rows (ModelOpt checkpoints keep the embedding unquantized): out = W[ids].
 extern "C" __global__ void __launch_bounds__(NT) tf_nemo_embed16(const int* __restrict__ ids, const __nv_bfloat16* __restrict__ W,
                                                                    __nv_bfloat16* __restrict__ out, int D) {