feat(cuda): optional arch-tiered PTX dispatch with compute_61 fallback

native/build.ps1

@@ -3,6 +3,30 @@
 # Output: native\ptx\*.ptx
 #
 # PTX is forward-compatible: compute_61 PTX runs on all GPUs from Pascal onward.
+#
+# ── Arch-tiered PTX (optional) ────────────────────────────────────────────────
+# In addition to the universal compute_61 "<kernel>.ptx" (always emitted), this
+# script can OPTIONALLY emit higher-arch PTX variants named "<kernel>.sm_<arch>.ptx"
+# for a curated subset of kernels. The runtime loader (CudaModule.LoadForArch)
+# picks the highest-arch variant whose arch is <= the device compute capability,
+# and falls back to the compute_61 "<kernel>.ptx" when no variant is present.
+#
+# This is opt-in. With no parameters the script produces EXACTLY today's output:
+# only compute_61 "<kernel>.ptx" files. To also emit higher-arch variants:
+#
+#   ./build.ps1 -ExtraArchs 80,86
+#   ./build.ps1 -ExtraArchs 80 -ExtraArchKernels quantized_gemv
+#
+# -ExtraArchs         SM numbers (e.g. 75,80,86,90). Empty = none.
+# -ExtraArchKernels   kernel base names to also build for -ExtraArchs. Defaults to
+#                     $archTieredKernels below. Only kernels with a genuinely
+#                     arch-specific implementation belong here; none exist yet, so
+#                     the default list is empty (true no-op).
+
+param(
+    [int[]] $ExtraArchs = @(),
+    [string[]] $ExtraArchKernels = $null
+)
 
 $ErrorActionPreference = "Stop"
 
@@ -14,11 +38,17 @@ if (-not (Test-Path $outDir)) { New-Item -ItemType Directory -Path $outDir | Out
 
 $arch = "compute_61"
 
+# Curated kernel list eligible for higher-arch variants. Empty until an
+# arch-specific kernel implementation actually exists.
+$archTieredKernels = @()
+if ($null -eq $ExtraArchKernels) { $ExtraArchKernels = $archTieredKernels }
+
 Write-Host "Compiling CUDA kernels -> PTX (target: $arch)..."
 
 foreach ($cuFile in Get-ChildItem "$kernelDir\*.cu") {
     $base = $cuFile.BaseName
 
+    # Universal compute_61 PTX — always emitted (today's behavior).
     & nvcc -ptx -arch=$arch `
         --use_fast_math `
         -o "$outDir\$base.ptx" `
@@ -29,6 +59,22 @@ foreach ($cuFile in Get-ChildItem "$kernelDir\*.cu") {
     }
 
     Write-Host "  $($cuFile.Name) -> $base.ptx"
+
+    # Optional higher-arch variants for the curated kernel list.
+    if ($ExtraArchs.Count -gt 0 -and $ExtraArchKernels -contains $base) {
+        foreach ($sm in $ExtraArchs) {
+            & nvcc -ptx -arch="compute_$sm" `
+                --use_fast_math `
+                -o "$outDir\$base.sm_$sm.ptx" `
+                $cuFile.FullName
+
+            if ($LASTEXITCODE -ne 0) {
+                throw "nvcc failed for $($cuFile.Name) (sm_$sm)"
+            }
+
+            Write-Host "  $($cuFile.Name) -> $base.sm_$sm.ptx (arch-tiered)"
+        }
+    }
 }
 
 Write-Host "Done. PTX files in $outDir\"

native/build.sh

@@ -5,6 +5,26 @@
 #
 # PTX is forward-compatible: compute_61 PTX runs on all GPUs from Pascal onward.
 # The CUDA driver JIT-compiles PTX → SASS for the specific GPU at load time.
