Add one-command B200 optimization suite

Author: shiloz-stack

README.md

@@ -2,7 +2,7 @@
 
 A kernel-level exploration of **persistent expert-tile scheduling** and **workload-aware dispatch** for irregular Mixture-of-Experts (MoE) inference on NVIDIA Blackwell GPUs.
 
-> **Status:** Work in progress. The CUTLASS BF16 grouped baseline is validated on NVIDIA B200. An independent SM100a-native dense TMA/tcgen05/TMEM path is now implemented and awaiting its B200 correctness/performance gate.
+> **Status:** Work in progress. The CUTLASS BF16 grouped baseline is validated on NVIDIA B200. Direct CuTe, 1-SM/2-SM native collectives, and project-owned GPU scheduler probes compile for SM100a in CI and await their B200 correctness/performance gate.
 
 ## Motivation
 
@@ -140,9 +140,9 @@ All benchmark reports will record GPU model, clocks or power mode when relevant,
 - [x] Implement expert-tile decomposition and CPU scheduler simulation
 - [ ] Validate the SM100a one-SM TMA/tcgen05/TMEM dense kernel on B200
 - [ ] Validate the direct CuTe TMA/tcgen05/TMEM kernel on B200
-- [ ] Implement static persistent scheduling
-- [ ] Implement dynamic work distribution
-- [ ] Add active-expert compaction
+- [x] Implement and test static persistent GPU tile assignment
+- [x] Implement and test dynamic GPU work distribution with chunked claims
+- [x] Add active-expert compaction to the generated device work list
 - [ ] Evaluate CLC-assisted work redistribution
 - [ ] Profile scheduler overhead and CTA tail effects
 - [ ] Derive and validate the crossover dispatch model
@@ -230,4 +230,11 @@ See [`docs/sm100_native_dense.md`](docs/sm100_native_dense.md).
 - [x] Native dense correctness test and small-`M` CUDA Event benchmark
 - [x] Direct CuTe TMA barriers, TMEM allocation, tcgen05 MMA, and TMEM-load epilogue
 - [x] Direct CuTe CPU-reference correctness test and benchmark harness
-- [ ] Project-owned device-side expert-tile persistent schedulers
+- [x] Two-SM TMA-multicast/tcgen05 collective reference and benchmark
+- [x] Project-owned device-side expert-tile persistent schedulers
+- [x] Exact-once GPU scheduler correctness and observed CTA-load metrics
+
+Run the complete native-kernel and scheduler matrix with
+[`tools/run_b200_optimization_suite.sh`](tools/run_b200_optimization_suite.sh).
+The output boundary and interpretation are documented in
+[`docs/b200_optimization_suite.md`](docs/b200_optimization_suite.md).

benchmarks/sm100_2sm_benchmark.cu

@@ -182,15 +182,18 @@ int main(int argc, char** argv) {
         (median_ms * 1.0e9);
     const auto stats = plan.stats();
     if (options.csv) {
-      std::cout << "kernel,device,cuda_runtime,cutlass,m,n,k,median_ms,p95_ms,"
-                   "tflops,tile_m,tile_n,tile_k,persistent_ctas,workspace_bytes\n"
+      std::cout << "kernel,device,compute_capability,cuda_runtime,cutlass,m,n,k,"
+                   "warmup,iterations,median_ms,p95_ms,tflops,tile_m,tile_n,"
+                   "tile_k,sm_count,persistent_ctas,workspace_bytes\n"
                 << sm100_2sm_kernel_name() << ',' << properties.name << ','
+                << properties.major << '.' << properties.minor << ','
                 << runtime_version << ',' << BLACKWELL_MOE_CUTLASS_REVISION
                 << ',' << options.m << ',' << options.n << ',' << options.k
-                << ',' << median_ms << ',' << p95_ms << ',' << tflops << ','
+                << ',' << options.warmup << ',' << options.iterations << ','
+                << median_ms << ',' << p95_ms << ',' << tflops << ','
                 << stats.tile_m << ',' << stats.tile_n << ',' << stats.tile_k
-                << ',' << stats.persistent_ctas << ',' << stats.workspace_bytes
-                << '\n';
+                << ',' << stats.sm_count << ',' << stats.persistent_ctas << ','
+                << stats.workspace_bytes << '\n';
     } else {
       std::cout << "kernel=" << sm100_2sm_kernel_name()
                 << " shape=" << options.m << 'x' << options.n << 'x'

docs/b200_optimization_suite.md

@@ -0,0 +1,93 @@
+# B200 Optimization Suite
+
+This suite separates three questions that must not be collapsed into one
+speedup number:
+
+1. Does each kernel produce correct output on SM100a?
+2. Which Blackwell math pipeline is best for a given dense expert shape?
+3. Does dynamic expert-tile assignment recover enough load balance to pay for
+   its queue overhead?
+
+## Implemented paths
+
+| Path | Purpose | Blackwell mechanisms |
+| --- | --- | --- |
+| CUTLASS grouped baseline | Irregular MoE reference | Portable grouped GEMM |
+| 1-SM collective | Optimized native reference | TMA, auto-staged warp specialization, tcgen05, TMEM, TMA epilogue, persistent CLC |
+| 2-SM collective | Cooperative native reference | 2-SM tcgen05, TMA multicast, auto-staged warp specialization, TMEM, persistent CLC |
+| Direct CuTe | Inspectable teaching kernel | Explicit TMA barriers, TMEM allocation, tcgen05 MMA, TMEM-to-register load |
+| Scheduler probe | Project-owned policy experiment | Persistent CTAs, static assignment, atomic dynamic queue, chunked claims |
+
+The scheduler probe intentionally uses deterministic integer work proportional
+to `valid_m * valid_n`; it does not execute GEMM. Its latency measures queue and
+load-balancing behavior in isolation. It must not be reported as a GEMM
+speedup. The native collective benchmarks measure the math pipeline, while a
+later integration combines project-owned expert-tile acquisition with that
+pipeline.
+
+## Build
+
+```bash
+cmake -S . -B build-gpu-make -G "Unix Makefiles" \
+  -DCMAKE_BUILD_TYPE=Release \
+  -DBLACKWELL_MOE_ENABLE_CUDA=ON \
+  -DBLACKWELL_MOE_ENABLE_SM100_NATIVE=ON \
+  -DBLACKWELL_MOE_CUDA_ARCHITECTURES=100 \
+  -DBLACKWELL_MOE_BUILD_TESTS=ON
+
+cmake --build build-gpu-make -j "$(nproc)"
+```
+
+The native translation units are compiled for `compute_100a/sm_100a`; the
+architecture-accelerated suffix is required for tcgen05 and TMEM.
+
+## Run everything once
+
+```bash
+./tools/run_b200_optimization_suite.sh \
+  build-gpu-make \
+  results/b200-optimization-suite
+```
+
+For a cheaper smoke run before the full matrix:
+
+```bash
+BLACKWELL_MOE_WARMUP=5 BLACKWELL_MOE_ITERATIONS=20 \
+  ./tools/run_b200_optimization_suite.sh \
+  build-gpu-make \
+  results/b200-optimization-smoke
+```
+
+The script runs all correctness tests first and stops on the first failure. It
+then writes:
+
+- `environment.txt`: GPU, driver, CUDA, and CMake metadata;
+- `correctness.log`: CPU, grouped-GEMM, scheduler, and SM100a gates;
+- `grouped_baseline.csv`: routing-distribution baseline;
+- `native_1sm_vs_2sm.csv`: 1-SM/2-SM crossover sweep;
+- `direct_cute.csv`: explicit primitive path;
+- `scheduler_probe.csv`: static/dynamic and claim-size sweep;
+- a compressed `.tar.gz` archive next to the result directory.
+
+## How to read the scheduler rows
+
+- `median_ms` and `p95_ms` include the probe kernel but exclude host workload
+  generation and the queue reset.
+- `observed_cta_work_cv` is measured from device counters after execution.
+- `tail_ratio = max(CTA work) / mean(CTA work)`; lower is better.
+- `observed_utilization = mean(CTA work) / max(CTA work)`; higher is better.
+- `claim_size=1` maximizes adaptability and atomic traffic; larger chunks
+  amortize atomics but can reintroduce tail imbalance.
+
+The desired result is a crossover, not a universal winner. Uniform workloads
+should expose dynamic-queue overhead. Heavy-hitter and Zipf workloads should
+show whether reduced tail imbalance compensates for that overhead. Sparse
+workloads additionally test active-expert compaction because zero-token experts
+produce no device work items.
+
+## Runtime validation boundary
+
+GitHub Actions compile every CUDA target for `sm_100a` in an official CUDA 13
+development container, but Actions do not provide a B200. A green CI check
+therefore proves compilation only. Correctness, latency, generated SASS, and
+performance-counter claims remain unvalidated until this suite runs on B200.

tools/run_b200_optimization_suite.sh

@@ -0,0 +1,105 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+build_dir="${1:-build-gpu-make}"
+results_dir="${2:-results/b200-optimization-suite}"
+warmup="${BLACKWELL_MOE_WARMUP:-20}"
+iterations="${BLACKWELL_MOE_ITERATIONS:-200}"
+
+mkdir -p "${results_dir}"
+
+append_csv() {
+  local output="$1"
+  shift
+  local temporary
+  temporary="$(mktemp)"
+  "$@" --csv >"${temporary}"
+  if [[ ! -s "${output}" ]]; then
+    cp "${temporary}" "${output}"
+  else
+    tail -n +2 "${temporary}" >>"${output}"
+  fi
+  rm "${temporary}"
+}
+
+for binary in \
+  test_cutlass_correctness \
+  test_scheduler_probe \
+  test_sm100_dense_correctness \
+  test_sm100_2sm_correctness \
+  test_direct_cute_correctness \
+  moe_gpu_bench \
+  moe_scheduler_gpu_bench \
+  sm100_dense_bench \
+  sm100_2sm_bench \
+  direct_cute_bench; do
+  if [[ ! -x "${build_dir}/${binary}" ]]; then
+    echo "missing executable: ${build_dir}/${binary}" >&2
+    echo "configure with BLACKWELL_MOE_ENABLE_CUDA=ON and " \
+         "BLACKWELL_MOE_ENABLE_SM100_NATIVE=ON, then rebuild" >&2
+    exit 1
+  fi
+done
+
+{
+  date -u
+  nvidia-smi --query-gpu=name,driver_version,memory.total,power.limit \
+    --format=csv,noheader
+  nvcc --version
+  cmake --version
+} >"${results_dir}/environment.txt"
+
+ctest --test-dir "${build_dir}" --output-on-failure \
+  2>&1 | tee "${results_dir}/correctness.log"
+
+baseline_csv="${results_dir}/grouped_baseline.csv"
+: >"${baseline_csv}"
+for distribution in uniform heavy_hitter sparse zipf; do
+  append_csv "${baseline_csv}" "${build_dir}/moe_gpu_bench" \
+    "--distribution=${distribution}" \
+    --experts=64 --tokens=4096 --n=7168 --k=2048 \
+    "--warmup=${warmup}" "--iterations=${iterations}"
+done
+
+native_csv="${results_dir}/native_1sm_vs_2sm.csv"
+: >"${native_csv}"
+for m in 128 256 512 1024 4096 8192; do
+  append_csv "${native_csv}" "${build_dir}/sm100_dense_bench" \
+    "--m=${m}" --n=7168 --k=2048 \
+    "--warmup=${warmup}" "--iterations=${iterations}"
+  append_csv "${native_csv}" "${build_dir}/sm100_2sm_bench" \
+    "--m=${m}" --n=7168 --k=2048 \
+    "--warmup=${warmup}" "--iterations=${iterations}"
+done
+
+direct_csv="${results_dir}/direct_cute.csv"
+: >"${direct_csv}"
+for m in 128 256 512 1024; do
+  append_csv "${direct_csv}" "${build_dir}/direct_cute_bench" \
+    "--m=${m}" --n=7168 --k=2048 \
+    "--warmup=${warmup}" "--iterations=${iterations}"
+done
+
+scheduler_csv="${results_dir}/scheduler_probe.csv"
+: >"${scheduler_csv}"
+for distribution in uniform heavy_hitter sparse zipf; do
+  append_csv "${scheduler_csv}" "${build_dir}/moe_scheduler_gpu_bench" \
+    "--distribution=${distribution}" --scheduler=static \
+    --experts=64 --tokens=4096 --n=7168 --tile-m=128 --tile-n=128 \
+    --ctas=120 --claim-size=1 --work-scale=8 \
+    "--warmup=${warmup}" "--iterations=${iterations}"
+  for claim_size in 1 2 4 8; do
+    append_csv "${scheduler_csv}" "${build_dir}/moe_scheduler_gpu_bench" \
+      "--distribution=${distribution}" --scheduler=dynamic \
+      --experts=64 --tokens=4096 --n=7168 --tile-m=128 --tile-n=128 \
+      --ctas=120 "--claim-size=${claim_size}" --work-scale=8 \
+      "--warmup=${warmup}" "--iterations=${iterations}"
+  done
+done
+
+tar -czf "${results_dir}.tar.gz" -C "$(dirname "${results_dir}")" \
+  "$(basename "${results_dir}")"
+
+echo "B200 optimization suite complete"
+echo "results: ${results_dir}"
+echo "archive: ${results_dir}.tar.gz"