+#
+# ── Arch-tiered PTX (optional) ────────────────────────────────────────────────
+# In addition to the universal compute_61 "<kernel>.ptx" (always emitted), this
+# script can OPTIONALLY emit higher-arch PTX variants named "<kernel>.sm_<arch>.ptx"
+# for a curated subset of kernels. The runtime loader (CudaModule.LoadForArch)
+# picks the highest-arch variant whose arch is <= the device compute capability,
+# and falls back to the compute_61 "<kernel>.ptx" when no variant is present.
+#
+# This is opt-in. With no extra arguments the script produces EXACTLY today's
+# output: only compute_61 "<kernel>.ptx" files. To also emit higher-arch variants:
+#
+#   EXTRA_ARCHS="80 86"            ./build.sh   # variants for the default kernel list
+#   EXTRA_ARCHS="80" \
+#   EXTRA_ARCH_KERNELS="quantized_gemv" ./build.sh
+#
+# EXTRA_ARCHS         space-separated SM numbers (e.g. "75 80 86 90"). Empty = none.
+# EXTRA_ARCH_KERNELS  space-separated kernel base names to also build for EXTRA_ARCHS.
+#                     Defaults to ARCH_TIERED_KERNELS below. Only kernels with a
+#                     genuinely arch-specific implementation belong here; none exist
+#                     yet, so the default list is empty (true no-op).
 
 set -e
 
@@ -18,6 +38,13 @@ mkdir -p "$OUT_DIR"
 # The driver will JIT to the actual GPU's native ISA at load time.
 ARCH="compute_61"
 
+# Optional higher-arch PTX variants (see header). Default: none → no-op.
+EXTRA_ARCHS="${EXTRA_ARCHS:-}"
+# Curated kernel list eligible for higher-arch variants. Empty until an
+# arch-specific kernel implementation actually exists.
+ARCH_TIERED_KERNELS=""
+EXTRA_ARCH_KERNELS="${EXTRA_ARCH_KERNELS:-$ARCH_TIERED_KERNELS}"
+
 # Kernels where --use_fast_math is safe (element-wise ops, no precision-sensitive math):
 FAST_MATH_KERNELS="add add_f32 swiglu swiglu_f32 convert bias_add bias_add_f32 embedding embedding_f32out dequant quant_kv"
 
@@ -36,24 +63,47 @@ is_fast_math_kernel() {
     return 1
 }
 
+is_in_list() {
+    local name="$1"; shift
+    for item in $@; do
+        [ "$item" = "$name" ] && return 0
+    done
+    return 1
+}
+
+# compile <cu_file> <arch> <out_ptx>
+compile_ptx() {
+    local cu_file="$1" arch="$2" out_ptx="$3"
+    local base
+    base=$(basename "$cu_file" .cu)
+    if is_fast_math_kernel "$base"; then
+        nvcc -ptx -arch="$arch" --use_fast_math -o "$out_ptx" "$cu_file"
+    else
+        nvcc -ptx -arch="$arch" -o "$out_ptx" "$cu_file"
+    fi
+}
+
 echo "Compiling CUDA kernels → PTX (target: $ARCH)..."
 
 for cu_file in "$KERNEL_DIR"/*.cu; do
     [ -f "$cu_file" ] || continue
     base=$(basename "$cu_file" .cu)
 
+    # Universal compute_61 PTX — always emitted (today's behavior).
+    compile_ptx "$cu_file" "$ARCH" "$OUT_DIR/$base.ptx"
     if is_fast_math_kernel "$base"; then
-        nvcc -ptx -arch="$ARCH" \
-             --use_fast_math \
-             -o "$OUT_DIR/$base.ptx" \
-             "$cu_file"
         echo "  $base.cu → $base.ptx (fast_math)"
     else
-        nvcc -ptx -arch="$ARCH" \
-             -o "$OUT_DIR/$base.ptx" \
-             "$cu_file"
         echo "  $base.cu → $base.ptx (precise)"
     fi
+
+    # Optional higher-arch variants for the curated kernel list.
+    if [ -n "$EXTRA_ARCHS" ] && is_in_list "$base" $EXTRA_ARCH_KERNELS; then
+        for sm in $EXTRA_ARCHS; do
+            compile_ptx "$cu_file" "compute_$sm" "$OUT_DIR/$base.sm_$sm.ptx"
+            echo "  $base.cu → $base.sm_$sm.ptx (arch-tiered)"
+        done
+    fi
 done
 
 echo "Done. PTX files in $OUT_DIR/"

src/DotLLM.Cuda/CudaKernels.cs

@@ -90,31 +90,45 @@ public sealed unsafe class CudaKernels : IDisposable
     /// Loads all PTX modules from the specified directory.
     /// </summary>
     /// <param name="ptxDir">Directory containing compiled .ptx files.</param>
-    public CudaKernels(string ptxDir)
+    /// <param name="ccMajor">
+    /// Device compute capability major version. When supplied (together with
+    /// <paramref name="ccMinor"/>), arch-tiered PTX variants named
+    /// <c>&lt;kernel&gt;.sm_&lt;arch&gt;.ptx</c> are preferred when shipped and <c>&lt;=</c> the device
+    /// arch; otherwise the universal <c>compute_61</c> <c>&lt;kernel&gt;.ptx</c> is used. The default
+    /// of <c>0</c> always selects the universal fallback (today's behavior).
+    /// </param>
+    /// <param name="ccMinor">Device compute capability minor version. See <paramref name="ccMajor"/>.</param>
+    public CudaKernels(string ptxDir, int ccMajor = 0, int ccMinor = 0)
     {
-        _rmsnormModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "rmsnorm.ptx"));
-        _ropeModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "rope.ptx"));
-        _swigluModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "swiglu.ptx"));
-        _addModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "add.ptx"));
-        _softmaxModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "softmax.ptx"));
-        _embeddingModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "embedding.ptx"));
-        _attentionModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "attention.ptx"));
-        _biasAddModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "bias_add.ptx"));
-        _perHeadRmsNormModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "per_head_rmsnorm.ptx"));
-        _convertModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "convert.ptx"));
-        _dequantModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "dequant.ptx"));
-        _quantizedGemvModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "quantized_gemv.ptx"));
-        _fusedAddRmsNormModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "fused_add_rmsnorm.ptx"));
-        _rmsnormF32InModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "rmsnorm_f32in.ptx"));
-        _addF32Module = CudaModule.LoadFromFile(Path.Combine(ptxDir, "add_f32.ptx"));
-        _embeddingF32OutModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "embedding_f32out.ptx"));
-        _ropeF32Module = CudaModule.LoadFromFile(Path.Combine(ptxDir, "rope_f32.ptx"));
-        _attentionF32Module = CudaModule.LoadFromFile(Path.Combine(ptxDir, "attention_f32.ptx"));
-        _swigluF32Module = CudaModule.LoadFromFile(Path.Combine(ptxDir, "swiglu_f32.ptx"));
-        _biasAddF32Module = CudaModule.LoadFromFile(Path.Combine(ptxDir, "bias_add_f32.ptx"));
-        _perHeadRmsNormF32Module = CudaModule.LoadFromFile(Path.Combine(ptxDir, "per_head_rmsnorm_f32.ptx"));
-        _rmsnormF32Module = CudaModule.LoadFromFile(Path.Combine(ptxDir, "rmsnorm_f32.ptx"));
-        _quantizedGemvF32InModule = CudaModule.LoadFromFile(Path.Combine(ptxDir, "quantized_gemv_f32in.ptx"));
+        // Arch-aware load: prefer a shipped "<kernel>.sm_<arch>.ptx" variant whose arch
+        // is <= the device compute capability; fall back to the universal compute_61
+        // "<kernel>.ptx". No-op (identical to LoadFromFile) when no variants are shipped
+        // or when ccMajor/ccMinor are 0.
+        CudaModule Load(string fileName) => CudaModule.LoadForArch(ptxDir, fileName, ccMajor, ccMinor);
+
+        _rmsnormModule = Load("rmsnorm.ptx");
+        _ropeModule = Load("rope.ptx");
+        _swigluModule = Load("swiglu.ptx");
+        _addModule = Load("add.ptx");
+        _softmaxModule = Load("softmax.ptx");
+        _embeddingModule = Load("embedding.ptx");
+        _attentionModule = Load("attention.ptx");
+        _biasAddModule = Load("bias_add.ptx");
+        _perHeadRmsNormModule = Load("per_head_rmsnorm.ptx");
+        _convertModule = Load("convert.ptx");
+        _dequantModule = Load("dequant.ptx");
+        _quantizedGemvModule = Load("quantized_gemv.ptx");
+        _fusedAddRmsNormModule = Load("fused_add_rmsnorm.ptx");
+        _rmsnormF32InModule = Load("rmsnorm_f32in.ptx");
+        _addF32Module = Load("add_f32.ptx");
+        _embeddingF32OutModule = Load("embedding_f32out.ptx");
+        _ropeF32Module = Load("rope_f32.ptx");
+        _attentionF32Module = Load("attention_f32.ptx");
+        _swigluF32Module = Load("swiglu_f32.ptx");
+        _biasAddF32Module = Load("bias_add_f32.ptx");
+        _perHeadRmsNormF32Module = Load("per_head_rmsnorm_f32.ptx");
+        _rmsnormF32Module = Load("rmsnorm_f32.ptx");
+        _quantizedGemvF32InModule = Load("quantized_gemv_f32in.ptx");
 
         _rmsnormFunc = _rmsnormModule.GetFunction("rmsnorm_f16");
         _rmsnormF32Func = _rmsnormF32Module.GetFunction("rmsnorm_f32");
@@ -156,10 +170,9 @@ public CudaKernels(string ptxDir)
         _dequantQ6_KFunc = _dequantModule.GetFunction("dequant_q6_k_f16");
 
         // KV-cache quantization (optional — PTX may not be compiled yet)
-        string quantKvPath = Path.Combine(ptxDir, "quant_kv.ptx");
-        if (File.Exists(quantKvPath))
+        if (File.Exists(Path.Combine(ptxDir, "quant_kv.ptx")))
         {
-            _quantKvModule = CudaModule.LoadFromFile(quantKvPath);
+            _quantKvModule = Load("quant_kv.ptx");
             _quantKvQ8_0Func = _quantKvModule.GetFunction("quant_f16_to_q8_0");
             _quantKvQ4_0Func = _quantKvModule.GetFunction("quant_f16_to_q4_0");
         }

src/DotLLM.Cuda/CudaModule.cs

@@ -22,6 +22,89 @@ public static CudaModule LoadFromFile(string ptxPath)
         return LoadFromBytes(ptxBytes);
     }
 
+    /// <summary>
+    /// Resolves the best-matching PTX variant for a kernel given the device compute
+    /// capability, then loads it.
+    /// </summary>
+    /// <remarks>
+    /// Looks for arch-tiered variants named <c>&lt;kernel&gt;.sm_&lt;arch&gt;.ptx</c> (e.g.
+    /// <c>rmsnorm.sm_80.ptx</c>) alongside the universal <c>&lt;kernel&gt;.ptx</c> built for
+    /// <c>compute_61</c>. The highest-arch variant whose architecture is &lt;= the device's
+    /// compute capability is selected; if none is present, the plain <c>compute_61</c>
+    /// <c>&lt;kernel&gt;.ptx</c> is used as the universal fallback. When no higher-arch
+    /// variants are shipped this is a no-op and resolves to the same file as
+    /// <see cref="LoadFromFile"/>.
+    /// </remarks>
+    /// <param name="ptxDir">Directory containing compiled .ptx files.</param>
+    /// <param name="baseFileName">Base kernel file name including extension (e.g. <c>rmsnorm.ptx</c>).</param>
+    /// <param name="ccMajor">Device compute capability major version.</param>
+    /// <param name="ccMinor">Device compute capability minor version.</param>
+    public static CudaModule LoadForArch(string ptxDir, string baseFileName, int ccMajor, int ccMinor)
+        => LoadFromFile(ResolveArchVariantPath(ptxDir, baseFileName, ccMajor, ccMinor));
+
+    /// <summary>
+    /// Selects the best-matching arch-tiered PTX variant path for a kernel without loading it.
+    /// Pure file-system logic — exposed for unit testing of variant selection.
+    /// </summary>
+    /// <remarks>
+    /// Given <paramref name="baseFileName"/> = <c>foo.ptx</c>, candidate variants are
+    /// <c>foo.sm_&lt;arch&gt;.ptx</c> where <c>&lt;arch&gt;</c> is a two- or three-digit SM number
+    /// (e.g. <c>75</c>, <c>80</c>, <c>86</c>, <c>90</c>). The variant with the highest arch value
+    /// that does not exceed the device compute capability (<c>ccMajor * 10 + ccMinor</c>) and that
+    /// actually exists on disk is returned. If no such variant exists, the plain
+    /// <paramref name="baseFileName"/> path (the <c>compute_61</c> universal build) is returned
+    /// unchanged.
+    /// </remarks>
+    /// <param name="ptxDir">Directory containing compiled .ptx files.</param>
+    /// <param name="baseFileName">Base kernel file name including extension (e.g. <c>rmsnorm.ptx</c>).</param>
+    /// <param name="ccMajor">Device compute capability major version.</param>
+    /// <param name="ccMinor">Device compute capability minor version.</param>
+    /// <returns>The full path of the variant to load.</returns>
+    public static string ResolveArchVariantPath(string ptxDir, string baseFileName, int ccMajor, int ccMinor)
+    {
+        string basePath = Path.Combine(ptxDir, baseFileName);
+
+        int deviceArch = (ccMajor * 10) + ccMinor;
+        if (deviceArch <= 0)
+            return basePath;
+
+        // "foo.ptx" -> stem "foo", suffix ".ptx"
+        string suffix = Path.GetExtension(baseFileName);              // ".ptx"
+        string stem = baseFileName[..^suffix.Length];                 // "foo"
+
+        // Enumerate sibling variants "foo.sm_<arch>.ptx" and pick the highest arch <= deviceArch.
+        string prefix = stem + ".sm_";
+        int bestArch = -1;
+        string bestPath = basePath;
+
+        IEnumerable<string> candidates;
+        try
+        {
+            candidates = Directory.EnumerateFiles(ptxDir, prefix + "*" + suffix);
+        }
+        catch (DirectoryNotFoundException)
+        {
+            return basePath;
+        }
+
+        foreach (string candidate in candidates)
+        {
+            string fileName = Path.GetFileName(candidate);
+            // Strip "<stem>.sm_" prefix and ".ptx" suffix to isolate the arch token.
+            string archToken = fileName[prefix.Length..^suffix.Length];
+            if (!int.TryParse(archToken, out int arch))
+                continue; // ignore malformed tokens (e.g. "foo.sm_80a.ptx")
+
+            if (arch <= deviceArch && arch > bestArch)
+            {
+                bestArch = arch;
+                bestPath = candidate;
+            }
+        }
+
+        return bestPath;
+    }
+
     /// <summary>
     /// Loads a PTX module from a byte array (UTF-8 text with null terminator).
     /// </summary>

src/DotLLM.Cuda/CudaTransformerModel.cs

@@ -87,7 +87,10 @@ public static CudaTransformerModel LoadFromGguf(GgufFile gguf, ModelConfig confi
 
         // Resolve PTX directory
         ptxDir ??= Path.Combine(AppContext.BaseDirectory, "ptx");
-        var kernels = new CudaKernels(ptxDir);
+        // Pass device compute capability so arch-tiered PTX variants can be selected
+        // when shipped; falls back to the universal compute_61 PTX otherwise.
+        var device = CudaDevice.GetDevice(deviceId);
+        var kernels = new CudaKernels(ptxDir, device.ComputeCapabilityMajor, device.ComputeCapabilityMinor);
 
         // Check VRAM before loading — warn if model likely exceeds available memory.
         // Estimate: sum of quantized byte sizes for all GGUF tensors.

src/DotLLM.Cuda/HybridTransformerModel.cs

@@ -134,7 +134,10 @@ public static HybridTransformerModel LoadFromGguf(
         cublas.SetStream(stream);
 
         string? ptxDir = Path.Combine(AppContext.BaseDirectory, "ptx");
-        var kernels = new CudaKernels(ptxDir);
+        // Pass device compute capability so arch-tiered PTX variants can be selected
+        // when shipped; falls back to the universal compute_61 PTX otherwise.
+        var device = CudaDevice.GetDevice(deviceId);
+        var kernels = new CudaKernels(ptxDir, device.ComputeCapabilityMajor, device.ComputeCapabilityMinor);
 
         // 3. Upload only GPU layers to VRAM
         var gpuWeights = CudaWeights.LoadFromGguf(cpuWeights, config, kernels, stream.Handle, numGpuLayers);