From cd4317a07cf29809de6071366678ebd60b5dd2ca Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Wed, 16 Sep 2026 17:04:27 -0700 Subject: [PATCH 1/2] Integrate fast.cu kernels with native harness pairs and B200 experiments --- code/ch09/README.md | 7 + code/ch10/README.md | 7 + code/ch11/README.md | 7 + code/ch14/README.md | 7 + code/labs/README.md | 1 + code/labs/fast_cu/README.md | 86 + code/labs/fast_cu/__init__.py | 1 + code/labs/fast_cu/b200_epilogue.cu | 199 ++ code/labs/fast_cu/b200_epilogue.py | 236 ++ code/labs/fast_cu/baseline_b200_epilogue.py | 13 + code/labs/fast_cu/baseline_h100_bf16_gemm.py | 43 + code/labs/fast_cu/baseline_int32_reduction.py | 47 + code/labs/fast_cu/baseline_nvfp4_gemm.py | 20 + code/labs/fast_cu/build.py | 89 + code/labs/fast_cu/h100_common.py | 352 +++ code/labs/fast_cu/h100_gemm_extension.cu | 371 +++ code/labs/fast_cu/h100_reduction_kernels.cu | 174 ++ code/labs/fast_cu/nvfp4.py | 334 ++ code/labs/fast_cu/nvfp4_native.cu | 569 ++++ .../fast_cu/nvfp4_performance_intake.yaml | 99 + code/labs/fast_cu/nvfp4_workload_spec.yaml | 125 + code/labs/fast_cu/optimization_guide.md | 81 + code/labs/fast_cu/optimized_b200_epilogue.py | 13 + code/labs/fast_cu/optimized_h100_bf16_gemm.py | 44 + .../labs/fast_cu/optimized_int32_reduction.py | 28 + code/labs/fast_cu/optimized_nvfp4_gemm.py | 24 + code/labs/fast_cu/performance_intake.yaml | 49 + code/labs/fast_cu/source.py | 54 + code/labs/fast_cu/upstream/.gitignore | 1 + code/labs/fast_cu/upstream/LICENSE | 21 + code/labs/fast_cu/upstream/Makefile | 56 + code/labs/fast_cu/upstream/README.md | 58 + .../fast_cu/upstream/gb300/nvfp4/README.md | 145 + .../fast_cu/upstream/gb300/nvfp4/gemm0.cuh | 1148 +++++++ .../fast_cu/upstream/gb300/nvfp4/gemm1.cuh | 1152 +++++++ .../fast_cu/upstream/gb300/nvfp4/gemm2.cuh | 1254 ++++++++ .../fast_cu/upstream/gb300/nvfp4/gemm3.cuh | 1256 ++++++++ .../fast_cu/upstream/gb300/nvfp4/gemm4.cuh | 1258 ++++++++ .../fast_cu/upstream/gb300/nvfp4/gemm5.cuh | 1263 ++++++++ .../fast_cu/upstream/gb300/nvfp4/gemm6.cuh | 1264 ++++++++ .../fast_cu/upstream/gb300/nvfp4/gemm7.cuh | 1421 +++++++++ .../fast_cu/upstream/gb300/nvfp4/gemm8.cuh | 1688 ++++++++++ .../fast_cu/upstream/gb300/nvfp4/gemm9.cuh | 2770 +++++++++++++++++ .../labs/fast_cu/upstream/gb300/nvfp4/main.cu | 1895 +++++++++++ .../upstream/gb300/nvfp4/run_bench_modes.sh | 83 + code/labs/fast_cu/upstream/h100/logs.txt | 8 + code/labs/fast_cu/upstream/h100/matmul.cu | 244 ++ .../upstream/h100/matmul/cublaslt_matmul.cu | 74 + .../fast_cu/upstream/h100/matmul/matmul_1.cuh | 180 ++ .../upstream/h100/matmul/matmul_10.cuh | 621 ++++ .../upstream/h100/matmul/matmul_11.cuh | 639 ++++ .../upstream/h100/matmul/matmul_12.cuh | 631 ++++ .../fast_cu/upstream/h100/matmul/matmul_2.cuh | 197 ++ .../fast_cu/upstream/h100/matmul/matmul_3.cuh | 410 +++ .../fast_cu/upstream/h100/matmul/matmul_4.cuh | 393 +++ .../fast_cu/upstream/h100/matmul/matmul_5.cuh | 407 +++ .../fast_cu/upstream/h100/matmul/matmul_6.cuh | 450 +++ .../fast_cu/upstream/h100/matmul/matmul_7.cuh | 482 +++ .../fast_cu/upstream/h100/matmul/matmul_8.cuh | 560 ++++ .../fast_cu/upstream/h100/matmul/matmul_9.cuh | 600 ++++ .../h100/matmul/pingpong_experimental.cuh | 761 +++++ code/labs/fast_cu/upstream/h100/sum.cu | 332 ++ code/labs/fast_cu/upstream_manifest.json | 123 + code/labs/fast_cu/validation.json | 79 + code/labs/fast_cu/validation.md | 94 + code/labs/fast_cu/workload_spec.yaml | 47 + code/tests/test_fast_cu_b200_epilogue.py | 266 ++ code/tests/test_fast_cu_build.py | 59 + code/tests/test_fast_cu_h100.py | 211 ++ code/tests/test_fast_cu_nvfp4.py | 235 ++ code/tests/test_fast_cu_source.py | 73 + 71 files changed, 27989 insertions(+) create mode 100644 code/labs/fast_cu/README.md create mode 100644 code/labs/fast_cu/__init__.py create mode 100644 code/labs/fast_cu/b200_epilogue.cu create mode 100644 code/labs/fast_cu/b200_epilogue.py create mode 100644 code/labs/fast_cu/baseline_b200_epilogue.py create mode 100644 code/labs/fast_cu/baseline_h100_bf16_gemm.py create mode 100644 code/labs/fast_cu/baseline_int32_reduction.py create mode 100644 code/labs/fast_cu/baseline_nvfp4_gemm.py create mode 100644 code/labs/fast_cu/build.py create mode 100644 code/labs/fast_cu/h100_common.py create mode 100644 code/labs/fast_cu/h100_gemm_extension.cu create mode 100644 code/labs/fast_cu/h100_reduction_kernels.cu create mode 100644 code/labs/fast_cu/nvfp4.py create mode 100644 code/labs/fast_cu/nvfp4_native.cu create mode 100644 code/labs/fast_cu/nvfp4_performance_intake.yaml create mode 100644 code/labs/fast_cu/nvfp4_workload_spec.yaml create mode 100644 code/labs/fast_cu/optimization_guide.md create mode 100644 code/labs/fast_cu/optimized_b200_epilogue.py create mode 100644 code/labs/fast_cu/optimized_h100_bf16_gemm.py create mode 100644 code/labs/fast_cu/optimized_int32_reduction.py create mode 100644 code/labs/fast_cu/optimized_nvfp4_gemm.py create mode 100644 code/labs/fast_cu/performance_intake.yaml create mode 100644 code/labs/fast_cu/source.py create mode 100644 code/labs/fast_cu/upstream/.gitignore create mode 100644 code/labs/fast_cu/upstream/LICENSE create mode 100644 code/labs/fast_cu/upstream/Makefile create mode 100644 code/labs/fast_cu/upstream/README.md create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/README.md create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/gemm0.cuh create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/gemm1.cuh create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/gemm2.cuh create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/gemm3.cuh create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/gemm4.cuh create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/gemm5.cuh create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/gemm6.cuh create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/gemm7.cuh create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/gemm8.cuh create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/gemm9.cuh create mode 100644 code/labs/fast_cu/upstream/gb300/nvfp4/main.cu create mode 100755 code/labs/fast_cu/upstream/gb300/nvfp4/run_bench_modes.sh create mode 100644 code/labs/fast_cu/upstream/h100/logs.txt create mode 100644 code/labs/fast_cu/upstream/h100/matmul.cu create mode 100644 code/labs/fast_cu/upstream/h100/matmul/cublaslt_matmul.cu create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_1.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_10.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_11.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_12.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_2.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_3.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_4.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_5.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_6.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_7.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_8.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/matmul_9.cuh create mode 100644 code/labs/fast_cu/upstream/h100/matmul/pingpong_experimental.cuh create mode 100644 code/labs/fast_cu/upstream/h100/sum.cu create mode 100644 code/labs/fast_cu/upstream_manifest.json create mode 100644 code/labs/fast_cu/validation.json create mode 100644 code/labs/fast_cu/validation.md create mode 100644 code/labs/fast_cu/workload_spec.yaml create mode 100644 code/tests/test_fast_cu_b200_epilogue.py create mode 100644 code/tests/test_fast_cu_build.py create mode 100644 code/tests/test_fast_cu_h100.py create mode 100644 code/tests/test_fast_cu_nvfp4.py create mode 100644 code/tests/test_fast_cu_source.py diff --git a/code/ch09/README.md b/code/ch09/README.md index 8355a41cc..5e733881f 100644 --- a/code/ch09/README.md +++ b/code/ch09/README.md @@ -87,3 +87,10 @@ python -m cli.aisp bench run --targets ch09 --profile minimal - `inline_ptx_example.cu` demonstrates how to wrap tcgen05 intrinsics safely with architecture guards. - `requirements.txt` includes the shared `requirements_latest.txt` pins for PyTorch 2.9.1 and its compatible Triton release. - `optimized_cublaslt_gemm_fp4` is intentionally capability-gated: if cuBLASLt cannot provide the native block-scaled NVFP4 heuristic, the benchmark reports a clean skip instead of silently falling back to a different FP4 mode. + +## Related fast.cu experiments + +The [fast.cu lab](../labs/fast_cu/README.md) preserves the complete H100 and GB300 +source progressions and adds native harness pairs, including B200 adaptations. +See its [chapter mapping and optimization guide](../labs/fast_cu/optimization_guide.md) +for the overlap with this chapter and the SM100 versus SM103 hardware boundaries. diff --git a/code/ch10/README.md b/code/ch10/README.md index 6236ec93b..8cb7789b7 100644 --- a/code/ch10/README.md +++ b/code/ch10/README.md @@ -88,3 +88,10 @@ python -m cli.aisp bench run --targets ch10 --profile minimal - `cufile_gds_example.py` demonstrates integrating GPUDirect Storage into tensor-core pipelines for IO-heavy training loops. - `requirements_cufile.txt` holds the optional `cufile` wheel; install it only on hosts with GPUDirect Storage enabled. - The CUTLASS-style warp-specialization pair provides a reference implementation aligned with `sm100_mma_array_warpspecialized` for performance comparison. + +## Related fast.cu experiments + +The [fast.cu lab](../labs/fast_cu/README.md) preserves the complete H100 and GB300 +source progressions and adds native harness pairs, including B200 adaptations. +See its [chapter mapping and optimization guide](../labs/fast_cu/optimization_guide.md) +for the overlap with this chapter and the SM100 versus SM103 hardware boundaries. diff --git a/code/ch11/README.md b/code/ch11/README.md index 16b565127..1a06b92dc 100644 --- a/code/ch11/README.md +++ b/code/ch11/README.md @@ -86,3 +86,10 @@ python -m cli.aisp bench run --targets ch11 --profile minimal - The README calls out the legacy target names explicitly so book-facing labels still point at the actual copy+elementwise overlap workload and runtime-adaptive scheduling pairs. - `warp_specialized_triton.py` provides a Triton analogue for the CUDA concurrency demos so you can compare compiler-generated schedules. - `kv_prefetch_pipeline_enhanced_demo.cu` builds on the DSMEM kernels bundled in this directory so you can study the entire pipeline locally. + +## Related fast.cu experiments + +The [fast.cu lab](../labs/fast_cu/README.md) preserves the complete H100 and GB300 +source progressions and adds native harness pairs, including B200 adaptations. +See its [chapter mapping and optimization guide](../labs/fast_cu/optimization_guide.md) +for the overlap with this chapter and the SM100 versus SM103 hardware boundaries. diff --git a/code/ch14/README.md b/code/ch14/README.md index c7bb630b0..6813beb50 100644 --- a/code/ch14/README.md +++ b/code/ch14/README.md @@ -86,3 +86,10 @@ python -m cli.aisp bench run --targets ch14 --profile minimal - `requirements.txt` includes the shared `requirements_latest.txt` pins for stable PyTorch and Triton wheels aligned with the CUDA 13 toolchain. - For repo-native supporting examples that fill the training hot-path gaps without changing this chapter's primary compile narrative, see `labs/training_hotpath`. - `cublas_vs_cutlass` is a supplementary comparison pair. Chapter-native performance claims stay anchored on `model_compile_reduced_precision`, `regional_triton`, and `triton_persistent`. + +## Related fast.cu experiments + +The [fast.cu lab](../labs/fast_cu/README.md) preserves the complete H100 and GB300 +source progressions and adds native harness pairs, including B200 adaptations. +See its [chapter mapping and optimization guide](../labs/fast_cu/optimization_guide.md) +for the overlap with this chapter and the SM100 versus SM103 hardware boundaries. diff --git a/code/labs/README.md b/code/labs/README.md index 0b4d0b247..2cd50d17a 100644 --- a/code/labs/README.md +++ b/code/labs/README.md @@ -66,6 +66,7 @@ python -m cli.aisp bench list-targets --chapter labs/moe_cuda | `labs/cutlass_profiler_kernel_selector/` | CUTLASS profiler-based kernel selection | ch06, ch09 | | `labs/decode_optimization/` | Decoder hot-path optimization | ch18, ch19 | | `labs/dynamic_router/` | Dynamic prefill/decode routing | ch17, ch19 | +| `labs/fast_cu/` | Pinned fast.cu H100 GEMM, B200 reduction and wide-store epilogue adaptations, and the complete GB300 NVFP4 ladder; see the [chapter mapping](fast_cu/optimization_guide.md) | ch09, ch10, ch11, ch14 | | `labs/flashattention4/` | FlashAttention-4 pipeline co-design, including a [Colfax diary guide](flashattention4/colfax_optimization_diaries.md) to decode S/P ping-pong and backward hdim-64 TMEM de-aliasing | ch10, ch11, ch18 | | `labs/flashattention_gluon/` | FlashAttention experimentation | ch18 | | `labs/flashinfer_attention/` | FlashInfer block-sparse attention lab | ch16 | diff --git a/code/labs/fast_cu/README.md b/code/labs/fast_cu/README.md new file mode 100644 index 000000000..69c01c1f1 --- /dev/null +++ b/code/labs/fast_cu/README.md @@ -0,0 +1,86 @@ +# Lab - fast.cu kernel optimization studies + +This lab brings [Pranjal Shankhdhar's fast.cu](https://github.com/pranjalssh/fast.cu) +into the repository: all twelve Hopper GEMM snapshots, four reduction variants, +and all ten GB300 NVFP4 snapshots live in [upstream/](upstream/). Repository-native +`baseline_*` / `optimized_*` pairs add explicit hardware gates, current-stream +execution, caller-owned random inputs, and full output verification. + +The [optimization guide](optimization_guide.md) maps the examples and lessons to +existing chapters and labs. The [source manifest](upstream_manifest.json) pins +revision `2dfe5e26aecfd9e5f27bf9d5837deea01acda24b` and the SHA256 of every upstream +file. The upstream MIT license is preserved in [upstream/LICENSE](upstream/LICENSE). + +## Native benchmark pairs + +| Target in `labs/fast_cu` | Baseline | Optimized | Required GPU/toolkit | +| --- | --- | --- | --- | +| `h100_bf16_gemm` | cuBLAS BF16 GEMM | fast.cu Hopper kernel 12 | H100, SM90, CUDA 12+ | +| `int32_reduction` | CUB sum | Corrected vectorized warp reduction | H100 SM90/CUDA 12+ or B200 SM100/CUDA 12.8+ | +| `b200_epilogue` | Two 128-bit stores per thread | One 256-bit store per thread | B200, SM100, CUDA 12.9+ | +| `nvfp4_gemm` | cuBLASLt NVFP4 GEMM | fast.cu r9 NVFP4 GEMM | B300/GB300, SM103, CUDA 13.1+ | + +The NVFP4 kernel uses architecture-specific SM103 instructions and is **not a +B200 kernel**. Changing the compiler target to SM100 does not port its K=96 MMA +pipeline. The B200 epilogue is a separate experiment applying the article's wide +store technique to the same FP32-to-FP16 conversion workload in both arms. It +does not stand in for a complete NVFP4 GEMM. + +## Run with the repository harness + +Run from `code/` in the normal CUDA-enabled repository environment (PyTorch, +matching local CUDA toolkit, Ninja, Nsight Compute and Nsight Systems): + +```bash +python -m cli.aisp bench list-targets --chapter labs/fast_cu + +# B200: architecture-compatible adaptations +python -m cli.aisp bench run --targets labs/fast_cu:int32_reduction --profile deep_dive --single-gpu +python -m cli.aisp bench run --targets labs/fast_cu:b200_epilogue --profile deep_dive --single-gpu + +# H100 +python -m cli.aisp bench run --targets labs/fast_cu:h100_bf16_gemm --profile deep_dive --single-gpu + +# B300/GB300 with CUDA 13.1+ +python -m cli.aisp bench run --targets labs/fast_cu:nvfp4_gemm --profile deep_dive --single-gpu +``` + +Select an allocated GPU using `CUDA_VISIBLE_DEVICES` before launching. Unsupported +hardware/toolkits produce explicit `SKIPPED:` diagnostics; no precision or +architecture fallback is substituted. Extension compilation and preparation run +in setup, outside timed execution. The source tree is checked before compilation, +and build cache names include all upstream file digests and adapter sources. + +Use the repository's strict validity profile and clock-lock/profiler machinery for +performance evidence. The original upstream Makefile and shell scripts are retained +as source evidence; their direct `sudo nvidia-smi`/`ncu` helpers are not the lab's +execution interface. Do not build inside the pinned `upstream/` directory: generated +files intentionally fail the exact file-set integrity check. + +## Correctness and measurement boundaries + +- Both arms use the same logical inputs, precision, shape, and full timed output. + Preparation, JIT, workspace allocation, and reference checks are separate from + the steady-state hot path. +- The reduction adaptation moves output initialization ahead of the grid on the + same stream, preventing the original cross-block reset race. Its bounded input + distribution avoids signed 32-bit overflow. +- The Hopper adapter normalizes the kernel's physical transposed output into a + logical matrix view for verification. A successful upstream process exit is not + treated as evidence of correctness. +- The NVFP4 native pair uses a single set of resident buffers. The article's + headline instead rotates input buffers and alternates five cooled-down timing + rounds. Those are different workload/cache protocols; this lab does not import + the article's speedup as an expectation. +- [Validation](validation.md) records what has actually run and what remains + hardware-gated. Source preservation and CPU tests do not prove GPU correctness. + +Focused checks, from `code/`: + +```bash +python -m pytest tests/test_fast_cu_*.py -q +python -m core.scripts.linting.check_benchmarks labs/fast_cu --fail-on-warnings +``` + +See [performance_intake.yaml](performance_intake.yaml) and +[workload_spec.yaml](workload_spec.yaml) for cost hypotheses and fixed workloads. diff --git a/code/labs/fast_cu/__init__.py b/code/labs/fast_cu/__init__.py new file mode 100644 index 000000000..693ee1bc1 --- /dev/null +++ b/code/labs/fast_cu/__init__.py @@ -0,0 +1 @@ +"""Pinned fast.cu CUDA kernels and repository-native benchmark adapters.""" diff --git a/code/labs/fast_cu/b200_epilogue.cu b/code/labs/fast_cu/b200_epilogue.cu new file mode 100644 index 000000000..6b429cacf --- /dev/null +++ b/code/labs/fast_cu/b200_epilogue.cu @@ -0,0 +1,199 @@ +#include +#include +#include +#include + +#include +#include +#include + +#include +#include + +#if __CUDACC_VER_MAJOR__ < 12 || \ + (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 9) +#error "fast.cu B200 epilogue requires CUDA 12.9 or newer for st.global.v8.b32" +#endif + +namespace { + +constexpr int kThreads = 256; +constexpr int64_t kElementsPerThread = 16; +constexpr uintptr_t kWideStoreAlignment = 32; + +__device__ __forceinline__ uint32_t pack_f16x2(float lo, float hi) { + uint32_t packed; + // PTX places its first f32 source in the high 16 bits. Reverse the operands + // so the lower-addressed FP16 value occupies the low half of the b32 word. + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" + : "=r"(packed) + : "f"(hi), "f"(lo)); + return packed; +} + +__device__ __forceinline__ uint64_t make_evict_first_policy() { + uint64_t policy; + asm volatile( + "createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" + : "=l"(policy)); + return policy; +} + +__device__ __forceinline__ void store_b128( + void* pointer, + uint4 value, + uint64_t policy) { + asm volatile( + "st.global.L1::no_allocate.L2::cache_hint.v4.b32 " + "[%0], {%1, %2, %3, %4}, %5;" + : + : "l"(pointer), "r"(value.x), "r"(value.y), "r"(value.z), + "r"(value.w), "l"(policy) + : "memory"); +} + +__device__ __forceinline__ void store_b256( + void* pointer, + uint4 lo, + uint4 hi, + uint64_t policy) { + asm volatile( + "st.global.L1::no_allocate.L2::cache_hint.v8.b32 " + "[%0], {%1, %2, %3, %4, %5, %6, %7, %8}, %9;" + : + : "l"(pointer), "r"(lo.x), "r"(lo.y), "r"(lo.z), "r"(lo.w), + "r"(hi.x), "r"(hi.y), "r"(hi.z), "r"(hi.w), "l"(policy) + : "memory"); +} + +template +__global__ void epilogue_kernel( + const float* accumulators, + __half* output, + int64_t work_items) { + const int64_t work_item = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (work_item >= work_items) { + return; + } + + const int64_t base = work_item * kElementsPerThread; + uint4 lo; + uint4 hi; + lo.x = pack_f16x2(accumulators[base + 0], accumulators[base + 1]); + lo.y = pack_f16x2(accumulators[base + 2], accumulators[base + 3]); + lo.z = pack_f16x2(accumulators[base + 4], accumulators[base + 5]); + lo.w = pack_f16x2(accumulators[base + 6], accumulators[base + 7]); + hi.x = pack_f16x2(accumulators[base + 8], accumulators[base + 9]); + hi.y = pack_f16x2(accumulators[base + 10], accumulators[base + 11]); + hi.z = pack_f16x2(accumulators[base + 12], accumulators[base + 13]); + hi.w = pack_f16x2(accumulators[base + 14], accumulators[base + 15]); + + const uint64_t policy = make_evict_first_policy(); + void* const destination = static_cast(output + base); + if constexpr (WideStore) { + store_b256(destination, lo, hi, policy); + } else { + store_b128(destination, lo, policy); + store_b128(static_cast(output + base + 8), hi, policy); + } +} + +void check_tensor_contract( + const torch::Tensor& accumulators, + const torch::Tensor& output) { + TORCH_CHECK( + accumulators.is_cuda() && output.is_cuda(), + "fast.cu B200 epilogue tensors must be CUDA tensors"); + TORCH_CHECK( + accumulators.device() == output.device(), + "fast.cu B200 epilogue tensors must be on the same CUDA device"); + TORCH_CHECK( + accumulators.scalar_type() == at::kFloat, + "fast.cu B200 epilogue accumulators must be float32"); + TORCH_CHECK( + output.scalar_type() == at::kHalf, + "fast.cu B200 epilogue output must be float16"); + TORCH_CHECK( + accumulators.is_contiguous() && output.is_contiguous(), + "fast.cu B200 epilogue tensors must be contiguous"); + TORCH_CHECK( + accumulators.sizes() == output.sizes(), + "fast.cu B200 epilogue input and output shapes must match"); + + const int64_t count = accumulators.numel(); + TORCH_CHECK(count > 0, "fast.cu B200 epilogue count must be positive"); + TORCH_CHECK( + count % kElementsPerThread == 0, + "fast.cu B200 epilogue count must be divisible by 16; found ", + count); + TORCH_CHECK( + reinterpret_cast(output.data_ptr()) % kWideStoreAlignment == 0, + "fast.cu B200 epilogue output must be 32-byte aligned"); + + const uintptr_t input_begin = + reinterpret_cast(accumulators.data_ptr()); + const uintptr_t input_end = + input_begin + static_cast(count) * sizeof(float); + const uintptr_t output_begin = reinterpret_cast(output.data_ptr()); + const uintptr_t output_end = + output_begin + static_cast(count) * sizeof(__half); + TORCH_CHECK( + input_end <= output_begin || output_end <= input_begin, + "fast.cu B200 epilogue input and output storage must not overlap"); + + const int64_t work_items = count / kElementsPerThread; + constexpr int64_t kMaxWorkItems = + static_cast(std::numeric_limits::max()) * kThreads; + TORCH_CHECK( + work_items <= kMaxWorkItems, + "fast.cu B200 epilogue tensor is too large for the launch grid"); +} + +template +torch::Tensor launch_epilogue( + const torch::Tensor& accumulators, + torch::Tensor output) { + check_tensor_contract(accumulators, output); + const c10::cuda::CUDAGuard guard(accumulators.device()); + const int64_t work_items = accumulators.numel() / kElementsPerThread; + const int blocks = static_cast((work_items + kThreads - 1) / kThreads); + const cudaStream_t stream = + at::cuda::getCurrentCUDAStream(accumulators.get_device()); + epilogue_kernel<<>>( + accumulators.data_ptr(), + reinterpret_cast<__half*>(output.data_ptr()), + work_items); + AT_CUDA_CHECK(cudaGetLastError()); + return output; +} + +void require_exact_sm100() { + int device = -1; + AT_CUDA_CHECK(cudaGetDevice(&device)); + cudaDeviceProp properties{}; + AT_CUDA_CHECK(cudaGetDeviceProperties(&properties, device)); + TORCH_CHECK( + properties.major == 10 && properties.minor == 0, + "fast.cu B200 epilogue requires exact SM100 (compute capability 10.0); found ", + properties.major, + ".", + properties.minor); +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def( + "baseline", + &launch_epilogue, + "FP32-to-FP16 epilogue using two 128-bit stores"); + module.def( + "optimized", + &launch_epilogue, + "FP32-to-FP16 epilogue using one 256-bit store"); + module.def( + "require_exact_sm100", + &require_exact_sm100, + "Reject every CUDA architecture except exact SM100"); +} diff --git a/code/labs/fast_cu/b200_epilogue.py b/code/labs/fast_cu/b200_epilogue.py new file mode 100644 index 000000000..109be671d --- /dev/null +++ b/code/labs/fast_cu/b200_epilogue.py @@ -0,0 +1,236 @@ +"""SM100 epilogue store-width isolation derived from fast.cu r4/r5. + +This benchmark is an FP32-accumulator-to-FP16 epilogue microbenchmark. It is +not an NVFP4 GEMM and cannot reproduce fast.cu's end-to-end headline result. +""" + +from __future__ import annotations + +import re +from collections.abc import Callable +from functools import lru_cache +from pathlib import Path +from typing import Any + +import torch + +from core.benchmark.verification_mixin import VerificationPayloadMixin +from core.harness.benchmark_harness import BaseBenchmark, BenchmarkConfig, WorkloadMetadata +from labs.fast_cu.build import load_cuda_extension + +CUDA_SOURCE = Path(__file__).with_suffix(".cu") +DEFAULT_NUM_ELEMENTS = 1 << 24 +ELEMENTS_PER_THREAD = 16 +OUTPUT_ALIGNMENT_BYTES = 32 +MINIMUM_CUDA = (12, 9) +EXPECTED_COMPUTE_CAPABILITY = (10, 0) + + +def parse_cuda_build(version: str | None) -> tuple[int, int]: + if not isinstance(version, str): + raise RuntimeError("SKIPPED: fast.cu B200 epilogue requires a CUDA-enabled PyTorch build") + match = re.match(r"^(\d+)\.(\d+)(?:\.|$)", version) + if match is None: + raise RuntimeError(f"Cannot parse PyTorch CUDA build version: {version!r}") + return int(match[1]), int(match[2]) + + +def validate_num_elements(num_elements: int) -> None: + if type(num_elements) is not int or num_elements <= 0: + raise ValueError("num_elements must be a positive integer") + if num_elements % ELEMENTS_PER_THREAD: + raise ValueError("num_elements must be divisible by 16") + + +def ensure_b200_capability_supported(capability: tuple[int, int]) -> None: + if capability != EXPECTED_COMPUTE_CAPABILITY: + raise RuntimeError( + "SKIPPED: fast.cu B200 epilogue requires exact SM100 " + f"(compute capability 10.0); found {capability[0]}.{capability[1]}" + ) + + +def ensure_cuda_build_supported(cuda_build: tuple[int, int]) -> None: + if cuda_build < MINIMUM_CUDA: + raise RuntimeError( + "SKIPPED: st.global.v8.b32 requires a CUDA 12.9+ PyTorch build; " + f"found {cuda_build[0]}.{cuda_build[1]}" + ) + + +def require_b200_runtime(device: torch.device) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("SKIPPED: fast.cu B200 epilogue requires CUDA") + device = torch.device(device) + if device.type != "cuda": + raise RuntimeError("SKIPPED: fast.cu B200 epilogue requires a CUDA device") + ensure_b200_capability_supported(torch.cuda.get_device_capability(device)) + ensure_cuda_build_supported(parse_cuda_build(torch.version.cuda)) + + +def validate_epilogue_tensors( + accumulators: torch.Tensor, + output: torch.Tensor, +) -> None: + if not isinstance(accumulators, torch.Tensor) or not isinstance(output, torch.Tensor): + raise TypeError("accumulators and output must be torch.Tensor instances") + if accumulators.dtype != torch.float32: + raise TypeError("accumulators must have dtype torch.float32") + if output.dtype != torch.float16: + raise TypeError("output must have dtype torch.float16") + if not accumulators.is_contiguous() or not output.is_contiguous(): + raise ValueError("accumulators and output must be contiguous") + if accumulators.shape != output.shape: + raise ValueError("accumulators and output must have matching shapes") + validate_num_elements(accumulators.numel()) + if output.data_ptr() % OUTPUT_ALIGNMENT_BYTES: + raise ValueError("output must be 32-byte aligned") + if accumulators.device != output.device: + raise ValueError("accumulators and output must be on the same device") + if not accumulators.is_cuda: + raise ValueError("accumulators and output must be CUDA tensors") + + +@lru_cache(maxsize=1) +def load_b200_epilogue_extension() -> Any: + return load_cuda_extension( + "fast_cu_b200_epilogue", + [CUDA_SOURCE], + extra_cuda_cflags=[ + "-O3", + "-std=c++17", + "-lineinfo", + "-gencode=arch=compute_100,code=sm_100", + ], + extra_ldflags=[], + minimum_cuda=MINIMUM_CUDA, + ) + + +class B200EpilogueBenchmark(VerificationPayloadMixin, BaseBenchmark): + """Compare equal epilogues whose only device-code difference is store width.""" + + def __init__(self, *, optimized: bool, num_elements: int = DEFAULT_NUM_ELEMENTS): + if type(optimized) is not bool: + raise TypeError("optimized must be bool") + validate_num_elements(num_elements) + super().__init__() + self.optimized = optimized + self.num_elements = num_elements + self.accumulators: torch.Tensor | None = None + self._output_buffer: torch.Tensor | None = None + self.output: torch.Tensor | None = None + self._verification_output: torch.Tensor | None = None + self._extension: Any = None + self._launch: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None + self._caller_seed: int | None = None + self._ran = False + self._workload = WorkloadMetadata( + requests_per_iteration=1.0, + samples_per_iteration=float(num_elements), + bytes_per_iteration=float(num_elements * (4 + 2)), + ) + + def setup(self) -> None: + self._ran = False + self._verification_output = None + self._verification_payload = None + self._output_buffer = None + self.output = None + require_b200_runtime(self.device) + self._extension = load_b200_epilogue_extension() + self._extension.require_exact_sm100() + self._caller_seed = int(torch.initial_seed()) + self.accumulators = torch.randn( + self.num_elements, + device=self.device, + dtype=torch.float32, + ) + self._output_buffer = torch.empty( + self.num_elements, + device=self.device, + dtype=torch.float16, + ) + validate_epilogue_tensors(self.accumulators, self._output_buffer) + self._launch = self._extension.optimized if self.optimized else self._extension.baseline + + def benchmark_fn(self) -> None: + if self.accumulators is None or self._output_buffer is None or self._launch is None: + raise RuntimeError("setup() must complete before benchmark_fn()") + label = ( + "fast_cu_b200_epilogue_optimized" + if self.optimized + else "fast_cu_b200_epilogue_baseline" + ) + with self._nvtx_range(label): + self._launch(self.accumulators, self._output_buffer) + self.output = self._output_buffer + self._ran = True + + def capture_verification_payload(self) -> None: + if not self._ran or self.accumulators is None or self.output is None: + raise RuntimeError("benchmark_fn() must run before verification capture") + self._verification_output = self.output.clone() + self._set_verification_payload( + inputs={"accumulators": self.accumulators}, + output=self._verification_output, + batch_size=1, + parameter_count=0, + precision_flags={"fp16": True, "bf16": False, "fp8": False, "tf32": False}, + output_tolerance=(0.0, 0.0), + ) + + def validate_result(self) -> str | None: + if not self._ran or self.accumulators is None or self.output is None: + return "No timed epilogue output" + expected = self.accumulators.to(torch.float16) + if not torch.equal(self.output, expected): + return "Epilogue output differs from full FP32-to-FP16 reference" + return None + + def get_config(self) -> BenchmarkConfig: + return BenchmarkConfig( + iterations=100, + warmup=20, + single_gpu=True, + measurement_timeout_seconds=180, + ) + + def get_workload_metadata(self) -> WorkloadMetadata: + return self._workload + + def get_custom_metrics(self) -> dict[str, float]: + return { + "fast_cu.elements_per_thread": float(ELEMENTS_PER_THREAD), + "fast_cu.store_width_bytes": 32.0 if self.optimized else 16.0, + "fast_cu.store_instructions_per_thread": 1.0 if self.optimized else 2.0, + "fast_cu.full_nvfp4_gemm": 0.0, + } + + def teardown(self) -> None: + self.accumulators = None + self._output_buffer = None + self.output = None + self._verification_output = None + self._extension = None + self._launch = None + self._caller_seed = None + self._ran = False + self._verification_payload = None + super().teardown() + + +__all__ = [ + "B200EpilogueBenchmark", + "DEFAULT_NUM_ELEMENTS", + "ELEMENTS_PER_THREAD", + "EXPECTED_COMPUTE_CAPABILITY", + "MINIMUM_CUDA", + "ensure_b200_capability_supported", + "ensure_cuda_build_supported", + "load_b200_epilogue_extension", + "parse_cuda_build", + "require_b200_runtime", + "validate_epilogue_tensors", + "validate_num_elements", +] diff --git a/code/labs/fast_cu/baseline_b200_epilogue.py b/code/labs/fast_cu/baseline_b200_epilogue.py new file mode 100644 index 000000000..b735e822a --- /dev/null +++ b/code/labs/fast_cu/baseline_b200_epilogue.py @@ -0,0 +1,13 @@ +"""SM100 control: FP32-to-FP16 epilogue with two 128-bit stores.""" + +from core.harness.benchmark_harness import BaseBenchmark +from labs.fast_cu.b200_epilogue import B200EpilogueBenchmark + + +class BaselineB200EpilogueBenchmark(B200EpilogueBenchmark): + def __init__(self): + super().__init__(optimized=False) + + +def get_benchmark() -> BaseBenchmark: + return BaselineB200EpilogueBenchmark() diff --git a/code/labs/fast_cu/baseline_h100_bf16_gemm.py b/code/labs/fast_cu/baseline_h100_bf16_gemm.py new file mode 100644 index 000000000..5f5ea273c --- /dev/null +++ b/code/labs/fast_cu/baseline_h100_bf16_gemm.py @@ -0,0 +1,43 @@ +"""Baseline: cuBLAS BF16 GEMM for the fast.cu H100 comparison.""" + +from __future__ import annotations + +import torch + +from core.harness.benchmark_harness import BaseBenchmark +from labs.fast_cu.h100_common import H100Bf16GemmBenchmarkBase + + +class BaselineH100Bf16GemmBenchmark(H100Bf16GemmBenchmarkBase): + """Compute the full A @ B.T output with cuBLAS on the active stream.""" + + nvtx_label = "fast_cu_baseline_h100_bf16_gemm" + + def setup(self) -> None: + self._setup_tensors() + if ( + self.extension is None + or self.matrix_a is None + or self.matrix_b is None + or self._physical_output is None + ): + raise RuntimeError("H100 GEMM setup did not initialize required state") + # Handle creation and its internal allocations are steady-state setup. + self.extension.prepare_cublas_gemm(self.matrix_a, self.matrix_b, self._physical_output) + + def benchmark_fn(self) -> None: + if ( + self.extension is None + or self.matrix_a is None + or self.matrix_b is None + or self._physical_output is None + or self._logical_output is None + ): + raise RuntimeError("setup() must run before benchmark_fn()") + with torch.inference_mode(), self._nvtx_range(self.nvtx_label): + self.extension.run_cublas_gemm(self.matrix_a, self.matrix_b, self._physical_output) + self.output = self._logical_output + + +def get_benchmark() -> BaseBenchmark: + return BaselineH100Bf16GemmBenchmark() diff --git a/code/labs/fast_cu/baseline_int32_reduction.py b/code/labs/fast_cu/baseline_int32_reduction.py new file mode 100644 index 000000000..dace44d53 --- /dev/null +++ b/code/labs/fast_cu/baseline_int32_reduction.py @@ -0,0 +1,47 @@ +"""Baseline: CUB int32 reduction for the fast.cu comparison.""" + +from __future__ import annotations + +import torch + +from core.harness.benchmark_harness import BaseBenchmark +from labs.fast_cu.h100_common import Int32ReductionBenchmarkBase + + +class BaselineInt32ReductionBenchmark(Int32ReductionBenchmarkBase): + """Reduce the complete int32 input with CUB on the active stream.""" + + nvtx_label = "fast_cu_baseline_int32_reduction" + + def __init__(self) -> None: + super().__init__() + self._temp_storage: torch.Tensor | None = None + + def setup(self) -> None: + self._setup_tensors() + if self.extension is None or self.input is None or self._output_buffer is None: + raise RuntimeError("int32 reduction setup did not initialize required state") + storage_bytes = self.extension.cub_temp_storage_bytes(self.input, self._output_buffer) + if storage_bytes <= 0: + raise RuntimeError(f"CUB returned invalid temporary storage size: {storage_bytes}") + self._temp_storage = torch.empty((storage_bytes,), device=self.device, dtype=torch.uint8) + + def benchmark_fn(self) -> None: + if ( + self.extension is None + or self.input is None + or self._output_buffer is None + or self._temp_storage is None + ): + raise RuntimeError("setup() must run before benchmark_fn()") + with torch.inference_mode(), self._nvtx_range(self.nvtx_label): + self.extension.run_cub_reduction(self.input, self._output_buffer, self._temp_storage) + self.output = self._output_buffer + + def teardown(self) -> None: + self._temp_storage = None + super().teardown() + + +def get_benchmark() -> BaseBenchmark: + return BaselineInt32ReductionBenchmark() diff --git a/code/labs/fast_cu/baseline_nvfp4_gemm.py b/code/labs/fast_cu/baseline_nvfp4_gemm.py new file mode 100644 index 000000000..1dcf54a66 --- /dev/null +++ b/code/labs/fast_cu/baseline_nvfp4_gemm.py @@ -0,0 +1,20 @@ +"""cuBLASLt control for the pinned fast.cu GB300 NVFP4 workload.""" + +from __future__ import annotations + +from core.harness.benchmark_harness import BaseBenchmark +from labs.fast_cu.nvfp4 import FastCuNvfp4Benchmark, Nvfp4Workload + + +class FastCuNvfp4CublasLtBenchmark(FastCuNvfp4Benchmark): + """First correctness-gated cuBLASLt NVFP4 algorithm on the shared inputs.""" + + def __init__(self, workload: Nvfp4Workload | None = None) -> None: + super().__init__(optimized=False, workload=workload) + + +def get_benchmark() -> BaseBenchmark: + return FastCuNvfp4CublasLtBenchmark() + + +__all__ = ["FastCuNvfp4CublasLtBenchmark", "get_benchmark"] diff --git a/code/labs/fast_cu/build.py b/code/labs/fast_cu/build.py new file mode 100644 index 000000000..6ccec8869 --- /dev/null +++ b/code/labs/fast_cu/build.py @@ -0,0 +1,89 @@ +"""Content-keyed, setup-only CUDA extension builds for the fast.cu lab.""" + +from __future__ import annotations + +import hashlib +import json +import re +import subprocess +from pathlib import Path + +from labs.fast_cu.source import verify_upstream + + +def parse_nvcc_release(output: str) -> tuple[int, int]: + match = re.search(r"\brelease\s+(\d+)\.(\d+)\b", output) + if match is None: + raise RuntimeError("Cannot determine CUDA toolkit release from nvcc --version") + return int(match[1]), int(match[2]) + + +def extension_name( + name: str, + sources: list[Path], + cuda_flags: list[str], + link_flags: list[str], + manifest: dict, + *, + toolchain_identity: str = "", +) -> str: + """Include transitive upstream headers in the JIT cache identity.""" + if not re.fullmatch(r"[a-zA-Z][a-zA-Z0-9_]*", name): + raise ValueError("Extension name must be a C/Python identifier") + if not sources: + raise ValueError("At least one CUDA extension source is required") + payload = { + "sources": [(path.name, hashlib.sha256(path.read_bytes()).hexdigest()) for path in sources], + "upstream_revision": manifest["revision"], + "upstream_files": manifest["files"], + "cuda_flags": cuda_flags, + "link_flags": link_flags, + "toolchain_identity": toolchain_identity, + } + digest = hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() + return f"{name}_{digest[:20]}" + + +def load_cuda_extension( + name: str, + sources: list[Path], + *, + extra_cuda_cflags: list[str], + extra_ldflags: list[str], + minimum_cuda: tuple[int, int], +): + """Build outside timing; never install dependencies or change global flags.""" + manifest = verify_upstream() + from torch.utils.cpp_extension import CUDA_HOME, load + + if CUDA_HOME is None: + raise RuntimeError("SKIPPED: fast.cu requires a local CUDA toolkit with nvcc") + nvcc = Path(CUDA_HOME) / "bin" / "nvcc" + if not nvcc.is_file(): + raise RuntimeError(f"SKIPPED: fast.cu nvcc is missing: {nvcc}") + compiler_version = subprocess.run( + [str(nvcc), "--version"], check=True, capture_output=True, text=True + ).stdout + version = parse_nvcc_release(compiler_version) + if version < minimum_cuda: + required = ".".join(map(str, minimum_cuda)) + raise RuntimeError( + f"SKIPPED: fast.cu requires CUDA toolkit {required} or newer; found {version}" + ) + module_name = extension_name( + name, + sources, + extra_cuda_cflags, + extra_ldflags, + manifest, + toolchain_identity=f"{nvcc.resolve()}\n{compiler_version}", + ) + return load( + name=module_name, + sources=[str(path) for path in sources], + extra_cflags=["-O3"], + extra_cuda_cflags=extra_cuda_cflags, + extra_ldflags=extra_ldflags, + with_cuda=True, + verbose=False, + ) diff --git a/code/labs/fast_cu/h100_common.py b/code/labs/fast_cu/h100_common.py new file mode 100644 index 000000000..062d33383 --- /dev/null +++ b/code/labs/fast_cu/h100_common.py @@ -0,0 +1,352 @@ +"""Shared runtime gates, JIT loaders, and benchmark contracts for fast.cu.""" + +from __future__ import annotations + +from functools import lru_cache +from pathlib import Path +from types import ModuleType + +import torch + +from core.benchmark.verification import InputSignature, PrecisionFlags +from core.benchmark.verification_mixin import VerificationPayloadMixin +from core.harness.benchmark_harness import BaseBenchmark, BenchmarkConfig +from labs.fast_cu.build import load_cuda_extension + +_LAB_DIR = Path(__file__).resolve().parent + +H100_CAPABILITY = (9, 0) +REDUCTION_CAPABILITIES = frozenset({H100_CAPABILITY, (10, 0)}) + +GEMM_M = 8192 +GEMM_N = 8192 +GEMM_K = 8192 +REDUCTION_ELEMENTS = 1 << 28 + + +def ensure_h100_capability_supported( + capability: tuple[int, int], *, module_name: str = "fast.cu H100 BF16 GEMM" +) -> None: + """Require the exact Hopper capability targeted by the upstream WGMMA kernel.""" + if capability == H100_CAPABILITY: + return + major, minor = capability + raise RuntimeError( + f"SKIPPED: {module_name} requires exact SM 9.0 (sm_90a); got SM {major}.{minor}." + ) + + +def ensure_reduction_capability_supported( + capability: tuple[int, int], *, module_name: str = "fast.cu int32 reduction" +) -> None: + """Require an architecture explicitly compiled and exercised by this adapter.""" + if capability in REDUCTION_CAPABILITIES: + return + major, minor = capability + raise RuntimeError( + f"SKIPPED: {module_name} supports exact SM 9.0 (H100) and SM 10.0 (B200); " + f"got SM {major}.{minor}." + ) + + +def _visible_capability(module_name: str) -> tuple[int, int]: + if not torch.cuda.is_available(): + raise RuntimeError(f"SKIPPED: {module_name} requires an NVIDIA CUDA GPU") + return tuple(torch.cuda.get_device_capability()) + + +def ensure_h100_gemm_supported() -> tuple[int, int]: + capability = _visible_capability("fast.cu H100 BF16 GEMM") + ensure_h100_capability_supported(capability) + return capability + + +def ensure_int32_reduction_supported() -> tuple[int, int]: + capability = _visible_capability("fast.cu int32 reduction") + ensure_reduction_capability_supported(capability) + return capability + + +def _gemm_cuda_flags(capability: tuple[int, int]) -> list[str]: + ensure_h100_capability_supported(capability) + return [ + "-O3", + "-std=c++17", + "-DNDEBUG", + "--expt-relaxed-constexpr", + # torch cpp_extension disables this conversion by default, but pinned + # matmul_12.cuh converts FP32 accumulators to BF16 at its store epilogue. + # Extra CUDA flags follow torch's defaults, so this narrowly restores + # the constructor without changing the unchanged upstream header. + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + "-gencode=arch=compute_90a,code=sm_90a", + f"-I{_LAB_DIR}", + ] + + +def _reduction_build_contract( + capability: tuple[int, int], +) -> tuple[list[str], tuple[int, int]]: + ensure_reduction_capability_supported(capability) + if capability == H100_CAPABILITY: + architecture = "90a" + minimum_cuda = (12, 0) + else: + architecture = "100a" + minimum_cuda = (12, 8) + return ( + [ + "-O3", + "-std=c++17", + f"-gencode=arch=compute_{architecture},code=sm_{architecture}", + ], + minimum_cuda, + ) + + +@lru_cache(maxsize=1) +def load_h100_gemm_extension() -> ModuleType: + capability = ensure_h100_gemm_supported() + return load_cuda_extension( + "fast_cu_h100_gemm", + [_LAB_DIR / "h100_gemm_extension.cu"], + extra_cuda_cflags=_gemm_cuda_flags(capability), + extra_ldflags=["-lcublas", "-lcuda"], + minimum_cuda=(12, 0), + ) + + +@lru_cache(maxsize=2) +def _load_reduction_extension(capability: tuple[int, int]) -> ModuleType: + flags, minimum_cuda = _reduction_build_contract(capability) + return load_cuda_extension( + "fast_cu_int32_reduction", + [_LAB_DIR / "h100_reduction_kernels.cu"], + extra_cuda_cflags=flags, + extra_ldflags=[], + minimum_cuda=minimum_cuda, + ) + + +def load_int32_reduction_extension() -> ModuleType: + return _load_reduction_extension(ensure_int32_reduction_supported()) + + +class H100Bf16GemmBenchmarkBase(VerificationPayloadMixin, BaseBenchmark): + """Common full-output contract for A[M,K] @ B[N,K].T.""" + + allow_cpu = False + matrix_rows = GEMM_M + matrix_cols = GEMM_N + shared_dim = GEMM_K + # Upstream's own full-matrix check uses an absolute 0.1 bound. Retain that + # bound here; target-H100 qualification must establish any future change. + output_tolerance = (1e-2, 1e-1) + + def __init__(self) -> None: + super().__init__() + self.matrix_a: torch.Tensor | None = None + self.matrix_b: torch.Tensor | None = None + self._physical_output: torch.Tensor | None = None + self._logical_output: torch.Tensor | None = None + self.output: torch.Tensor | None = None + self.extension: ModuleType | None = None + element_size = torch.finfo(torch.bfloat16).bits // 8 + self.register_workload_metadata( + bytes_per_iteration=float( + ( + self.matrix_rows * self.shared_dim + + self.matrix_cols * self.shared_dim + + self.matrix_rows * self.matrix_cols + ) + * element_size + ), + custom_units_per_iteration=float( + 2 * self.matrix_rows * self.matrix_cols * self.shared_dim + ), + custom_unit_name="FLOPs", + ) + + def _clear_runtime_state(self) -> None: + """Release tensors retained directly or through the verification payload.""" + self._verification_payload = None + self.matrix_a = None + self.matrix_b = None + self._physical_output = None + self._logical_output = None + self.output = None + self.extension = None + + def _setup_tensors(self) -> None: + # VerifyRunner can set up the same instance again with fresh inputs. Drop + # the previous full payload before allocating another multi-GiB workload. + self._clear_runtime_state() + ensure_h100_gemm_supported() + self.extension = load_h100_gemm_extension() + # The harness owns the active RNG seed. Do not reset it here. + self.matrix_a = torch.empty( + (self.matrix_rows, self.shared_dim), + device=self.device, + dtype=torch.bfloat16, + ).uniform_(-1.0, 1.0) + self.matrix_b = torch.empty( + (self.matrix_cols, self.shared_dim), + device=self.device, + dtype=torch.bfloat16, + ).uniform_(-1.0, 1.0) + # The upstream TMA store writes C physically as [N, M]. The transpose is + # a metadata-only logical [M, N] view and is shared by both variants. + self._physical_output = torch.empty( + (self.matrix_cols, self.matrix_rows), + device=self.device, + dtype=torch.bfloat16, + ) + self._logical_output = self._physical_output.transpose(0, 1) + self.output = None + + def capture_verification_payload(self) -> None: + if self.matrix_a is None or self.matrix_b is None or self.output is None: + raise RuntimeError("benchmark_fn() must run before capture_verification_payload()") + self._set_verification_payload( + inputs={"matrix_a": self.matrix_a, "matrix_b": self.matrix_b}, + output=self.output, + batch_size=self.matrix_rows, + parameter_count=self.matrix_a.numel() + self.matrix_b.numel(), + precision_flags={"fp16": False, "bf16": True, "fp8": False, "tf32": False}, + output_tolerance=self.output_tolerance, + ) + + def get_input_signature(self) -> InputSignature: + return InputSignature( + shapes={ + "matrix_a": (self.matrix_rows, self.shared_dim), + "matrix_b": (self.matrix_cols, self.shared_dim), + "output": (self.matrix_rows, self.matrix_cols), + }, + dtypes={ + "matrix_a": str(torch.bfloat16), + "matrix_b": str(torch.bfloat16), + "output": str(torch.bfloat16), + }, + batch_size=self.matrix_rows, + parameter_count=( + self.matrix_rows * self.shared_dim + self.matrix_cols * self.shared_dim + ), + precision_flags=PrecisionFlags(bf16=True, tf32=False), + ) + + def validate_result(self) -> str | None: + if self.output is None: + return "benchmark_fn() did not produce output" + if self.output.shape != (self.matrix_rows, self.matrix_cols): + return f"unexpected output shape: {tuple(self.output.shape)}" + if self.output.dtype is not torch.bfloat16: + return f"unexpected output dtype: {self.output.dtype}" + return None + + def get_config(self) -> BenchmarkConfig: + return BenchmarkConfig(iterations=20, warmup=5, use_subprocess=False) + + def teardown(self) -> None: + self._clear_runtime_state() + super().teardown() + + +class Int32ReductionBenchmarkBase(VerificationPayloadMixin, BaseBenchmark): + """Common overflow-safe input and full scalar-output reduction contract.""" + + allow_cpu = False + element_count = REDUCTION_ELEMENTS + + def __init__(self) -> None: + super().__init__() + self.input: torch.Tensor | None = None + self._output_buffer: torch.Tensor | None = None + self.output: torch.Tensor | None = None + self.extension: ModuleType | None = None + self.register_workload_metadata( + bytes_per_iteration=float(self.element_count * torch.iinfo(torch.int32).bits // 8), + custom_units_per_iteration=float(self.element_count), + custom_unit_name="elements", + ) + + def _clear_runtime_state(self) -> None: + """Release tensors retained directly or through the verification payload.""" + self._verification_payload = None + self.input = None + self._output_buffer = None + self.output = None + self.extension = None + + def _setup_tensors(self) -> None: + # Verification may repeat setup on this instance with a fresh input. + # Clear the previous 1 GiB input payload before allocating its successor. + self._clear_runtime_state() + ensure_int32_reduction_supported() + self.extension = load_int32_reduction_extension() + # Caller-owned RNG, with a bounded domain: every possible sum fits int32, + # so both CUB and the upstream-derived atomic tree have defined arithmetic. + self.input = torch.randint( + -1, + 2, + (self.element_count,), + device=self.device, + dtype=torch.int32, + ) + self._output_buffer = torch.empty((1,), device=self.device, dtype=torch.int32) + self.output = None + + def capture_verification_payload(self) -> None: + if self.input is None or self.output is None: + raise RuntimeError("benchmark_fn() must run before capture_verification_payload()") + self._set_verification_payload( + inputs={"input": self.input}, + output=self.output, + batch_size=self.element_count, + parameter_count=0, + precision_flags={"fp16": False, "bf16": False, "fp8": False, "tf32": False}, + output_tolerance=(0.0, 0.0), + ) + + def get_input_signature(self) -> InputSignature: + return InputSignature( + shapes={"input": (self.element_count,), "output": (1,)}, + dtypes={"input": str(torch.int32), "output": str(torch.int32)}, + batch_size=self.element_count, + parameter_count=0, + precision_flags=PrecisionFlags(tf32=False), + ) + + def validate_result(self) -> str | None: + if self.output is None: + return "benchmark_fn() did not produce output" + if self.output.shape != (1,): + return f"unexpected output shape: {tuple(self.output.shape)}" + if self.output.dtype is not torch.int32: + return f"unexpected output dtype: {self.output.dtype}" + return None + + def get_config(self) -> BenchmarkConfig: + return BenchmarkConfig(iterations=50, warmup=10, use_subprocess=False) + + def teardown(self) -> None: + self._clear_runtime_state() + super().teardown() + + +__all__ = [ + "GEMM_K", + "GEMM_M", + "GEMM_N", + "H100Bf16GemmBenchmarkBase", + "H100_CAPABILITY", + "Int32ReductionBenchmarkBase", + "REDUCTION_CAPABILITIES", + "REDUCTION_ELEMENTS", + "ensure_h100_capability_supported", + "ensure_h100_gemm_supported", + "ensure_int32_reduction_supported", + "ensure_reduction_capability_supported", + "load_h100_gemm_extension", + "load_int32_reduction_extension", +] diff --git a/code/labs/fast_cu/h100_gemm_extension.cu b/code/labs/fast_cu/h100_gemm_extension.cu new file mode 100644 index 000000000..490819209 --- /dev/null +++ b/code/labs/fast_cu/h100_gemm_extension.cu @@ -0,0 +1,371 @@ +// Repository adapter for pranjalssh/fast.cu's latest Hopper BF16 GEMM. +// +// The vendored matmul_12 kernel is included unchanged. This translation unit +// supplies strict tensor validation, setup-only TMA/schedule construction, +// cuBLAS parity, and current-PyTorch-stream launches. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +using std::max; +using bf16 = __nv_bfloat16; + +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +namespace { + +void checked_cuda(cudaError_t status, const char* file, int line) { + TORCH_CHECK( + status == cudaSuccess, + "CUDA failure in vendored fast.cu host helper at ", + file, + ":", + line, + ": ", + cudaGetErrorString(status)); +} + +} // namespace + +#define cudaCheck(error) checked_cuda((error), __FILE__, __LINE__) + +// Upstream revision 2dfe5e26aecfd9e5f27bf9d5837deea01acda24b. +// NDEBUG is intentionally supplied by h100_common.py because the pinned header +// contains an upstream-only debug assert with an undeclared identifier. Every +// launch precondition hidden by that release-mode build is checked below. +#include "upstream/h100/matmul/matmul_12.cuh" + +namespace { + +constexpr int kBlockM = 128; +constexpr int kBlockN = 256; +constexpr int kBlockK = 64; +constexpr int kThreads = 128 * 3; +constexpr int kStages = 3; +constexpr int kClusterM = 2; +constexpr int kClusterN = 1; +constexpr int kGridCtas = 128; +constexpr int kScheduleClusters = kGridCtas / (kClusterM * kClusterN); +constexpr int kScheduleEntries = M12::SPACE_LEN; + +void check_cublas(cublasStatus_t status, const char* operation) { + TORCH_CHECK( + status == CUBLAS_STATUS_SUCCESS, + operation, + " failed with cuBLAS status ", + static_cast(status)); +} + +void check_exact_h100(int device) { + cudaDeviceProp properties{}; + C10_CUDA_CHECK(cudaGetDeviceProperties(&properties, device)); + TORCH_CHECK( + properties.major == 9 && properties.minor == 0, + "fast.cu H100 BF16 GEMM requires exact SM 9.0 (sm_90a); got SM ", + properties.major, + ".", + properties.minor); +} + +void check_gemm_tensors( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& physical_output) { + TORCH_CHECK(a.is_cuda() && b.is_cuda() && physical_output.is_cuda(), "all tensors must be CUDA tensors"); + TORCH_CHECK(a.device() == b.device() && a.device() == physical_output.device(), "all tensors must use one CUDA device"); + TORCH_CHECK(a.scalar_type() == at::kBFloat16, "A must have dtype torch.bfloat16"); + TORCH_CHECK(b.scalar_type() == at::kBFloat16, "B must have dtype torch.bfloat16"); + TORCH_CHECK(physical_output.scalar_type() == at::kBFloat16, "output must have dtype torch.bfloat16"); + TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && physical_output.dim() == 2, "A, B, and output must be rank-2 tensors"); + TORCH_CHECK(a.is_contiguous() && b.is_contiguous() && physical_output.is_contiguous(), "A, B, and physical output must be contiguous"); + TORCH_CHECK(a.size(1) == b.size(1), "A[M,K] and B[N,K] must share K"); + + const int64_t m = a.size(0); + const int64_t n = b.size(0); + const int64_t k = a.size(1); + TORCH_CHECK(m > 0 && n > 0 && k > 0, "GEMM dimensions must be positive"); + TORCH_CHECK( + m <= std::numeric_limits::max() && + n <= std::numeric_limits::max() && + k <= std::numeric_limits::max(), + "GEMM dimensions exceed the upstream int32 ABI"); + TORCH_CHECK( + physical_output.size(0) == n && physical_output.size(1) == m, + "physical output must have transposed [N,M] layout; expected [", + n, + ",", + m, + "]"); + TORCH_CHECK(m % (kBlockM * kClusterM) == 0, "M must be divisible by 256"); + TORCH_CHECK(n % kBlockN == 0, "N must be divisible by 256"); + TORCH_CHECK(k % kBlockK == 0, "K must be divisible by 64"); + TORCH_CHECK( + std::max(m / (kBlockM * kClusterM), n / kBlockN) > 1, + "upstream Hilbert schedule requires at least two tiles along M or N"); + const int64_t cluster_tiles = + (m / (kBlockM * kClusterM)) * (n / (kBlockN * kClusterN)); + TORCH_CHECK( + cluster_tiles <= static_cast(kScheduleClusters) * kScheduleEntries, + "upstream schedule capacity exceeded: ", + cluster_tiles, + " cluster tiles for ", + kScheduleClusters * kScheduleEntries, + " slots"); + TORCH_CHECK(reinterpret_cast(a.data_ptr()) % 16 == 0, "A must be 16-byte aligned"); + TORCH_CHECK(reinterpret_cast(b.data_ptr()) % 16 == 0, "B must be 16-byte aligned"); + TORCH_CHECK(reinterpret_cast(physical_output.data_ptr()) % 16 == 0, "output must be 16-byte aligned"); + check_exact_h100(a.get_device()); +} + +template +CUtensorMap checked_tensor_map(bf16* pointer, int height, int width) { + static_assert(BlockMinorSize >= 64); + TORCH_CHECK(width % 64 == 0, "TMA width must be divisible by 64"); + CUtensorMap tensor_map{}; + uint64_t global_shape[5] = { + 64, static_cast(height), static_cast(width / 64), 1, 1}; + uint64_t global_stride[5] = { + sizeof(bf16) * static_cast(width), 64 * sizeof(bf16), 0, 0, 0}; + uint32_t box_shape[5] = { + Padding ? 72u : 64u, + static_cast(BlockMajorSize), + static_cast(BlockMinorSize / 64), + 1, + 1}; + uint32_t element_stride[5] = {1, 1, 1, 1, 1}; + const CUresult result = cuTensorMapEncodeTiled( + &tensor_map, + CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + 3, + pointer, + global_shape, + global_stride, + box_shape, + element_stride, + CU_TENSOR_MAP_INTERLEAVE_NONE, + Swizzle ? CU_TENSOR_MAP_SWIZZLE_128B : CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + const char* error = nullptr; + if (result != CUDA_SUCCESS) { + cuGetErrorString(result, &error); + } + TORCH_CHECK( + result == CUDA_SUCCESS, + "cuTensorMapEncodeTiled failed: ", + error == nullptr ? "unknown CUDA driver error" : error); + return tensor_map; +} + +struct CublasState { + cublasHandle_t handle = nullptr; + int device = -1; + bool ready = false; +}; + +struct UpstreamGemmPlan { + bool ready = false; + int device = -1; + int m = 0; + int n = 0; + int k = 0; + const void* a_pointer = nullptr; + const void* b_pointer = nullptr; + void* output_pointer = nullptr; + CUtensorMap a_map{}; + CUtensorMap b_map{}; + CUtensorMap output_map{}; + torch::Tensor schedule; +}; + +CublasState g_cublas; +UpstreamGemmPlan g_upstream; + +void prepare_cublas_gemm( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& physical_output) { + check_gemm_tensors(a, b, physical_output); + const c10::cuda::CUDAGuard guard(a.device()); + const int device = a.get_device(); + if (g_cublas.handle != nullptr && g_cublas.device != device) { + check_cublas(cublasDestroy(g_cublas.handle), "cublasDestroy"); + g_cublas = CublasState{}; + } + if (g_cublas.handle == nullptr) { + check_cublas(cublasCreate(&g_cublas.handle), "cublasCreate"); + check_cublas( + cublasSetMathMode(g_cublas.handle, CUBLAS_TENSOR_OP_MATH), + "cublasSetMathMode"); + } + g_cublas.device = device; + g_cublas.ready = true; +} + +void run_cublas_gemm( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& physical_output) { + check_gemm_tensors(a, b, physical_output); + const c10::cuda::CUDAGuard guard(a.device()); + TORCH_CHECK( + g_cublas.ready && g_cublas.handle != nullptr && g_cublas.device == a.get_device(), + "prepare_cublas_gemm() must run on this device before timing"); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device()); + check_cublas(cublasSetStream(g_cublas.handle, stream), "cublasSetStream"); + + const int m = static_cast(a.size(0)); + const int n = static_cast(b.size(0)); + const int k = static_cast(a.size(1)); + const float alpha = 1.0f; + const float beta = 0.0f; + check_cublas( + cublasGemmEx( + g_cublas.handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + m, + n, + k, + &alpha, + a.data_ptr(), + CUDA_R_16BF, + k, + b.data_ptr(), + CUDA_R_16BF, + k, + &beta, + physical_output.data_ptr(), + CUDA_R_16BF, + m, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP), + "cublasGemmEx"); +} + +void prepare_upstream_gemm( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& physical_output) { + check_gemm_tensors(a, b, physical_output); + const c10::cuda::CUDAGuard guard(a.device()); + const int m = static_cast(a.size(0)); + const int n = static_cast(b.size(0)); + const int k = static_cast(a.size(1)); + + UpstreamGemmPlan next{}; + next.device = a.get_device(); + next.m = m; + next.n = n; + next.k = k; + next.a_pointer = a.data_ptr(); + next.b_pointer = b.data_ptr(); + next.output_pointer = physical_output.data_ptr(); + next.a_map = checked_tensor_map( + reinterpret_cast(a.data_ptr()), m, k); + next.b_map = checked_tensor_map( + reinterpret_cast(b.data_ptr()), n, k); + next.output_map = checked_tensor_map< + kBlockN, kBlockM / ((kThreads / 128) - 1), false, true>( + reinterpret_cast(physical_output.data_ptr()), n, m); + + std::vector host_schedule(kGridCtas * kScheduleEntries, -1); + M12::createHilbert( + m / (kBlockM * kClusterM), + n / (kBlockN * kClusterN), + kScheduleClusters, + host_schedule.data()); + next.schedule = torch::empty( + {static_cast(host_schedule.size())}, + a.options().dtype(torch::kInt32)); + C10_CUDA_CHECK(cudaMemcpy( + next.schedule.data_ptr(), + host_schedule.data(), + host_schedule.size() * sizeof(int), + cudaMemcpyHostToDevice)); + + auto* kernel = M12::matmulKernel12< + kBlockM, + kBlockN, + kBlockK, + kThreads, + kStages, + kGridCtas, + kClusterM, + kClusterN>; + constexpr size_t shared_bytes = + sizeof(M12::SMem); + static_assert(shared_bytes < 256 * 1024); + C10_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes)); + next.ready = true; + g_upstream = std::move(next); +} + +void run_upstream_gemm( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& physical_output) { + check_gemm_tensors(a, b, physical_output); + const c10::cuda::CUDAGuard guard(a.device()); + TORCH_CHECK(g_upstream.ready, "prepare_upstream_gemm() must run before timing"); + TORCH_CHECK(g_upstream.device == a.get_device(), "prepared GEMM device changed"); + TORCH_CHECK( + g_upstream.m == a.size(0) && + g_upstream.n == b.size(0) && + g_upstream.k == a.size(1), + "prepared GEMM dimensions changed"); + TORCH_CHECK( + g_upstream.a_pointer == a.data_ptr() && + g_upstream.b_pointer == b.data_ptr() && + g_upstream.output_pointer == physical_output.data_ptr(), + "prepared GEMM storage changed; call prepare_upstream_gemm() again"); + + auto* kernel = M12::matmulKernel12< + kBlockM, + kBlockN, + kBlockK, + kThreads, + kStages, + kGridCtas, + kClusterM, + kClusterN>; + constexpr size_t shared_bytes = + sizeof(M12::SMem); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device()); + kernel<<>>( + g_upstream.m, + g_upstream.n, + g_upstream.k, + g_upstream.output_map, + g_upstream.a_map, + g_upstream.b_map, + g_upstream.schedule.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("prepare_cublas_gemm", &prepare_cublas_gemm); + module.def("run_cublas_gemm", &run_cublas_gemm); + module.def("prepare_upstream_gemm", &prepare_upstream_gemm); + module.def("run_upstream_gemm", &run_upstream_gemm); +} diff --git a/code/labs/fast_cu/h100_reduction_kernels.cu b/code/labs/fast_cu/h100_reduction_kernels.cu new file mode 100644 index 000000000..43ddd9040 --- /dev/null +++ b/code/labs/fast_cu/h100_reduction_kernels.cu @@ -0,0 +1,174 @@ +// Safe, stream-correct adapter for pranjalssh/fast.cu h100/sum.cu. +// +// The upstream vectorized reduction resets the global output from one CTA while +// other CTAs can already atomicAdd into it. This adaptation keeps its int4 load, +// per-thread batching, warp shuffle, shared warp sums, and global atomic, while +// moving the reset to cudaMemsetAsync before launch on the same PyTorch stream. +// It intentionally lives outside the WGMMA translation unit so the same source +// can compile natively for both H100 SM90a and B200 SM100a. + +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace { + +constexpr int kWarpSize = 32; +constexpr int kBlockSize = 1024; +constexpr int kBatchVectors = 16; + +void check_supported_device(int device) { + cudaDeviceProp properties{}; + C10_CUDA_CHECK(cudaGetDeviceProperties(&properties, device)); + const bool supported = + (properties.major == 9 && properties.minor == 0) || + (properties.major == 10 && properties.minor == 0); + TORCH_CHECK( + supported, + "fast.cu int32 reduction supports exact SM 9.0 (H100) and SM 10.0 (B200); got SM ", + properties.major, + ".", + properties.minor); +} + +void check_reduction_tensors( + const torch::Tensor& input, const torch::Tensor& output) { + TORCH_CHECK(input.is_cuda() && output.is_cuda(), "input and output must be CUDA tensors"); + TORCH_CHECK(input.device() == output.device(), "input and output must use one CUDA device"); + TORCH_CHECK(input.scalar_type() == at::kInt, "input must have dtype torch.int32"); + TORCH_CHECK(output.scalar_type() == at::kInt, "output must have dtype torch.int32"); + TORCH_CHECK(input.dim() == 1 && output.dim() == 1, "input and output must be rank-1 tensors"); + TORCH_CHECK(input.is_contiguous() && output.is_contiguous(), "input and output must be contiguous"); + TORCH_CHECK(output.numel() == 1, "output must contain exactly one int32 value"); + TORCH_CHECK(input.numel() > 0, "input must be non-empty"); + TORCH_CHECK(input.numel() % 4 == 0, "input length must be divisible by four for int4 loads"); + TORCH_CHECK( + input.numel() <= std::numeric_limits::max(), + "input length exceeds the upstream int32 indexing ABI"); + TORCH_CHECK( + reinterpret_cast(input.data_ptr()) % alignof(int4) == 0, + "input storage must be 16-byte aligned for int4 loads"); + check_supported_device(input.get_device()); +} + +__device__ __forceinline__ int warp_reduce_sum(int value) { +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset); + } + return value; +} + +// Derived from upstream sumKernel2<1024, 16>. The only semantic change inside +// the kernel is removal of `if (i == 0) *d_out = 0`; reset ordering belongs to +// the host launch sequence below, where it cannot race another CTA. +__global__ __launch_bounds__(kBlockSize) void safe_vectorized_sum_kernel( + const int4* input, int* output, int vector_count) { + __shared__ __align__(16) int warp_sums[kWarpSize]; + + const unsigned int tid = threadIdx.x; + const unsigned int index = + blockIdx.x * kBlockSize * kBatchVectors + tid; + int4 values = index < vector_count + ? input[index] + : make_int4(0, 0, 0, 0); + int sum = values.x + values.y + values.z + values.w; +#pragma unroll + for (int batch = 1; batch < kBatchVectors; ++batch) { + const unsigned int next = index + batch * kBlockSize; + if (next < vector_count) { + values = input[next]; + sum += values.x + values.y + values.z + values.w; + } + } + + sum = warp_reduce_sum(sum); + if ((tid & (kWarpSize - 1)) == 0) { + warp_sums[tid / kWarpSize] = sum; + } + __syncthreads(); + + if (tid < kWarpSize) { + sum = warp_reduce_sum(warp_sums[tid]); + } + if (tid == 0) { + atomicAdd(output, sum); + } +} + +int64_t cub_temp_storage_bytes( + const torch::Tensor& input, const torch::Tensor& output) { + check_reduction_tensors(input, output); + const c10::cuda::CUDAGuard guard(input.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device()); + size_t storage_bytes = 0; + C10_CUDA_CHECK(cub::DeviceReduce::Sum( + nullptr, + storage_bytes, + input.data_ptr(), + output.data_ptr(), + static_cast(input.numel()), + stream)); + TORCH_CHECK(storage_bytes > 0, "CUB returned zero temporary storage bytes"); + TORCH_CHECK( + storage_bytes <= static_cast(std::numeric_limits::max()), + "CUB temporary storage size exceeds Python int64"); + return static_cast(storage_bytes); +} + +void run_cub_reduction( + const torch::Tensor& input, + const torch::Tensor& output, + const torch::Tensor& temp_storage) { + check_reduction_tensors(input, output); + TORCH_CHECK(temp_storage.is_cuda(), "CUB temporary storage must be a CUDA tensor"); + TORCH_CHECK(temp_storage.device() == input.device(), "CUB temporary storage must use the input device"); + TORCH_CHECK(temp_storage.scalar_type() == at::kByte, "CUB temporary storage must have dtype torch.uint8"); + TORCH_CHECK(temp_storage.dim() == 1 && temp_storage.is_contiguous(), "CUB temporary storage must be contiguous and rank 1"); + TORCH_CHECK(temp_storage.numel() > 0, "CUB temporary storage must be non-empty"); + + const c10::cuda::CUDAGuard guard(input.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device()); + size_t storage_bytes = static_cast(temp_storage.numel()); + C10_CUDA_CHECK(cub::DeviceReduce::Sum( + temp_storage.data_ptr(), + storage_bytes, + input.data_ptr(), + output.data_ptr(), + static_cast(input.numel()), + stream)); +} + +void run_safe_vectorized_reduction( + const torch::Tensor& input, const torch::Tensor& output) { + check_reduction_tensors(input, output); + const c10::cuda::CUDAGuard guard(input.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(input.get_device()); + + // This stream-ordered reset is part of every measured reduction. It must not + // be hoisted into setup because each launch accumulates into the same buffer. + C10_CUDA_CHECK(cudaMemsetAsync(output.data_ptr(), 0, sizeof(int), stream)); + const int vector_count = static_cast(input.numel() / 4); + const int vectors_per_block = kBlockSize * kBatchVectors; + const int blocks = (vector_count + vectors_per_block - 1) / vectors_per_block; + safe_vectorized_sum_kernel<<>>( + reinterpret_cast(input.data_ptr()), + output.data_ptr(), + vector_count); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("cub_temp_storage_bytes", &cub_temp_storage_bytes); + module.def("run_cub_reduction", &run_cub_reduction); + module.def("run_safe_vectorized_reduction", &run_safe_vectorized_reduction); +} diff --git a/code/labs/fast_cu/nvfp4.py b/code/labs/fast_cu/nvfp4.py new file mode 100644 index 000000000..44c81a12d --- /dev/null +++ b/code/labs/fast_cu/nvfp4.py @@ -0,0 +1,334 @@ +"""Native harness adapter for fast.cu's GB300 NVFP4 optimization ladder.""" + +from __future__ import annotations + +import re +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +from core.benchmark.verification_mixin import VerificationPayloadMixin +from core.harness.benchmark_harness import BaseBenchmark, BenchmarkConfig +from labs.fast_cu.build import load_cuda_extension, parse_nvcc_release +from labs.fast_cu.source import UPSTREAM_DIR, verify_upstream + +DEFAULT_M = 8192 +DEFAULT_N = 8192 +DEFAULT_K = 8192 +DEFAULT_RUNG = 9 +CUDA_VERSION = (13, 1) +COMPUTE_CAPABILITY = (10, 3) +OUTPUT_RTOL = 2e-3 +OUTPUT_ATOL = 0.5 +_NATIVE_SOURCE = Path(__file__).with_name("nvfp4_native.cu") +_UPSTREAM_NVFP4_DIR = UPSTREAM_DIR / "gb300" / "nvfp4" + + +@dataclass(frozen=True) +class Nvfp4Workload: + """The exact MxNxK workload shared by both comparison arms.""" + + m: int = DEFAULT_M + n: int = DEFAULT_N + k: int = DEFAULT_K + + def validate(self) -> None: + for name in ("m", "n", "k"): + value = getattr(self, name) + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive integer") + if self.k % 32: + raise ValueError( + "k must be divisible by 32 because the cuBLASLt NVFP4 baseline " + "has no supported algorithm otherwise" + ) + + +def validate_rung(rung: int) -> int: + if type(rung) is not int or not 0 <= rung <= 9: + raise ValueError("fast.cu NVFP4 rung must be an integer in [0, 9]") + return rung + + +def workload_bytes(workload: Nvfp4Workload) -> int: + """Packed A/B + VEC16 scales read and full FP16 C written per launch.""" + workload.validate() + ab_bytes = (workload.m + workload.n) * (workload.k // 2) + packed_scale_inner = ((workload.k + 63) // 64) * 4 + packed_scale_rows = ((workload.m + 127) // 128) * 128 + ((workload.n + 127) // 128) * 128 + scale_bytes = packed_scale_inner * packed_scale_rows + output_bytes = workload.m * workload.n * 2 + return ab_bytes + scale_bytes + output_bytes + + +def workload_flops(workload: Nvfp4Workload) -> int: + """Conventional GEMM operation count: one multiply and add per K.""" + workload.validate() + return 2 * workload.m * workload.n * workload.k + + +def _version_tuple(version: str | None) -> tuple[int, int] | None: + if not isinstance(version, str): + return None + match = re.fullmatch(r"(\d+)\.(\d+)(?:\.\d+)?", version) + if match is None: + return None + return int(match[1]), int(match[2]) + + +def require_nvfp4_runtime() -> int: + """Reject runtimes older than CUDA 13.1 and every GPU except SM103.""" + if not torch.cuda.is_available(): + raise RuntimeError("SKIPPED: fast.cu NVFP4 requires a CUDA 13.1 GB300/B300 (SM103) GPU") + runtime = _version_tuple(torch.version.cuda) + if runtime is None or runtime < CUDA_VERSION: + found = torch.version.cuda if torch.version.cuda is not None else "none" + raise RuntimeError( + f"SKIPPED: fast.cu NVFP4 requires CUDA runtime 13.1 or newer; found {found}" + ) + device = torch.cuda.current_device() + capability = torch.cuda.get_device_capability(device) + if capability != COMPUTE_CAPABILITY: + raise RuntimeError( + "SKIPPED: fast.cu NVFP4 requires exact SM103 GB300/B300; " + f"found sm_{capability[0]}{capability[1]} on " + f"{torch.cuda.get_device_name(device)}" + ) + return device + + +def require_nvfp4_toolkit() -> Path: + """Require nvcc 13.1+ before starting the content-keyed JIT build.""" + from torch.utils.cpp_extension import CUDA_HOME + + if CUDA_HOME is None: + raise RuntimeError("SKIPPED: fast.cu NVFP4 requires CUDA toolkit 13.1 with nvcc") + nvcc = Path(CUDA_HOME) / "bin" / "nvcc" + if not nvcc.is_file(): + raise RuntimeError(f"SKIPPED: fast.cu NVFP4 nvcc is missing: {nvcc}") + completed = subprocess.run( + [str(nvcc), "--version"], + check=True, + capture_output=True, + text=True, + ) + version = parse_nvcc_release(completed.stdout) + if version < CUDA_VERSION: + raise RuntimeError( + "SKIPPED: fast.cu NVFP4 requires CUDA toolkit 13.1 or newer; " + f"found {version[0]}.{version[1]}" + ) + return nvcc + + +def load_nvfp4_extension(rung: int = DEFAULT_RUNG): + """Verify provenance and runtime gates, then compile one explicit rung.""" + rung = validate_rung(rung) + manifest = verify_upstream() + hardware = manifest.get("hardware", {}).get("nvfp4", {}) + if hardware.get("compute_capability") != [10, 3] or hardware.get("minimum_cuda") != "13.1": + raise RuntimeError("Pinned fast.cu NVFP4 hardware manifest is inconsistent") + require_nvfp4_runtime() + require_nvfp4_toolkit() + return load_cuda_extension( + f"fast_cu_nvfp4_r{rung}", + [_NATIVE_SOURCE], + extra_cuda_cflags=[ + "-std=c++17", + "-O3", + "-DNDEBUG", + f"-DFAST_CU_NVFP4_RUNG={rung}", + f"-I{_UPSTREAM_NVFP4_DIR}", + "-gencode=arch=compute_103a,code=sm_103a", + ], + extra_ldflags=["-lcublasLt", "-lcuda"], + minimum_cuda=CUDA_VERSION, + ) + + +class FastCuNvfp4Benchmark(VerificationPayloadMixin, BaseBenchmark): + """One arm of the native single-buffer cuBLASLt-versus-fast.cu pair. + + The r9 route table is translation-unit global. Pybind keeps the GIL for all + context calls, setup synchronizes before replacing that table, and launch + rejects a table configured for another shape. Concurrent native calls that + bypass this Python/GIL boundary are unsupported. + """ + + allow_cpu = False + multi_gpu_required = False + required_world_size = 1 + + def __init__( + self, + *, + optimized: bool, + workload: Nvfp4Workload | None = None, + rung: int = DEFAULT_RUNG, + ) -> None: + super().__init__() + self.optimized = bool(optimized) + self.workload = workload or Nvfp4Workload() + self.workload.validate() + self.rung = validate_rung(rung) + self.output: torch.Tensor | None = None + self._output_buffer: torch.Tensor | None = None + self._native_module: Any | None = None + self._native_context: Any | None = None + self._inputs: dict[str, torch.Tensor] = {} + self._caller_seed: int | None = None + self._setup_gate_passed = False + self.register_workload_metadata( + requests_per_iteration=1.0, + bytes_per_iteration=float(workload_bytes(self.workload)), + custom_units_per_iteration=float(workload_flops(self.workload)), + custom_unit_name="nvfp4_flops", + ) + + def setup(self) -> None: + self._verification_payload = None + self.output = None + self._output_buffer = None + self._inputs = {} + self._native_context = None + self._native_module = None + self._setup_gate_passed = False + require_nvfp4_runtime() + self._native_module = load_nvfp4_extension(self.rung) + self._caller_seed = int(torch.initial_seed()) + m, n, k = self.workload.m, self.workload.n, self.workload.k + self._output_buffer = torch.empty((m, n), device=self.device, dtype=torch.float16) + self.output = None + self._native_context = self._native_module.Nvfp4Context( + m, + n, + k, + self._caller_seed, + self.optimized, + ) + if int(self._native_context.rung) != self.rung: + raise RuntimeError( + f"Compiled fast.cu rung {self._native_context.rung} does not match requested r{self.rung}" + ) + self._inputs = { + "a_packed_e2m1": self._native_context.a_packed, + "b_packed_e2m1": self._native_context.b_packed, + "sfa_vec16_ue4m3": self._native_context.sfa_packed, + "sfb_vec16_ue4m3": self._native_context.sfb_packed, + } + self._run_untimed_arm_gate() + + def _run_untimed_arm_gate(self) -> None: + """Prove full overwrite and determinism on separately poisoned outputs.""" + if self._native_context is None: + raise RuntimeError("native context is unavailable for setup validation") + m, n = self.workload.m, self.workload.n + first = torch.full((m, n), float("nan"), device=self.device, dtype=torch.float16) + second = torch.full((m, n), float("nan"), device=self.device, dtype=torch.float16) + launch = ( + self._native_context.launch_fast + if self.optimized + else self._native_context.launch_cublaslt + ) + launch(first) + launch(second) + torch.cuda.synchronize(self.device) + if not bool(torch.isfinite(first).all().item()) or not bool( + torch.isfinite(second).all().item() + ): + raise RuntimeError("untimed NVFP4 setup gate found an incomplete or non-finite output") + if not torch.equal(first, second): + raise RuntimeError("untimed NVFP4 setup gate found nondeterministic output") + if self.optimized: + self._native_context.validate_schedule() + self._setup_gate_passed = True + + def benchmark_fn(self) -> None: + if self._native_context is None or self._output_buffer is None: + raise RuntimeError("setup() must complete before benchmark_fn()") + if self.optimized: + self._native_context.launch_fast(self._output_buffer) + else: + self._native_context.launch_cublaslt(self._output_buffer) + self.output = self._output_buffer + + def capture_verification_payload(self) -> None: + if self._native_context is None or self.output is None or not self._inputs: + raise RuntimeError("benchmark_fn() must run before verification capture") + if not self._setup_gate_passed: + raise RuntimeError("untimed NVFP4 setup correctness gate did not pass") + if self.optimized: + self._native_context.validate_schedule() + self._set_verification_payload( + inputs=self._inputs, + output=self.output, + batch_size=self.workload.m, + parameter_count=0, + precision_flags={"fp16": True, "bf16": False, "fp8": False, "tf32": False}, + output_tolerance=(OUTPUT_RTOL, OUTPUT_ATOL), + signature_overrides={"quantization_mode": "nvfp4_e2m1_vec16_ue4m3"}, + ) + + def validate_result(self) -> str | None: + if self.output is None: + return "benchmark_fn() did not produce an output" + if self.output.dtype != torch.float16: + return f"expected FP16 output, found {self.output.dtype}" + if tuple(self.output.shape) != (self.workload.m, self.workload.n): + return f"unexpected output shape {tuple(self.output.shape)}" + if not bool(torch.isfinite(self.output).all().item()): + return "NVFP4 GEMM output contains non-finite values" + return None + + def get_config(self) -> BenchmarkConfig: + return BenchmarkConfig( + iterations=25, + warmup=5, + min_run_time_ms=100.0, + setup_timeout_seconds=1800, + measurement_timeout_seconds=600, + single_gpu=True, + ncu_metric_set="minimal", + ncu_replay_mode="kernel", + ncu_replay_mode_override=True, + ) + + def get_input_signature(self): + return super().get_input_signature() + + def get_verify_output(self): + return super().get_verify_output() + + def get_output_tolerance(self): + return super().get_output_tolerance() + + def teardown(self) -> None: + self._inputs = {} + self._native_context = None + self._native_module = None + self.output = None + self._output_buffer = None + self._verification_payload = None + self._setup_gate_passed = False + super().teardown() + + +__all__ = [ + "COMPUTE_CAPABILITY", + "CUDA_VERSION", + "DEFAULT_K", + "DEFAULT_M", + "DEFAULT_N", + "DEFAULT_RUNG", + "FastCuNvfp4Benchmark", + "Nvfp4Workload", + "load_nvfp4_extension", + "require_nvfp4_runtime", + "require_nvfp4_toolkit", + "validate_rung", + "workload_bytes", + "workload_flops", +] diff --git a/code/labs/fast_cu/nvfp4_native.cu b/code/labs/fast_cu/nvfp4_native.cu new file mode 100644 index 000000000..4046d2a8c --- /dev/null +++ b/code/labs/fast_cu/nvfp4_native.cu @@ -0,0 +1,569 @@ +#include + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef FAST_CU_NVFP4_RUNG +#error "FAST_CU_NVFP4_RUNG must select one upstream optimization rung" +#elif FAST_CU_NVFP4_RUNG == 0 +#define NVFP4_GEMM_HEADER "gemm0.cuh" +#elif FAST_CU_NVFP4_RUNG == 1 +#define NVFP4_GEMM_HEADER "gemm1.cuh" +#elif FAST_CU_NVFP4_RUNG == 2 +#define NVFP4_GEMM_HEADER "gemm2.cuh" +#elif FAST_CU_NVFP4_RUNG == 3 +#define NVFP4_GEMM_HEADER "gemm3.cuh" +#elif FAST_CU_NVFP4_RUNG == 4 +#define NVFP4_GEMM_HEADER "gemm4.cuh" +#elif FAST_CU_NVFP4_RUNG == 5 +#define NVFP4_GEMM_HEADER "gemm5.cuh" +#elif FAST_CU_NVFP4_RUNG == 6 +#define NVFP4_GEMM_HEADER "gemm6.cuh" +#elif FAST_CU_NVFP4_RUNG == 7 +#define NVFP4_GEMM_HEADER "gemm7.cuh" +#elif FAST_CU_NVFP4_RUNG == 8 +#define NVFP4_GEMM_HEADER "gemm8.cuh" +#elif FAST_CU_NVFP4_RUNG == 9 +#define NVFP4_GEMM_HEADER "gemm9.cuh" +#else +#error "FAST_CU_NVFP4_RUNG must be in [0, 9]" +#endif + +// Reuse the pinned upstream host packing, tensor-map, cuBLASLt, and kernel +// implementation without editing the vendored source. The standalone entrypoint +// is renamed because this translation unit supplies a Python module instead. +#define main fast_cu_nvfp4_upstream_standalone_main +#include "main.cu" +#undef main + +namespace py = pybind11; + +namespace { + +std::mutex small_gate_mutex; +std::vector small_gated_devices; + +void check_cuda(cudaError_t status, const char* operation) { + TORCH_CHECK( + status == cudaSuccess, + operation, + " failed: ", + cudaGetErrorName(status), + " (", cudaGetErrorString(status), ")"); +} + +void check_cublas(cublasStatus_t status, const char* operation) { + TORCH_CHECK( + status == CUBLAS_STATUS_SUCCESS, + operation, + " failed with cuBLASLt status ", + static_cast(status)); +} + +cudaStream_t current_stream(int device) { + int current_device = -1; + check_cuda(cudaGetDevice(¤t_device), "cudaGetDevice"); + TORCH_CHECK( + current_device == device, + "fast.cu NVFP4 context was created on CUDA device ", + device, + " but current device is ", + current_device); + return at::cuda::getCurrentCUDAStream(device); +} + +void require_exact_runtime(cudaDeviceProp* properties, int* device) { + int runtime_version = 0; + check_cuda(cudaRuntimeGetVersion(&runtime_version), "cudaRuntimeGetVersion"); + TORCH_CHECK( + runtime_version >= 13010, + "SKIPPED: fast.cu NVFP4 requires CUDA runtime 13.1 or newer; found ", + runtime_version); + check_cuda(cudaGetDevice(device), "cudaGetDevice"); + check_cuda(cudaGetDeviceProperties(properties, *device), "cudaGetDeviceProperties"); + TORCH_CHECK( + properties->major == 10 && properties->minor == 3, + "SKIPPED: fast.cu NVFP4 requires exact SM103 GB300/B300; found sm_", + properties->major, + properties->minor, + " on ", + properties->name); + TORCH_CHECK( + properties->multiProcessorCount > 0 && properties->multiProcessorCount % 2 == 0, + "fast.cu NVFP4 requires a positive even SM count for two-CTA clusters; found ", + properties->multiProcessorCount); +} + +torch::Tensor cuda_bytes(int64_t count, int device) { + return torch::empty( + {count}, + torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA, device)); +} + +#if NVFP4_HAS_SCHEDULE +// Pybind methods intentionally retain the GIL. Setup synchronizes inside the +// census before replacing these translation-unit-global constants, and every +// launch checks the active device/shape. Native callers bypassing pybind must +// provide their own serialization. +std::mutex schedule_mutex; +int scheduled_device = -1; +std::array scheduled_shape{-1, -1, -1}; + +void strict_setup_r9_schedule( + int device, + const cudaDeviceProp& properties, + dim3 grid, + int clusters, + int M, + int N, + int K, + size_t operand_read_bytes) { + std::lock_guard lock(schedule_mutex); + const int mblocks = (M + nvfp4::CLUSTER_M - 1) / nvfp4::CLUSTER_M; + const int nblocks = (N + nvfp4::BLOCK_N - 1) / nvfp4::BLOCK_N; + const int total_work = mblocks * nblocks; + TORCH_CHECK( + total_work <= nvfp4::L2A_ROUTE_WORK_CAP, + "fast.cu r9 schedule has ", + total_work, + " tiles but the pinned table holds only ", + nvfp4::L2A_ROUTE_WORK_CAP); + + const sched::ScheduleMode mode = + sched::pick_schedule(operand_read_bytes, mblocks, nblocks); + TORCH_CHECK( + sched::requires_side_census(mode), + "fast.cu r9 AUTO schedule unexpectedly disabled L2-side ownership"); + + constexpr int64_t probe_bytes = int64_t(2) << 20; + torch::Tensor probe_storage = cuda_bytes(2 * probe_bytes, device); + const uintptr_t raw = reinterpret_cast(probe_storage.data_ptr()); + char* arena = reinterpret_cast( + (raw + probe_bytes - 1) & ~(uintptr_t(probe_bytes) - 1)); + const l2side::RuntimeMap map = + l2side::probe_stable(arena, probe_bytes, /*repeats=*/3); + TORCH_CHECK( + map.nsm == properties.multiProcessorCount, + "fast.cu r9 L2 map covered ", map.nsm, + " SMs but the device reports ", properties.multiProcessorCount); + TORCH_CHECK( + map.hash == l2side::kExpectedHash && map.model_mismatches == 0, + "fast.cu r9 L2 address model failed: hash=", + map.hash, + " mismatches=", + map.model_mismatches); + + const size_t census_count = size_t(clusters) * 2; + torch::Tensor device_smids = torch::empty( + {static_cast(census_count)}, + torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA, device)); + cudaStream_t stream = current_stream(device); + check_cuda( + cudaMemsetAsync( + device_smids.data_ptr(), + 0xff, + census_count * sizeof(unsigned), + stream), + "cudaMemsetAsync(cluster census)"); + nvfp4::l2a_cluster_probe<<>>( + reinterpret_cast(device_smids.data_ptr())); + check_cuda(cudaGetLastError(), "fast.cu r9 cluster census launch"); + std::vector smids(census_count); + check_cuda( + cudaMemcpyAsync( + smids.data(), + device_smids.data_ptr(), + census_count * sizeof(unsigned), + cudaMemcpyDeviceToHost, + stream), + "cudaMemcpyAsync(cluster census)"); + check_cuda(cudaStreamSynchronize(stream), "cudaStreamSynchronize(cluster census)"); + + const sched::Census census = + sched::classify_census(smids, map, properties.multiProcessorCount); + sched::upload_side_constants(map, census); + std::vector table = sched::build_schedule(census, mblocks, nblocks, mode); + TORCH_CHECK( + !table.empty(), + "fast.cu r9 could not build a strict L2-owned schedule for ", + M, "x", N, "x", K, + "; raster fallback is forbidden"); + sched::upload_table(table, bench::MAIN_TABLE); + scheduled_device = device; + scheduled_shape = {M, N, K}; +} +#endif + +class Nvfp4Context { +public: + Nvfp4Context(int M, int N, int K, uint64_t seed, bool prepare_fast) + : M_(M), N_(N), K_(K), seed_(seed) { + TORCH_CHECK(M > 0 && N > 0 && K > 0, "M, N, and K must be positive"); + TORCH_CHECK( + K % 32 == 0, + "cuBLASLt NVFP4 comparison requires K divisible by 32; found ", + K); + + cudaDeviceProp properties{}; + require_exact_runtime(&properties, &device_); + c10::cuda::CUDAGuard guard(device_); + stream_ = current_stream(device_); + + host::Fixture fixture = host::make_fixture(M, N, K, seed, /*keep_deq=*/false); + TORCH_CHECK( + fixture.inputs.A == fixture.A_logical && + fixture.inputs.B == fixture.B_logical, + "fast.cu and cuBLASLt input bytes differ"); + + a_ = cuda_bytes(static_cast(fixture.inputs.A.size()), device_); + b_ = cuda_bytes(static_cast(fixture.inputs.B.size()), device_); + sfa_ = cuda_bytes(static_cast(fixture.inputs.SFA.size()), device_); + sfb_ = cuda_bytes(static_cast(fixture.inputs.SFB.size()), device_); + copy_to_device(a_, fixture.inputs.A, "A"); + copy_to_device(b_, fixture.inputs.B, "B"); + copy_to_device(sfa_, fixture.inputs.SFA, "SFA"); + copy_to_device(sfb_, fixture.inputs.SFB, "SFB"); + check_cuda(cudaStreamSynchronize(stream_), "cudaStreamSynchronize(input setup)"); + + fast_slot_.M = M; + fast_slot_.N = N; + fast_slot_.K = K; + fast_slot_.dA = a_.data_ptr(); + fast_slot_.dB = b_.data_ptr(); + fast_slot_.dSFA = sfa_.data_ptr(); + fast_slot_.dSFB = sfb_.data_ptr(); + fast_slot_.a_sz = fixture.inputs.A.size(); + fast_slot_.b_sz = fixture.inputs.B.size(); + fast_slot_.sfa_sz = fixture.inputs.SFA.size(); + fast_slot_.sfb_sz = fixture.inputs.SFB.size(); + fast_slot_.A_t = host::make_ab_tmap( + fast_slot_.dA, M, K, fixture.inputs.ab_row_stride); + fast_slot_.B_t = host::make_ab_tmap( + fast_slot_.dB, N, K, fixture.inputs.ab_row_stride); + fast_slot_.SFA_t = host::make_sf_tmap(fast_slot_.dSFA, M, K); + fast_slot_.SFB_t = host::make_sf_tmap(fast_slot_.dSFB, N, K); + + sms_ = properties.multiProcessorCount; + clusters_ = sms_ / 2; + grid_ = nvfp4::launch_grid(clusters_); + if (prepare_fast) { + prepare_fast_kernel(properties, fixture.inputs.A.size() + + fixture.inputs.B.size() + fixture.inputs.SFA.size() + + fixture.inputs.SFB.size()); + } else { + plan_ = std::make_unique(M, N, K); + initialize_vendor_descriptor(); + plan_->initialize_algorithm(vendor_); + if (plan_->workspace_bytes > 0) { + workspace_ = cuda_bytes( + static_cast(plan_->workspace_bytes), device_); + vendor_.workspace = workspace_.data_ptr(); + TORCH_CHECK( + (reinterpret_cast(vendor_.workspace) & 255u) == 0, + "cuBLASLt workspace is not 256-byte aligned"); + } + } + } + + ~Nvfp4Context() { + if (vendor_.op != nullptr) { + cublasLtMatmulDescDestroy(vendor_.op); + vendor_.op = nullptr; + } + } + + Nvfp4Context(const Nvfp4Context&) = delete; + Nvfp4Context& operator=(const Nvfp4Context&) = delete; + + void launch_cublaslt(torch::Tensor output) { + validate_output(output); + TORCH_CHECK(plan_ != nullptr, "cuBLASLt was not prepared for this context"); + c10::cuda::CUDAGuard guard(device_); + cudaStream_t stream = current_stream(device_); + float alpha = 1.0f; + float beta = 0.0f; + void* output_ptr = output.data_ptr(); + check_cublas( + cublasLtMatmul( + plan_->handle, + vendor_.op, + &alpha, + vendor_.a, + plan_->a_layout, + vendor_.b, + plan_->b_layout, + &beta, + output_ptr, + plan_->out_layout, + output_ptr, + plan_->out_layout, + &plan_->algo, + vendor_.workspace, + plan_->workspace_bytes, + stream), + "cublasLtMatmul"); + } + + void launch_fast(torch::Tensor output) { + validate_output(output); + TORCH_CHECK(fast_prepared_, "fast.cu kernel was not prepared for this context"); + c10::cuda::CUDAGuard guard(device_); +#if NVFP4_HAS_SCHEDULE + TORCH_CHECK( + scheduled_device == device_ && scheduled_shape == std::array{M_, N_, K_}, + "fast.cu r9 schedule no longer matches this context"); +#endif + cudaStream_t stream = current_stream(device_); + nvfp4::nvfp4_gemm_kernel + <<>>( + fast_slot_.A_t, + fast_slot_.B_t, + fast_slot_.SFA_t, + fast_slot_.SFB_t, + reinterpret_cast<__half*>(output.data_ptr()), + M_, + N_, + K_ +#if NVFP4_HAS_SCHEDULE + , bench::MAIN_TABLE * nvfp4::L2A_ROUTE_WORK_CAP +#endif + ); + check_cuda(cudaGetLastError(), "fast.cu NVFP4 kernel launch"); + } + + void validate_schedule() { +#if NVFP4_HAS_SCHEDULE + TORCH_CHECK(fast_prepared_, "fast.cu r9 schedule was not prepared"); + c10::cuda::CUDAGuard guard(device_); + check_cuda( + cudaStreamSynchronize(current_stream(device_)), + "cudaStreamSynchronize(schedule audit)"); + unsigned placement_errors = 0; + check_cuda( + cudaMemcpyFromSymbol( + &placement_errors, + nvfp4::l2a_placement_errors, + sizeof(placement_errors)), + "cudaMemcpyFromSymbol(placement audit)"); + TORCH_CHECK( + placement_errors == 0, + "fast.cu r9 placement audit found ", + placement_errors, + " cluster/L2 ownership mismatches; result is invalid"); +#else + TORCH_CHECK(fast_prepared_, "fast.cu kernel was not prepared"); +#endif + } + + torch::Tensor a_packed() const { return a_; } + torch::Tensor b_packed() const { return b_; } + torch::Tensor sfa_packed() const { return sfa_; } + torch::Tensor sfb_packed() const { return sfb_; } + int rung() const { return FAST_CU_NVFP4_RUNG; } + uint64_t seed() const { return seed_; } + +private: + void copy_to_device( + torch::Tensor& destination, + const std::vector& source, + const char* label) { + TORCH_CHECK( + static_cast(destination.numel()) == source.size(), + label, + " source/destination size mismatch"); + check_cuda( + cudaMemcpyAsync( + destination.data_ptr(), + source.data(), + source.size(), + cudaMemcpyHostToDevice, + stream_), + label); + } + + void initialize_vendor_descriptor() { + vendor_.a = a_.data_ptr(); + vendor_.b = b_.data_ptr(); + vendor_.sfa = sfa_.data_ptr(); + vendor_.sfb = sfb_.data_ptr(); + check_cublas( + cublasLtMatmulDescCreate( + &vendor_.op, + CUBLAS_COMPUTE_32F, + CUDA_R_32F), + "cublasLtMatmulDescCreate"); + cublasOperation_t transa = CUBLAS_OP_T; + cublasOperation_t transb = CUBLAS_OP_N; + int32_t sf_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; + int8_t fast_accum = 0; + check_cublas( + cublasLtMatmulDescSetAttribute( + vendor_.op, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa)), + "set cuBLASLt TRANSA"); + check_cublas( + cublasLtMatmulDescSetAttribute( + vendor_.op, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transb)), + "set cuBLASLt TRANSB"); + check_cublas( + cublasLtMatmulDescSetAttribute( + vendor_.op, + CUBLASLT_MATMUL_DESC_FAST_ACCUM, + &fast_accum, + sizeof(fast_accum)), + "set cuBLASLt FAST_ACCUM"); + check_cublas( + cublasLtMatmulDescSetAttribute( + vendor_.op, + CUBLASLT_MATMUL_DESC_A_SCALE_MODE, + &sf_mode, + sizeof(sf_mode)), + "set cuBLASLt A scale mode"); + check_cublas( + cublasLtMatmulDescSetAttribute( + vendor_.op, + CUBLASLT_MATMUL_DESC_B_SCALE_MODE, + &sf_mode, + sizeof(sf_mode)), + "set cuBLASLt B scale mode"); + void* sfa_pointer = sfa_.data_ptr(); + void* sfb_pointer = sfb_.data_ptr(); + check_cublas( + cublasLtMatmulDescSetAttribute( + vendor_.op, + CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, + &sfa_pointer, + sizeof(sfa_pointer)), + "set cuBLASLt A scale pointer"); + check_cublas( + cublasLtMatmulDescSetAttribute( + vendor_.op, + CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, + &sfb_pointer, + sizeof(sfb_pointer)), + "set cuBLASLt B scale pointer"); + } + + void prepare_fast_kernel( + const cudaDeviceProp& properties, + size_t operand_read_bytes) { + bench::g.sms = sms_; + bench::g.clusters = clusters_; + bench::g.l2_bytes = properties.l2CacheSize; + bench::g.grid = grid_; + check_cuda( + cudaFuncSetAttribute( + nvfp4::nvfp4_gemm_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(sizeof(nvfp4::SmemCD))), + "cudaFuncSetAttribute(fast.cu kernel)"); +#if NVFP4_HAS_SCHEDULE + check_cuda( + cudaFuncSetAttribute( + nvfp4::l2a_cluster_probe, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(sizeof(nvfp4::SmemCD))), + "cudaFuncSetAttribute(fast.cu census)"); + cudaLaunchConfig_t occupancy_config{}; + cudaLaunchAttribute cluster_attribute{}; + cluster_attribute.id = cudaLaunchAttributeClusterDimension; + cluster_attribute.val.clusterDim.x = 2; + cluster_attribute.val.clusterDim.y = 1; + cluster_attribute.val.clusterDim.z = 1; + occupancy_config.blockDim = dim3(nvfp4::TB_SIZE, 1, 1); + occupancy_config.dynamicSmemBytes = sizeof(nvfp4::SmemCD); + occupancy_config.attrs = &cluster_attribute; + occupancy_config.numAttrs = 1; + int active_clusters = 0; + check_cuda( + cudaOccupancyMaxActiveClusters( + &active_clusters, + nvfp4::nvfp4_gemm_kernel, + &occupancy_config), + "cudaOccupancyMaxActiveClusters(fast.cu r9)"); + TORCH_CHECK( + active_clusters == clusters_, + "fast.cu r9 requires all ", + clusters_, + " two-CTA clusters resident; occupancy reports ", + active_clusters); + strict_setup_r9_schedule( + device_, properties, grid_, clusters_, M_, N_, K_, operand_read_bytes); +#endif + { + std::lock_guard lock(small_gate_mutex); + const bool already_gated = std::find( + small_gated_devices.begin(), small_gated_devices.end(), device_) + != small_gated_devices.end(); + if (!already_gated) { + TORCH_CHECK( + bench::small_correctness_gates(), + "fast.cu upstream small host-reference, guard, or determinism gate failed"); + small_gated_devices.push_back(device_); + } + } + fast_prepared_ = true; + } + + void validate_output(const torch::Tensor& output) const { + TORCH_CHECK(output.is_cuda(), "NVFP4 output must be a CUDA tensor"); + TORCH_CHECK( + output.get_device() == device_, + "NVFP4 output is on CUDA device ", + output.get_device(), + " but context is on ", + device_); + TORCH_CHECK(output.scalar_type() == torch::kFloat16, "NVFP4 output must be FP16"); + TORCH_CHECK(output.is_contiguous(), "NVFP4 output must be contiguous"); + TORCH_CHECK( + output.dim() == 2 && output.size(0) == M_ && output.size(1) == N_, + "NVFP4 output must have shape [", M_, ", ", N_, "]"); + } + + int M_ = 0; + int N_ = 0; + int K_ = 0; + uint64_t seed_ = 0; + int device_ = -1; + int sms_ = 0; + int clusters_ = 0; + dim3 grid_{}; + cudaStream_t stream_ = nullptr; + bool fast_prepared_ = false; + std::unique_ptr plan_; + bench::LtSlot vendor_{}; + host::DeviceSlot fast_slot_{}; + torch::Tensor a_; + torch::Tensor b_; + torch::Tensor sfa_; + torch::Tensor sfb_; + torch::Tensor workspace_; +}; + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + py::class_(module, "Nvfp4Context") + .def(py::init()) + .def("launch_cublaslt", &Nvfp4Context::launch_cublaslt) + .def("launch_fast", &Nvfp4Context::launch_fast) + .def("validate_schedule", &Nvfp4Context::validate_schedule) + .def_property_readonly("a_packed", &Nvfp4Context::a_packed) + .def_property_readonly("b_packed", &Nvfp4Context::b_packed) + .def_property_readonly("sfa_packed", &Nvfp4Context::sfa_packed) + .def_property_readonly("sfb_packed", &Nvfp4Context::sfb_packed) + .def_property_readonly("rung", &Nvfp4Context::rung) + .def_property_readonly("seed", &Nvfp4Context::seed); +} diff --git a/code/labs/fast_cu/nvfp4_performance_intake.yaml b/code/labs/fast_cu/nvfp4_performance_intake.yaml new file mode 100644 index 000000000..e8003a656 --- /dev/null +++ b/code/labs/fast_cu/nvfp4_performance_intake.yaml @@ -0,0 +1,99 @@ +objective: + primary_kpi: median_launch_latency_us + secondary_kpis: [p95_launch_latency_us, effective_pflops] + +benchmark: + layer: component + grade: realism + variable_under_test: kernel_implementation + +workload: + type: gemm + model: "8192x8192x8192 NVFP4 E2M1 GEMM, UE4M3 VEC16 scales, FP16 output" + batch_shape: + p50: [8192, 8192, 8192] + p95: [8192, 8192, 8192] + max: [8192, 8192, 8192] + constraints: + - identical packed input bytes and scale buffers for both arms + - full FP16 output compared with rtol 0.002 and atol 0.5 + - exact SM103 GB300/B300 and CUDA 13.1 or newer + - no B200 or alternate-precision substitution + +slos: + latency_ms: {p50: null, p95: null, p99: null} + availability: null + cost_budget: null + +environment: + hardware: + gpu: ["GB300/B300 SM103", 1, null] + cpu: "recorded by harness provenance" + ram_gb: null + storage: null + interconnect: ["single GPU"] + topology: "single exact-SM103 device" + software: + os: "recorded by harness provenance" + kernel: "recorded by harness provenance" + driver: "recorded by harness provenance" + cuda: "13.1 or newer; extension compiled only for sm_103a" + pytorch: "recorded by harness provenance" + triton: "not used" + frameworks: ["CUDA C++", "cuBLASLt", "PyTorch extension"] + nccl: "not used" + container_orchestrator: "recorded by harness provenance" + data_pipeline: + - caller seed is passed to the pinned upstream host quantizer during setup + - packed E2M1 values and UE4M3 VEC16 scales remain device resident + +current_baseline: + p50_ms: null + p99_ms: null + gpu_util_percent: null + mem_bw_percent: null + goodput_percent: null + +hot_path_model: + path: "labs.fast_cu.baseline_nvfp4_gemm vs labs.fast_cu.optimized_nvfp4_gemm" + phase: steady_state + invocation_frequency: per_request + invocations_per_primary_unit: 1 + estimated_costs: + bytes_read: 75497472 + bytes_written: 134217728 + host_device_bytes: 0 + allocations_or_materializations: 0 + kernel_launches: 1 + synchronizations_or_locks: 0 + api_process_or_rpc_crossings: 1 + storage_or_network_operations: 0 + operation_cost_source: "Analytical counts from the pinned 8192^3 packed layouts; runtime cost unmeasured" + expected_dominant_cost: "NVFP4 tensor-core GEMM plus FP16 output traffic" + best_case_primary_kpi_improvement_pct: null + +evidence: + baseline_artifact: "pending exact-SM103 harness run" + profile_artifacts: [] + allocation_or_memory_profile: "setup-only allocations enforced by adapter structure" + hardware_counter_artifacts: [] + static_review_findings: + - "r9 depends on a runtime L2-side census and owned tile schedule" + - "adapter rejects census/table failure instead of accepting upstream raster fallback" + - "both cuBLASLt and fast.cu launch on the active PyTorch CUDA stream" + +candidate_optimizations: + - "r0-r9 pinned fast.cu ladder, explicitly selected at extension build time" + - "r9 L2-side ownership and per-shape tile ordering" + +known_issues_or_hypotheses: + - "No exact-SM103 runtime measurement is available in the local development environment" + - "This native pair uses one resident operand set and does not reproduce the upstream rotating five-round headline protocol" + +risk_guardrails: + correctness_tests: + - "full FP16 cuBLASLt-versus-fast.cu output comparison" + - "same packed A/B/SFA/SFB bytes from the caller seed" + - "post-timing r9 cluster/L2 placement counter audit" + change_windows: + - "promote performance claims only after locked-clock repeated comparison plus NCU and NSYS evidence" diff --git a/code/labs/fast_cu/nvfp4_workload_spec.yaml b/code/labs/fast_cu/nvfp4_workload_spec.yaml new file mode 100644 index 000000000..e9efb6588 --- /dev/null +++ b/code/labs/fast_cu/nvfp4_workload_spec.yaml @@ -0,0 +1,125 @@ +run_metadata: + benchmark_layer: component + benchmark_grade: realism + use_case: engineering_debug + owner: "AI Systems Performance Engineering" + ticket_or_epic: "fast.cu NVFP4 native integration" + +comparison_axis: + variable_under_test: kernel_implementation + control: "cuBLASLt first usable NVFP4 algorithm" + candidate: "pinned fast.cu rung; r9 by default" + constants_locked: + - shape_8192x8192x8192 + - packed_e2m1_inputs + - vec16_ue4m3_scales + - fp32_accumulation_policy + - fp16_full_output + - caller_seed + - current_cuda_stream + +workload: + type: inference + name: "fast.cu GB300 NVFP4 GEMM" + model: + name: "dense matrix multiplication" + parameter_count_b: null + precision: nvfp4_e2m1_vec16_ue4m3_to_fp16 + weights_source: "fast.cu@2dfe5e26aecfd9e5f27bf9d5837deea01acda24b" + input_distribution: + sequence_length_mix: {p50: null, p95: null, max: null} + prompt_tokens: {p50: null, p95: null, max: null} + completion_tokens: {p50: null, p95: null, max: null} + batching_policy: + mode: static + max_batch_size: 1 + max_num_seqs: 1 + max_tokens_in_batch: null + concurrency_model: + mode: closed_loop + values: [1] + serving_contract: + stack_under_test: "native PyTorch CUDA extension" + protocol: "in-process benchmark_fn" + deployment_model: "single exact-SM103 GPU" + fairness_constraints: + protocol_fixed_across_candidates: true + deployment_model_fixed_across_candidates: true + api_behavior_fixed_across_candidates: true + quality_constraints: + correctness_suite: + - "full output torch.allclose rtol=0.002 atol=0.5" + - "finite FP16 output" + - "r9 placement audit equals zero" + max_accuracy_delta_pct: 0.0 + allowed_quality_exceptions: [] + +execution: + image: "record in run manifest" + runtime_versions: + driver: "record in run manifest" + cuda: ">=13.1" + nccl: "not used" + pytorch: "record in run manifest" + runtime_or_server: "pinned native extension" + hardware_tier: + gpu_model: "GB300/B300" + gpu_count: 1 + cpu_model: "record in run manifest" + ram_gb: null + network_fabric: "not used" + storage_tier: "not used during timing" + topology: + gpu_to_gpu: "single GPU" + gpu_to_nic_affinity: "not applicable" + numa_layout: "record in run manifest" + pcie_layout: "record in run manifest" + scheduler_path: + platform: bare_metal + queue_or_scheduler: "repository benchmark queue" + placement_policy: "exclusive exact-SM103 GPU" + control_plane_path: "aisp bench run" + +measurement_policy: + trials: 5 + warmup_trials: 5 + report_statistics: [median, mean, stddev, p95, p99, cv_pct] + outlier_policy: + method: tukey_iqr_and_mad + keep_raw_samples: true + auto_flag_only: true + discard_requires_written_reason: true + confidence_policy: + method: bootstrap_95_ci + min_effect_size_pct: 5.0 + environment_controls: + gpu_clocks_locked: true + app_clocks_recorded: true + topology_recorded: true + background_load_recorded: true + +metrics: + primary: median_launch_latency_us + secondary: [p95_launch_latency_us, effective_pflops] + +diagnosis_policy: + bottleneck_classes: [compute_bound, memory_bound, launch_bound] + required_signals: [gpu_counters, nsys_trace, correctness_receipt, placement_audit] + distributed: + rank_level_visibility: false + isolate_suspected_straggler_node: false + cordon_on_repeated_failure: false + +artifacts: + artifact_root: "artifacts/runs/" + raw_artifacts_immutable: true + require_manifest: true + require_profiler_artifacts: true + require_sha256: true + provenance_signing: + required: true + status: gap_in_current_repo + publication_packet: + include_methodology_writeup: true + include_comparison_report: true + include_audit_ledger: true diff --git a/code/labs/fast_cu/optimization_guide.md b/code/labs/fast_cu/optimization_guide.md new file mode 100644 index 000000000..98e7aa070 --- /dev/null +++ b/code/labs/fast_cu/optimization_guide.md @@ -0,0 +1,81 @@ +# Where fast.cu fits in this repository + +The source is [fast.cu](https://github.com/pranjalssh/fast.cu) and its accompanying +[H100 worklog](https://cudaforfun.substack.com/p/outperforming-cublas-on-h100-a-worklog) +and [NVFP4 article](https://cudaforfun.substack.com/p/outperforming-cublas-on-nvfp4). +The supplied PDF of the NVFP4 article was consulted, including its optimization +ladder figure and measurement protocol. The PDF itself is not redistributed. + +## Existing coverage and the added experiments + +| Subject | Existing home | What this lab adds | +| --- | --- | --- | +| Arithmetic intensity, tiling, fusion | [ch09](../../ch09/README.md), including `tcgen05_tma_pipeline` and `micro_tiling_matmul` | Complete Hopper kernel progression and a cuBLAS comparison with the original kernel 12 | +| Tensor-core pipelines, TMA, warp roles, clusters | [ch10](../../ch10/README.md), [blackwell_matmul](../blackwell_matmul/README.md) | Complete plain CUDA/PTX GB300 r0-r9 ladder with separate A/B and scale queues and two-CTA MMA | +| Warp reductions and bandwidth | [ch10 atomic/DSMEM reduction](../../ch10/README.md), [memory_bandwidth_patterns](../memory_bandwidth_patterns/README.md) | CUB versus a safe adaptation of fast.cu's vectorized int32 reduction, including B200 execution | +| Stream ordering and overlap | [ch11](../../ch11/README.md) | Explicit current-stream native adapters and stream-ordered reduction reset | +| Compiler-sensitive instruction selection | [ch14](../../ch14/README.md) | `%laneid`, exact block geometry, wide stores, and an isolated B200 epilogue experiment | +| B200 block-scaled FP4 GEMM | [nvfp4_gemm](../nvfp4_gemm/README.md) | Distinct GB300 SM103/K96 implementation; it complements the existing SM100 path | + +Existing chapter implementations remain the teaching entrypoints for common +concepts. The new lab keeps the source progression and controlled comparisons +together so that an unfamiliar reader can inspect the complete examples locally. + +## Every upstream example is present + +- [Hopper GEMM driver](upstream/h100/matmul.cu) and + [matmul_1.cuh through matmul_12.cuh](upstream/h100/matmul/) preserve the progression + from elementary matrix multiplication through asynchronous, warp-specialized, + persistent scheduling. The supplementary `cublaslt_matmul.cu` and + `pingpong_experimental.cuh` are also retained as upstream reference experiments. +- [Hopper reduction](upstream/h100/sum.cu) preserves all four source variants and + its CUB comparison. The native adapter uses a corrected variant; the original + program has a grid-wide output-reset race and weak correctness handling. +- [GB300 driver](upstream/gb300/nvfp4/main.cu) and the ten headers below preserve + the entire NVFP4 progression, scheduler controls, and correctness runner. + +## NVFP4 ladder and what transfers to B200 + +NVFP4 stores E2M1 values with an FP8 scale for each group of 16 along K. Values +and scales have different layouts and traffic patterns. The source uses TMA +queues for both, FP32 accumulators in TMEM, and FP16 output. Correctness depends +on descriptor layout, queue lifetime, barrier phase, and output coverage together. + +| Rung | Local source | Mechanism | B200 interpretation | +| --- | --- | --- | --- | +| r0 | [gemm0.cuh](upstream/gb300/nvfp4/gemm0.cuh) | Working two-CTA NVFP4 pipeline with K192 feed | K96 instructions require SM103; use the existing SM100 GEMM for K64 | +| r1 | [gemm1.cuh](upstream/gb300/nvfp4/gemm1.cuh) | Direct `%laneid` for single-lane work | Inspect uniform-register and convergence code generation; avoid assuming a speedup | +| r2 | [gemm2.cuh](upstream/gb300/nvfp4/gemm2.cuh) | K256 TMA windows with independently released buffers | Queue/barrier design transfers; the K96/K256 decomposition does not transfer unchanged | +| r3 | [gemm3.cuh](upstream/gb300/nvfp4/gemm3.cuh) | Overlap two TMEM accumulator buffers | Relates to ch10 ping-pong pipelines; check TMEM capacity and accumulator lifetime | +| r4 | [gemm4.cuh](upstream/gb300/nvfp4/gemm4.cuh) | 256-bit output stores | Implemented as the native `b200_epilogue` store-width pair | +| r5 | [gemm5.cuh](upstream/gb300/nvfp4/gemm5.cuh) | Avoid L1 allocation and favor early L2 eviction for output | Both B200 epilogue arms use the same policy to isolate width; policy itself is not ablated | +| r6 | [gemm6.cuh](upstream/gb300/nvfp4/gemm6.cuh) | Exact block shape improves compiler knowledge | Exact launch geometry must remain true; requires suitable compiler support | +| r7 | [gemm7.cuh](upstream/gb300/nvfp4/gemm7.cuh) | Remove dead tail work | Transfer only after proving coverage for the actual tile/K decomposition | +| r8 | [gemm8.cuh](upstream/gb300/nvfp4/gemm8.cuh) | Fold compatible tails into K64 MMAs | K64 is supported on SM100; the mixed K96/K64 schedule remains SM103-specific | +| r9 | [gemm9.cuh](upstream/gb300/nvfp4/gemm9.cuh) | L2-side ownership and shape-dependent visit order | Measure actual topology; do not transplant a GB300 route/census assumption onto B200 | + +The r9 scheduling lesson has two parts: decide which cluster population reuses a +row, then place those reuses close together in visitation order. A tiled or Hilbert +order alone does not establish ownership. Topology measurements and route coverage +checks are prerequisites; a silent raster fallback cannot count as r9 evidence. + +## Lessons carried into the adapters + +1. Preserve the entire arithmetic contract. Storage layout changes must not change + the logical matrix or hide output elements from verification. +2. Separate buffer preparation from execution and expose the real output written + by the timed kernel. Hashes verify source identity, not numerical correctness. +3. Treat CUDA architecture suffixes as requirements. SM100 and SM103 are not + interchangeable simply because both are Blackwell. +4. Use a controlled ablation for a portable idea. The B200 epilogue holds loads, + conversion, cache policy, launch geometry, and output size constant while + changing only store width. +5. Keep input distribution and integer range explicit. Repairing an atomic reset + race is insufficient if the reduction can still overflow. +6. Distinguish a source author's measurements from this repository's measurements. + Buffer rotation, order, cooldown, clocks, power, and compiler versions can + change a small performance margin. Record those choices before comparing. + +No expected speedup is installed from the blog. A performance conclusion requires +equivalent workloads, full-output correctness, repeated paired measurements, +clock/provenance records, and Nsight evidence from the target hardware. diff --git a/code/labs/fast_cu/optimized_b200_epilogue.py b/code/labs/fast_cu/optimized_b200_epilogue.py new file mode 100644 index 000000000..310ebb7a5 --- /dev/null +++ b/code/labs/fast_cu/optimized_b200_epilogue.py @@ -0,0 +1,13 @@ +"""SM100 candidate: FP32-to-FP16 epilogue with one 256-bit store.""" + +from core.harness.benchmark_harness import BaseBenchmark +from labs.fast_cu.b200_epilogue import B200EpilogueBenchmark + + +class OptimizedB200EpilogueBenchmark(B200EpilogueBenchmark): + def __init__(self): + super().__init__(optimized=True) + + +def get_benchmark() -> BaseBenchmark: + return OptimizedB200EpilogueBenchmark() diff --git a/code/labs/fast_cu/optimized_h100_bf16_gemm.py b/code/labs/fast_cu/optimized_h100_bf16_gemm.py new file mode 100644 index 000000000..638b53ea1 --- /dev/null +++ b/code/labs/fast_cu/optimized_h100_bf16_gemm.py @@ -0,0 +1,44 @@ +"""Optimized: pranjalssh/fast.cu's latest H100 BF16 WGMMA kernel.""" + +from __future__ import annotations + +import torch + +from core.harness.benchmark_harness import BaseBenchmark +from labs.fast_cu.h100_common import H100Bf16GemmBenchmarkBase + + +class OptimizedH100Bf16GemmBenchmark(H100Bf16GemmBenchmarkBase): + """Run upstream matmul_12 with repository timing and verification semantics.""" + + nvtx_label = "fast_cu_optimized_h100_bf16_gemm" + + def setup(self) -> None: + self._setup_tensors() + if ( + self.extension is None + or self.matrix_a is None + or self.matrix_b is None + or self._physical_output is None + ): + raise RuntimeError("H100 GEMM setup did not initialize required state") + # TMA descriptors, the Hilbert schedule, device scratch, and the dynamic + # shared-memory attribute are prepared before the steady-state timer. + self.extension.prepare_upstream_gemm(self.matrix_a, self.matrix_b, self._physical_output) + + def benchmark_fn(self) -> None: + if ( + self.extension is None + or self.matrix_a is None + or self.matrix_b is None + or self._physical_output is None + or self._logical_output is None + ): + raise RuntimeError("setup() must run before benchmark_fn()") + with torch.inference_mode(), self._nvtx_range(self.nvtx_label): + self.extension.run_upstream_gemm(self.matrix_a, self.matrix_b, self._physical_output) + self.output = self._logical_output + + +def get_benchmark() -> BaseBenchmark: + return OptimizedH100Bf16GemmBenchmark() diff --git a/code/labs/fast_cu/optimized_int32_reduction.py b/code/labs/fast_cu/optimized_int32_reduction.py new file mode 100644 index 000000000..3f345ac5d --- /dev/null +++ b/code/labs/fast_cu/optimized_int32_reduction.py @@ -0,0 +1,28 @@ +"""Optimized: safe adaptation of fast.cu's vectorized int32 reduction.""" + +from __future__ import annotations + +import torch + +from core.harness.benchmark_harness import BaseBenchmark +from labs.fast_cu.h100_common import Int32ReductionBenchmarkBase + + +class OptimizedInt32ReductionBenchmark(Int32ReductionBenchmarkBase): + """Run the upstream-derived int4 reduction without its cross-CTA reset race.""" + + nvtx_label = "fast_cu_optimized_int32_reduction" + + def setup(self) -> None: + self._setup_tensors() + + def benchmark_fn(self) -> None: + if self.extension is None or self.input is None or self._output_buffer is None: + raise RuntimeError("setup() must run before benchmark_fn()") + with torch.inference_mode(), self._nvtx_range(self.nvtx_label): + self.extension.run_safe_vectorized_reduction(self.input, self._output_buffer) + self.output = self._output_buffer + + +def get_benchmark() -> BaseBenchmark: + return OptimizedInt32ReductionBenchmark() diff --git a/code/labs/fast_cu/optimized_nvfp4_gemm.py b/code/labs/fast_cu/optimized_nvfp4_gemm.py new file mode 100644 index 000000000..6771cc692 --- /dev/null +++ b/code/labs/fast_cu/optimized_nvfp4_gemm.py @@ -0,0 +1,24 @@ +"""Selected fast.cu rung for the pinned GB300 NVFP4 workload.""" + +from __future__ import annotations + +from core.harness.benchmark_harness import BaseBenchmark +from labs.fast_cu.nvfp4 import DEFAULT_RUNG, FastCuNvfp4Benchmark, Nvfp4Workload + + +class FastCuNvfp4KernelBenchmark(FastCuNvfp4Benchmark): + """Native fast.cu r0-r9 kernel; r9 is the explicit default.""" + + def __init__( + self, + workload: Nvfp4Workload | None = None, + rung: int = DEFAULT_RUNG, + ) -> None: + super().__init__(optimized=True, workload=workload, rung=rung) + + +def get_benchmark() -> BaseBenchmark: + return FastCuNvfp4KernelBenchmark(rung=DEFAULT_RUNG) + + +__all__ = ["FastCuNvfp4KernelBenchmark", "get_benchmark"] diff --git a/code/labs/fast_cu/performance_intake.yaml b/code/labs/fast_cu/performance_intake.yaml new file mode 100644 index 000000000..074c4efa0 --- /dev/null +++ b/code/labs/fast_cu/performance_intake.yaml @@ -0,0 +1,49 @@ +objective: + primary_kpi: latency_ms + secondary_kpis: [effective_gbps, gemm_tflops] +benchmark: + layer: component + grade: realism + variable_under_test: kernel_implementation_or_output_store_width +workload: + h100_bf16_gemm: {m: 8192, n: 8192, k: 8192, input: bf16, output: bf16} + int32_reduction: {elements: 268435456, input: int32, output: int32, input_values: [-1, 0, 1]} + b200_epilogue: {elements: 16777216, input: fp32, output: fp16, values_per_thread: 16} + nvfp4_gemm: {m: 8192, n: 8192, k: 8192, input: nvfp4, output: fp16, default_rung: 9} + constraints: [same_inputs_per_pair, full_output_verification, current_cuda_stream, no_architecture_fallback] +hot_path_model: + phase: steady_state + invocation_frequency: per_batch + invocations_per_primary_unit: 1 + estimated_costs: + h100_gemm_flops: 1099511627776 + h100_gemm_minimum_bytes: 402653184 + reduction_input_bytes: 1073741824 + reduction_output_bytes: 4 + epilogue_input_bytes: 67108864 + epilogue_output_bytes: 33554432 + nvfp4_gemm_flops: 1099511627776 + nvfp4_operand_and_scale_bytes: 75497472 + nvfp4_output_bytes: 134217728 + host_device_bytes: 0 + storage_or_network_operations: 0 + operation_cost_source: tensor shapes and scalar formats; not measured hardware costs + expected_dominant_cost: tensor-core feed and epilogue for GEMM; input bandwidth and atomics for reduction; load/convert/store instruction and memory throughput for epilogue + best_case_primary_kpi_improvement_pct: null +evidence: + baseline_artifact: null + profile_artifacts: [] + hardware_counter_artifacts: [] + static_review_findings: + - Exact SM103 K96 instructions do not execute on SM100. + - Vectorized reduction requires output reset before the grid and bounded sums. + - Wide-store epilogue isolates 256-bit versus 128-bit stores with identical cache hints. + - NVFP4 single-buffer measurements differ from upstream rotating-buffer headline results. +candidate_optimizations: + - H100 kernel12 versus cuBLAS at equal BF16 precision and shape. + - Vectorized warp sum versus CUB with identical int32 inputs. + - B200 output store width with fixed conversions, launch geometry, and cache policy. + - GB300 r9 versus cuBLASLt using identical packed NVFP4 inputs and FP16 outputs. +risk_guardrails: + correctness_tests: [tests/test_fast_cu_source.py, tests/test_fast_cu_build.py, tests/test_fast_cu_h100.py, tests/test_fast_cu_b200_epilogue.py, tests/test_fast_cu_nvfp4.py] + publication_boundary: CPU/source checks and noncanonical GPU correctness tests do not establish a performance win. diff --git a/code/labs/fast_cu/source.py b/code/labs/fast_cu/source.py new file mode 100644 index 000000000..cf6ff9229 --- /dev/null +++ b/code/labs/fast_cu/source.py @@ -0,0 +1,54 @@ +"""Verify the complete vendored source snapshot before compiling CUDA code.""" + +from __future__ import annotations + +import hashlib +import json +import re +from pathlib import Path, PurePosixPath + +LAB_DIR = Path(__file__).resolve().parent +UPSTREAM_DIR = LAB_DIR / "upstream" +MANIFEST_PATH = LAB_DIR / "upstream_manifest.json" +UPSTREAM_REVISION = "2dfe5e26aecfd9e5f27bf9d5837deea01acda24b" +UPSTREAM_REPOSITORY = "https://github.com/pranjalssh/fast.cu" + + +def verify_upstream(root: Path = UPSTREAM_DIR, manifest_path: Path = MANIFEST_PATH) -> dict: + """Fail on missing, changed, extra, or symlinked source files. + + The explicit paths support offline integrity checks on copied source trees. + No download, build, GPU initialization, or benchmark happens here. + """ + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + if manifest.get("schema_version") != 1: + raise ValueError("Unsupported fast.cu source manifest schema") + if manifest.get("revision") != UPSTREAM_REVISION: + raise ValueError("Unexpected fast.cu source revision") + if manifest.get("repository", "").removesuffix(".git") != UPSTREAM_REPOSITORY: + raise ValueError("Unexpected fast.cu source repository") + files = manifest.get("files") + if not isinstance(files, dict) or not files or "LICENSE" not in files: + raise ValueError("fast.cu manifest must include source files and LICENSE") + if root.is_symlink() or not root.is_dir(): + raise ValueError(f"fast.cu source root must be a real directory: {root}") + actual = set() + for path in root.rglob("*"): + if path.is_symlink(): + raise ValueError(f"Symlink in fast.cu source tree: {path.relative_to(root)}") + if path.is_file(): + actual.add(path.relative_to(root).as_posix()) + if actual != set(files): + missing = sorted(set(files) - actual) + extra = sorted(actual - set(files)) + raise ValueError(f"fast.cu source file set changed: missing={missing}, extra={extra}") + for relative, expected in files.items(): + path = PurePosixPath(relative) + if path.is_absolute() or ".." in path.parts or path.as_posix() != relative: + raise ValueError(f"Invalid fast.cu manifest path: {relative}") + if not isinstance(expected, str) or not re.fullmatch(r"[0-9a-f]{64}", expected): + raise ValueError(f"Invalid fast.cu SHA256 for {relative}") + digest = hashlib.sha256((root / relative).read_bytes()).hexdigest() + if digest != expected: + raise ValueError(f"fast.cu source integrity mismatch: {relative}") + return manifest diff --git a/code/labs/fast_cu/upstream/.gitignore b/code/labs/fast_cu/upstream/.gitignore new file mode 100644 index 000000000..89f9ac04a --- /dev/null +++ b/code/labs/fast_cu/upstream/.gitignore @@ -0,0 +1 @@ +out/ diff --git a/code/labs/fast_cu/upstream/LICENSE b/code/labs/fast_cu/upstream/LICENSE new file mode 100644 index 000000000..9a4f580a2 --- /dev/null +++ b/code/labs/fast_cu/upstream/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Pranjal Shankhdhar + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/code/labs/fast_cu/upstream/Makefile b/code/labs/fast_cu/upstream/Makefile new file mode 100644 index 000000000..af47a48d1 --- /dev/null +++ b/code/labs/fast_cu/upstream/Makefile @@ -0,0 +1,56 @@ +NVCC_FLAGS = -std=c++17 -O3 -DNDEBUG -w +NVCC_LDFLAGS = -lcublas -lcuda +OUT_DIR = out + +CUDA_OUTPUT_FILE = -o $(OUT_DIR)/$@ +NCU_PATH := $(shell which ncu) +NCU_COMMAND = sudo $(NCU_PATH) --set full --import-source yes + +NVCC_FLAGS += --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -Xcompiler=-fPIE -Xcompiler=-Wno-psabi -Xcompiler=-fno-strict-aliasing +NVCC_FLAGS += -gencode arch=compute_90a,code=sm_90a + +NVCC_BASE = nvcc $(NVCC_FLAGS) $(NVCC_LDFLAGS) -lineinfo + +NVFP4_NVCC ?= /usr/local/cuda/bin/nvcc +NVFP4_NVCC_FLAGS = -std=c++17 -O3 -DNDEBUG \ + -gencode arch=compute_103a,code=sm_103a +NVFP4_LDFLAGS = -lcublasLt -lcuda +NVFP4_DIR = gb300/nvfp4 +NVFP4_SOURCE = $(NVFP4_DIR)/main.cu +NVFP4_RUNGS = 0 1 2 3 4 5 6 7 8 9 +NVFP4_TARGETS = $(addprefix nvfp4-r,$(NVFP4_RUNGS)) + +sum: h100/sum.cu + $(NVCC_BASE) $^ $(CUDA_OUTPUT_FILE) + +sumprofile: sum + $(NCU_COMMAND) -o $@ -f $(OUT_DIR)/$^ + +matmul: h100/matmul.cu + mkdir -p $(OUT_DIR) + $(NVCC_BASE) $^ $(CUDA_OUTPUT_FILE) + +matmulprofile: matmul + $(NCU_COMMAND) -o $@ -f $(OUT_DIR)/$^ + +nvfp4: $(OUT_DIR)/nvfp4 + +$(OUT_DIR)/nvfp4: $(NVFP4_SOURCE) $(NVFP4_DIR)/gemm9.cuh | $(OUT_DIR) + $(NVFP4_NVCC) $(NVFP4_NVCC_FLAGS) $< -o $@ $(NVFP4_LDFLAGS) + +$(NVFP4_TARGETS): nvfp4-r%: $(OUT_DIR)/nvfp4-r% + +$(OUT_DIR)/nvfp4-r%: $(NVFP4_SOURCE) $(NVFP4_DIR)/gemm%.cuh | $(OUT_DIR) + $(NVFP4_NVCC) $(NVFP4_NVCC_FLAGS) \ + -DNVFP4_GEMM_HEADER='"gemm$*.cuh"' \ + $< -o $@ $(NVFP4_LDFLAGS) + +nvfp4-ladder: $(NVFP4_TARGETS) + +$(OUT_DIR): + mkdir -p $(OUT_DIR) + +clean: + rm $(OUT_DIR)/* + +.PHONY: nvfp4 nvfp4-ladder $(NVFP4_TARGETS) diff --git a/code/labs/fast_cu/upstream/README.md b/code/labs/fast_cu/upstream/README.md new file mode 100644 index 000000000..67f547a9a --- /dev/null +++ b/code/labs/fast_cu/upstream/README.md @@ -0,0 +1,58 @@ +# Fastest GPU kernels, written from scratch. + +## H100 BF16 Matrix Multiplication + +Matrix multiplication of square bf16 matrices, accumulated in fp32. + +``` +N=4096 +Kernel: 763 TFLOPs +cuBLAS: 716 TFLOPs + +N=8192 +Kernel: 808 TFLOPs +cuBLAS: 795 TFLOPs +``` + +Explanation in https://cudaforfun.substack.com/p/outperforming-cublas-on-h100-a-worklog + +##### To run: +``` +make matmul && out/matmul +``` +Example kernels are in [`h100/matmul/`](h100/matmul), and orchestration is +in [`h100/matmul.cu`](h100/matmul.cu). + +## GB300 NVFP4 Matrix Multiplication + +The NVFP4 example contains ten complete kernel snapshots, a shared correctness +and benchmark runner, and a cuBLASLt baseline. + +```bash +make nvfp4 +./out/nvfp4 8192 8192 8192 --rounds 5 --cooldown-sec 8 + +make nvfp4-ladder +./out/nvfp4-r0 8192 8192 8192 --rounds 5 --cooldown-sec 8 +./out/nvfp4-r9 8192 8192 8192 --rounds 5 --cooldown-sec 8 +``` + +This requires a GB300/B300 and CUDA 13.1. See +[`gb300/nvfp4/README.md`](gb300/nvfp4/README.md) for every optimization +rung, scheduler controls, correctness-only runs, and benchmark modes. + +## H100 Sum Reduction + +We compute sum of 2^30 elements. + +##### To run: +``` +make sum && out/sum +``` + +``` +Kernel: 3240.11 GB/s +cub Library: 3193 GB/s +``` + +The kernel is in [`h100/sum.cu`](h100/sum.cu). diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/README.md b/code/labs/fast_cu/upstream/gb300/nvfp4/README.md new file mode 100644 index 000000000..2327a52ac --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/README.md @@ -0,0 +1,145 @@ +# NVFP4 GEMM on GB300 + +This directory contains the complete optimization ladder for an NVFP4 GEMM +with FP16 output. The kernels are plain CUDA C++ plus inline PTX and target +NVIDIA GB300/B300 (`sm_103a`). cuBLASLt is used only as the reference and +benchmark baseline. + +## Build and run + +From the repository root, build the final kernel: + +```bash +make nvfp4 +./out/nvfp4 8192 8192 8192 --rounds 5 --cooldown-sec 8 +``` + +Every invocation checks correctness and determinism before benchmarking. The +final result includes our PFLOP/s, cuBLASLt PFLOP/s, their ratio, and paired +rounds won. + +Build every optimization rung: + +```bash +make nvfp4-ladder +``` + +Run one rung, or run the whole ladder: + +```bash +./out/nvfp4-r0 8192 8192 8192 --rounds 5 --cooldown-sec 8 + +for rung in {0..9}; do + ./out/nvfp4-r${rung} 8192 8192 8192 \ + --rounds 5 --cooldown-sec 8 +done +``` + +Use `--check-only` to stop after the correctness gates: + +```bash +./out/nvfp4-r9 8192 8192 8192 --check-only +``` + +For a quick kernel-only timing signal, keep the correctness gates but skip the +timed cuBLASLt leg: + +```bash +./out/nvfp4-r9 8192 8192 8192 \ + --ours-only --rounds 5 --cooldown-sec 8 +``` + +## Optimization ladder + +Each header is a complete, directly editable implementation snapshot. + +| Rung | Change | +| --- | --- | +| `r0` | Simple Blackwell baseline: K=192 feed and `threadIdx.x % 32` lane checks | +| `r1` | Read `%laneid` directly for single-lane work | +| `r2` | Fill K=256 A/B windows and pipeline their releases | +| `r3` | Overlap two TMEM accumulator buffers | +| `r4` | Use 256-bit output stores | +| `r5` | Add output cache hints | +| `r6` | Promise the exact 224-thread block shape | +| `r7` | Skip dead tail work | +| `r8` | Fold compatible tails with K64 MMAs | +| `r9` | Add L2-side ownership and per-shape tile ordering | + +To modify a rung, edit its header and rebuild only that target: + +```bash +$EDITOR gb300/nvfp4/gemm3.cuh +make nvfp4-r3 +``` + +## Benchmark modes + +The final kernel supports all four measurement modes used for the published +comparison: + +```bash +gb300/nvfp4/run_bench_modes.sh current +gb300/nvfp4/run_bench_modes.sh clock-pinned +gb300/nvfp4/run_bench_modes.sh sustained +gb300/nvfp4/run_bench_modes.sh triton +``` + +Run all four with: + +```bash +gb300/nvfp4/run_bench_modes.sh all +``` + +The shape defaults to `8192 8192 8192`; pass another shape after the mode. +`clock-pinned` requires passwordless `sudo nvidia-smi` and always resets the +clock through an exit trap. `sustained` takes several minutes. + +The modes are: + +- `current`: rotating inputs, alternating ours/cuBLASLt order, five paired rounds. +- `clock-pinned`: the same protocol with the SM clock fixed at 1305 MHz. +- `sustained`: alternating 60-second blocks and post-warmup tail latency. +- `triton`: hot-loop timing with a 256 MiB L2 clear before each measured launch. + +## Schedulers + +`r9` exposes the scheduling controls used to isolate the final optimization: + +```bash +for schedule in \ + raster pocket-8x4 pocket-8x8 hilbert \ + owned-plain owned-pocket-8x4 owned-pocket-8x8 hilbert-in-owned auto +do + ./out/nvfp4-r9 8192 8192 8192 \ + --schedule "$schedule" --ours-only --rounds 5 --cooldown-sec 8 +done +``` + +Unprefixed schedules change only tile order. `owned-*` schedules also assign A +reuse according to the runtime L2-side census. `auto` selects the final +per-shape policy; for the 8192 cube it uses owned 8x8 pockets. + +## Benchmark contract + +- Every rung and cuBLASLt receives byte-identical packed E2M1 A/B inputs and + byte-identical standard `VEC16_UE4M3` scale buffers. +- Both implementations produce row-major FP16 `[M, N]` output. +- Timed launches rotate enough operand sets to avoid self-warming L2. +- Ours and the first cuBLASLt recommendation alternate first/second order + across paired rounds. +- Throughput uses logical `2MNK` work, including for early rungs that execute + zero-filled tail MMAs. +- `--cooldown-sec 8` waits for the SM clock to return to its reference range; + it aborts rather than recording a result if the clock does not recover. + +## Requirements + +- NVIDIA GB300/B300 +- CUDA 13.1 or a compatible toolkit with the required `tcgen05` forms +- cuBLASLt and the CUDA Driver library +- GNU Make + +The build deliberately uses +`-gencode arch=compute_103a,code=sm_103a`. Plain `-arch=sm_103a` can emit +family-generic PTX that rejects architecture-specific `tcgen05` instructions. diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/gemm0.cuh b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm0.cuh new file mode 100644 index 000000000..6965ebb9e --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm0.cuh @@ -0,0 +1,1148 @@ +#pragma once +// r0: K192 baseline with threadIdx lane checks. +// This is a complete, directly editable ladder snapshot. +// It intentionally duplicates the implementation in the other gemm*.cuh files. + +// Conditional branches are already materialized; edit this implementation directly. +#define NVFP4_BUILD_NAME "r0" +// Read-only metadata for main.cu; these do not select code in this header. +#define NVFP4_HAS_K64_TAIL 0 +#define NVFP4_HAS_SCHEDULE 0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(expr) do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \ + cudaGetErrorName(_e), __FILE__, __LINE__, \ + cudaGetErrorString(_e)); \ + std::abort(); \ + } \ +} while (0) + +#define CU_CHECK(expr) do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + std::fprintf(stderr, "CU error at %s:%d: %s\n", \ + __FILE__, __LINE__, _s ? _s : "?"); \ + std::abort(); \ + } \ +} while (0) + + +// ============================================================================ +// PTX helpers. Thin zero-cost wrappers over the exact instruction forms the +// kernel needs — nothing generic, nothing speculative. +// ============================================================================ +namespace ptx { + +// Generic pointer -> 32-bit `.shared` window byte address. +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Local `.shared` address -> the same offset in CTA `cta_rank`'s smem +// (`.shared::cluster` instructions take this mapped form). +static __device__ __forceinline__ uint32_t mapa_shared_cluster( + uint32_t local_addr, uint32_t cta_rank) { + uint32_t mapped; + asm("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(mapped) : "r"(local_addr), "r"(cta_rank)); + return mapped; +} + +static __device__ __forceinline__ void bar_sync(uint32_t id, uint32_t count) { + asm volatile("barrier.cta.sync %0, %1;" :: "r"(id), "r"(count)); +} + +// Cluster barrier with release/acquire — the publish boundary between the +// mbarrier inits and any cross-CTA use of those mbarriers. +static __device__ __forceinline__ void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// %cluster_ctarank: 0 or 1 in this kernel's (2,1,1) clusters. +static __device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} + + +// ---- mbarriers ------------------------------------------------------------- + +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + :: "r"(to_shared(bar)), "r"(count)); +} + +// Publish the inits to the async proxy + the peer CTA before any TMA +// completes on them. +static __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// Cross-CTA arrive with release ordering (teardown handshake). +static __device__ __forceinline__ void mbar_arrive_cluster_release( + uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" + :: "r"(mapped)); +} + +// Fused `tcgen05.wait::ld` + "buffer free" arrive, one asm block. The arrive +// must not become visible while TMEM loads are still in flight (it frees the +// accumulator buffer to the next tile's MMA — the bring-up chapter's +// drain overlap); the fused block +// keeps the arrive pinned behind the wait. (An earlier form also manufactured +// a register dependency on a drained value, `dep * 0` folded into the +// address; verified dead on nvcc 13.1 — ptxas const-folds it, SASS +// byte-identical without it — and retired.) +static __device__ __forceinline__ void tcgen05_wait_ld_then_arrive_release( + uint64_t* bar, uint32_t cta_rank, bool do_arrive) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile( + "{\n\t" + ".reg .pred %%pdo;\n\t" + "setp.ne.s32 %%pdo, %1, 0;\n\t" + "tcgen05.wait::ld.sync.aligned;\n\t" + "@%%pdo mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;\n\t" + "}\n\t" + :: "r"(mapped), "r"((int)do_arrive) : "memory"); +} + + +// Blocking parity wait at the consuming site (producer ring/tail waits + +// teardown). The 10 ms suspend hint is the polite form: a full-rate spin +// competes for issue slots and measurably loses; a nanosleep-backoff spin +// measured the same as this hint across five windows, so the simpler form +// ships. The carried one-stage-ahead try_wait peeks that used to gate these +// waits measured as pure overhead and were removed. +static __device__ __forceinline__ void mbar_wait_parity( + uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(to_shared(bar)), "r"(parity)); +} + +// ---- TMA ------------------------------------------------------------------- + +// Producer expect_tx arrive on the cluster-shared ring barrier: both +// peers' transfer bytes account against the even CTA's barrier, so one +// barrier per ring slot covers the whole cluster-wide gather. (Mechanism: +// `& 0xFEFFFFFF` clears bit 24 of the shared address — the hardware's +// peer-CTA bit — re-pointing the arrive at the even CTA's copy.) +static __device__ __forceinline__ void mbar_arrive_expect_tx_cluster( + uint64_t* bar, uint32_t bytes) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes)); +} + +// 3D tiled load, cta_group::2 + multicast (the A feed; the mask is the +// issuing CTA itself in a (2,1) cluster). Completion routes to the even +// CTA's mbarrier via the bit-24 clear. +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z) + : "memory"); +} + +// 3D tiled load, cta_group::2, no multicast (the B feed — each peer loads +// its own 128-wide N-block). +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_2sm_bit24( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst_smem), "l"(tmap), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) + : "memory"); +} + +// 4D tiled load, cta_group::2 + multicast (the scale-factor feed). The +// 512-byte SF tiles of one ring slot form a 4D box; `w` selects the slot. +static __device__ __forceinline__ void cp_async_bulk_tensor_4d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z2, int32_t w, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.4d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z2), "r"(w) + : "memory"); +} + + +// ---- tcgen05 (tensor memory + MMA plumbing) --------------------------------- + +static __device__ __forceinline__ void tcgen05_alloc_2sm( + uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_dealloc_2sm( + uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" + :: "r"(taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_relinquish_2sm() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); +} + +// The widest single-instruction TMEM read: 32 b32 registers per lane +// (blog: the epilogue — a quarter the load/wait count of the narrow forms). +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t taddr, uint32_t* r) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), + "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), + "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), + "=r"(r[16]), "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), + "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), + "=r"(r[24]), "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), + "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr)); +} + +// ---- stores + conversion ----------------------------------------------------- + +static __device__ __forceinline__ uint32_t cvt_pack_f16x2(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// Baseline output store: use the default cache policy. +static __device__ __forceinline__ void st_global_b128(uint4* ptr, uint4 v) { + asm volatile("st.global.v4.b32 [%0], {%1, %2, %3, %4};" + :: "l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)); +} + + + +// ---- MMA descriptors --------------------------------------------------------- + +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class ScaleFormat : uint8_t { E4M3 = 0, E8M0 = 1 }; + +// Source descriptor for the tcgen05.cp that stages scale factors from smem +// into TMEM: unswizzled, SBO=128, LBO=16. The low word later carries the +// per-slot byte offset (in 16-byte units). +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_sf_cp_src_desc_cd(uint32_t sf_smem_addr) { + uint64_t addr_field = uint64_t((sf_smem_addr >> 4) & 0x3FC0u); + return addr_field | 0x400800010000ull; +} + +// Operand descriptor for a K=96 (48-byte) slice inside a 128-byte-swizzled +// smem K-window. OR-constant 0x4010404000000000: version=1, absolute-address +// LBO mode (bit 52), 128B swizzle, SBO=1024. Both the address field and the +// absolute-LBO field encode the window's own 1024-aligned base; the +// intra-window offset advances the address field only. +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k96_sw128(uint32_t window_base_1024aligned, int intra_window_off = 0) { + uint64_t addr_field = uint64_t(((window_base_1024aligned >> 4) & 0x7FC0u) + + uint32_t(intra_window_off >> 4)); + uint64_t lbo_field = uint64_t((window_base_1024aligned << 12) & 0x7FC00000u); + return addr_field | lbo_field | 0x4010404000000000ull; +} + +// Instruction descriptor for tcgen05.mma.kind::mxf4nvf4.block_scale. +// The K extent is selected by BIT 31, not by the asm mnemonic: k_dim=0 runs +// K=64 (0x10400480), k_dim=1 runs K=96 (0x90400480, sm_103a-only) — the +// cheapest quarter of the gap if your kernel is built on the K=64 form +// (blog: the MMA warp). Bit 13 flips which TMEM scale-factor column the MMA reads. +__host__ __device__ static __forceinline__ constexpr uint32_t +mma_inst_desc_mxf4nvf4_block16( + uint32_t M, uint32_t N, + uint32_t tmem_sfa_addr, uint32_t tmem_sfb_addr, + Major a_major, Major b_major, + bool negate_a, bool negate_b, + bool k_dim, ScaleFormat sf) { + constexpr uint32_t MXF4_E2M1 = 1u; + uint32_t d = 0; + d |= ((tmem_sfb_addr & 0xC0000000u) >> 30) << 4; // b_sf_id + d |= MXF4_E2M1 << 7; // a_format = E2M1 + d |= MXF4_E2M1 << 10; // b_format = E2M1 + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; // N in units of 8 + d |= (static_cast(sf) & 0x1u) << 23; // 0 = UE4M3 + d |= ((M >> 4) & 0x1Fu) << 24; // m_dim at bits 27-28 + d |= ((tmem_sfa_addr & 0xC0000000u) >> 30) << 29; // a_sf_id + if (k_dim) d |= 1u << 31; // 1 = K=96 + return d; +} + +} // namespace ptx + +// ============================================================================ +// Kernel geometry from the blog's opening kernel section. One tile shape, +// one cluster shape: +// 256x256 output per 2-CTA cluster (128 rows per CTA), K advanced in +// 96-element MMA steps staged as 768-element K-groups of 8 steps, persistent +// grid of 76 clusters = 152 CTAs = one per SM. +// ============================================================================ +namespace nvfp4 { + +constexpr int CTA_GROUP = 2; // cta_group::2 — the MMA spans both CTAs +constexpr int BLOCK_M = 128; // A rows per CTA +constexpr int BLOCK_N = 256; // output tile N extent (per cluster) +constexpr int BLOCK_N_PEER = BLOCK_N / CTA_GROUP; // 128 — B cols per CTA +constexpr int CLUSTER_M = CTA_GROUP * BLOCK_M; // 256 — tile M extent +[[maybe_unused]] constexpr int MMA_K = 96; // K per tcgen05.mma (idesc bit 31) +constexpr int WIN_BYTES = 128; // one 128B-swizzled smem K-window +constexpr int AB_LOAD_BYTES = MMA_K; +constexpr int AB_TMA_SUBS = 4; +static_assert(AB_LOAD_BYTES == 96, + "K=192 packed-FP4 stage must be 96 bytes wide"); +constexpr int TB_SIZE = 224; // 7 warps (blog: the warp table) +constexpr int N_EPI_WARPS = 4; // warps 0-3 + +constexpr int SF_TMEM_COLS = 9 * 4; // 3 SFA + 6 SFB cp tiles, 4 cols each +// Bring-up layout: one accumulator followed immediately by the scale factors. +constexpr int D_STRIDE = BLOCK_N; +constexpr int D_DB = 1; +constexpr int N_D_COLS = BLOCK_N; +constexpr int SF_TMEM_BASE = BLOCK_N; +constexpr int N_ALLOC_COLS = 512; +static_assert(N_D_COLS <= SF_TMEM_BASE && + SF_TMEM_BASE + SF_TMEM_COLS <= N_ALLOC_COLS, + "accumulators + scale factors exceed the 512 TMEM columns"); +constexpr int EPI_CHUNK = 32; // cols per regular drain + +// ---- the shared-memory map (blog: the feed) ----------------------------------- +// Every offset is a compile-time expression. +constexpr int CD_HDR = 1024; // mbarrier header pad +constexpr int CD_AB_SLOTS = 6; +constexpr int CD_SF_SLOTS = 7; +constexpr int CD_AB_STRIDE = BLOCK_M * WIN_BYTES; // 16384 per A or B slot +constexpr int CD_SF_TILE = 512; // one (32 lanes x 16B) x4 tile +constexpr int CD_SFA_TILES = 3; // per SFA slot (M = 1 block) +constexpr int CD_SFB_TILES = 6; // per SFB slot (N = 2 blocks) +constexpr int CD_SFA_STRIDE = CD_SFA_TILES * CD_SF_TILE; // 1536 +constexpr int CD_SFB_STRIDE = CD_SFB_TILES * CD_SF_TILE; // 3072 + +// Expected transfer bytes are derived, never tuned: a wrong-by-one +// expectation doesn't crash, it deadlocks (blog: the feed). +constexpr uint32_t AB_TX = uint32_t(CTA_GROUP) * + uint32_t(BLOCK_M + BLOCK_N_PEER) * uint32_t(AB_LOAD_BYTES); +constexpr uint32_t SF_TX = uint32_t(CD_SFA_STRIDE + CD_SFB_STRIDE) * + uint32_t(CTA_GROUP); // 9216 +static_assert(AB_TX == 49152u && SF_TX == 9216u, + "expect_tx derivation drifted"); + +// The map IS the struct: member sizes and alignment produce every byte +// offset; the raw-address constants the mainloop uses read off it with +// offsetof, and the static_asserts pin the published values. +struct SmemCD { + // mbarrier header: four ring regions in handshake order, then the + // singleton accumulator + teardown barriers (alignas pads it to 1024). + struct alignas(1024) Header { + uint64_t ab_ready[CD_AB_SLOTS]; // TMA -> MMA + uint64_t ab_free [CD_AB_SLOTS]; // MMA -> TMA + uint64_t sf_ready[CD_SF_SLOTS]; // TMA -> MMA + uint64_t sf_free [CD_SF_SLOTS]; // MMA -> TMA + uint64_t acc_ready; // MMA -> epi + uint64_t acc_free; // epi -> MMA + uint64_t dealloc; // 2-CTA teardown + uint32_t tmem_addr; // TMEM base addr storage + } hdr; + uint8_t a [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t b [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t sfa[CD_SF_SLOTS][CD_SFA_STRIDE]; + // SFB starts 1024-aligned (the cp source descriptor drops address bits 4-9). + alignas(1024) uint8_t sfb[CD_SF_SLOTS][CD_SFB_STRIDE]; +}; +static_assert(sizeof(SmemCD::Header) == CD_HDR, "mbarrier header overflows the pad"); + +constexpr int CD_A_OFF = offsetof(SmemCD, a); +constexpr int CD_B_OFF = offsetof(SmemCD, b); +constexpr int CD_SFA_OFF = offsetof(SmemCD, sfa); +constexpr int CD_SFB_OFF = offsetof(SmemCD, sfb); +constexpr int CD_TOTAL = sizeof(SmemCD); +static_assert(CD_TOTAL <= 232448, "smem map exceeds the opt-in budget"); +static_assert(CD_AB_SLOTS != 6 || + (CD_B_OFF == 99328 && CD_SFA_OFF == 197632 && + CD_SFB_OFF == 208896 && CD_TOTAL == 230400), + "the 6-slot byte map drifted from its published values"); + +// The mainloop computes raw `.shared` byte addresses from these. +constexpr int MB_AB_READY = offsetof(SmemCD, hdr.ab_ready); +constexpr int MB_AB_FREE = offsetof(SmemCD, hdr.ab_free); +constexpr int MB_SF_READY = offsetof(SmemCD, hdr.sf_ready); +constexpr int MB_SF_FREE = offsetof(SmemCD, hdr.sf_free); +constexpr int MB_ACC_READY = offsetof(SmemCD, hdr.acc_ready); +constexpr int MB_ACC_FREE = offsetof(SmemCD, hdr.acc_free); + + + +} // namespace nvfp4 + +// ============================================================================ +// The consumer mainloop from the blog's pipeline section. One warp — warp 4 +// of the even CTA — +// issues everything: the tcgen05.cp copies that stage scale factors from the +// smem ring into TMEM, immediately followed by the tcgen05.mma stream, in the +// SAME warp. A same-warp tcgen05 stream is one ordered async pipeline, so the +// write-after-read hazard on the single reused scale-factor TMEM region needs +// no fence and no cross-warp barrier. +// +// The baseline stages four K=192 slices. Each slot supplies two K=96 MMAs at +// packed-byte offsets 0 and 48, both inside the same swizzle window. +// ============================================================================ +namespace nvfp4 { + + +// Ring/parity state carried across tiles by the MMA warp. +struct MainloopState { + uint32_t d_parity; // which TMEM accumulator buffer the next tile writes + uint32_t ab_phase, ab_slot; // A/B ring cursor + uint32_t sf_phase, sf_slot; // scale-factor ring cursor +}; + + +// Relaxed blocking wait (no .acquire): the tcgen05 pipeline's own ordering +// covers these ring waits, and the .acquire form perturbs the scoreboard +// cadence on the issuing warp. Waits sit AT the consuming site — the carried +// one-stage-ahead try_wait peeks that used to gate them measured as pure +// overhead and were removed. +static __device__ __forceinline__ void mbar_wait_relaxed( + uint32_t bar_smem, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(bar_smem), "r"(parity) : "memory"); +} + +// Gated form for the @216 buffer-free wait, which fires once per tile: the +// gate branches over the spin for every group past the first. +static __device__ __forceinline__ void mbar_wait_gated_relaxed( + uint32_t bar_smem, uint32_t parity, bool peek) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p bra DONE_%=;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t" + "DONE_%=:\n\t}\n" + :: "r"(bar_smem), "r"(parity), "r"((uint32_t)peek) : "memory"); +} + + +// One scale-factor staging copy, smem -> TMEM. Single-lane semantics behind +// the converged elect predicate: tcgen05.cp is a single-thread instruction — +// issued bare from a converged warp you get 32 initiations and garbage. +static __device__ __forceinline__ void cp_sf( + uint32_t el, uint32_t dst, uint32_t lo, uint32_t hi) { + const uint64_t desc = (uint64_t(hi) << 32) | uint64_t(lo); + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;\n\t}\n" + :: "r"(dst), "l"(desc), "r"(el) : "memory"); +} + +// Ring-slot "consumed" arrive, multicast to both CTAs of the pair. +// arrive::one must fire exactly once — hence the elect predicate. +static __device__ __forceinline__ void commit_consumed( + uint32_t el, uint32_t smbase, uint32_t mb_off, uint16_t mask) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t}\n" + :: "r"(smbase + mb_off), "h"(mask), "r"(el) : "memory"); +} + +// One K=96 NVFP4 MMA. The asm token says scale_vec::4X regardless — K=96 is +// selected by idesc bit 31 (blog: the MMA warp). acc_pred = 0 overwrites the +// accumulator (the tile's first step), nonzero accumulates. +static __device__ __forceinline__ void mma_cell( + uint32_t el, uint32_t d, + uint32_t desc_a_lo, uint32_t desc_a_hi, + uint32_t desc_b_lo, uint32_t desc_b_hi, + uint32_t idesc, uint32_t sf_a, uint32_t sf_b, uint32_t acc_pred) { + const uint64_t desc_a = (uint64_t(desc_a_hi) << 32) | uint64_t(desc_a_lo); + const uint64_t desc_b = (uint64_t(desc_b_hi) << 32) | uint64_t(desc_b_lo); + asm volatile( + "{\n\t.reg .pred p, acc;\n\t" + "setp.ne.u32 p, %7, 0;\n\t" + "setp.ne.b32 acc, %4, 0;\n\t" + "@p tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X" + " [%0], %1, %2, %3, [%5], [%6], acc;\n\t}\n" + :: "r"(d), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(acc_pred), "r"(sf_a), "r"(sf_b), "r"(el) : "memory"); +} + +// One SF ring slot -> TMEM: 9 tcgen05.cp tiles (3 SFA + 6 SFB). Source +// address advances are in 16-byte descriptor units. The SFB tile order +// 0,3,1,4,2,5 interleaves the two 128-wide N-blocks. +static __device__ __forceinline__ void cp_sf_buffer( + uint32_t el, uint32_t slot, + uint32_t sfa0, uint32_t sfa1, uint32_t sfa2, + uint32_t sfb0, uint32_t sfb1, uint32_t sfb2, + uint32_t sfb3, uint32_t sfb4, uint32_t sfb5, + uint32_t sfa_hi, uint32_t sfa_lo, uint32_t sfb_hi, uint32_t sfb_lo) { + constexpr uint32_t SFA_SLOT_U = CD_SFA_STRIDE / 16u; // 96 + constexpr uint32_t SFB_SLOT_U = CD_SFB_STRIDE / 16u; // 192 + constexpr uint32_t TILE_U = CD_SF_TILE / 16u; // 32 + const uint32_t a = slot * SFA_SLOT_U + sfa_lo; + cp_sf(el, sfa0, a, sfa_hi); + cp_sf(el, sfa1, a + 1u*TILE_U, sfa_hi); + cp_sf(el, sfa2, a + 2u*TILE_U, sfa_hi); + const uint32_t b = slot * SFB_SLOT_U + sfb_lo; + cp_sf(el, sfb0, b + 0u*TILE_U, sfb_hi); + cp_sf(el, sfb1, b + 3u*TILE_U, sfb_hi); + cp_sf(el, sfb2, b + 1u*TILE_U, sfb_hi); + cp_sf(el, sfb3, b + 4u*TILE_U, sfb_hi); + cp_sf(el, sfb4, b + 2u*TILE_U, sfb_hi); + cp_sf(el, sfb5, b + 5u*TILE_U, sfb_hi); +} + +// Operand descriptor LOW word for A/B ring slot `slot`, K-step `k_atom` +// (0..7 inside the 768-element group; each step is 48 bytes = +3 in the +// 16-byte address field). +static __device__ __forceinline__ uint32_t ab_desc_lo( + uint32_t smbase, uint32_t ab_base_off, uint32_t slot, uint32_t k_atom) { + const uint64_t d = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + slot * CD_AB_STRIDE, int(k_atom) * 16); + return uint32_t(d & 0xFFFFFFFFu); +} + + +// One output tile: `full_groups` complete 768-element K-groups. Coordinates +// past logical K are zero-filled by TMA from the unchanged canonical inputs. +static __device__ __forceinline__ MainloopState mainloop_tile( + uint32_t taddr, uint32_t smbase, MainloopState in, uint16_t commit_mask, + uint32_t full_groups + ) { + constexpr uint32_t AB_RING_SLOTS = uint32_t(CD_AB_SLOTS); + constexpr uint32_t SF_RING_SLOTS = uint32_t(CD_SF_SLOTS); + constexpr uint32_t AB_FREE_DELTA = uint32_t(MB_AB_FREE - MB_AB_READY); + constexpr uint32_t SF_READY_OFF = uint32_t(MB_SF_READY); + constexpr uint32_t SF_FREE_DELTA = uint32_t(MB_SF_FREE - MB_SF_READY); + constexpr uint32_t ACC_FREE_OFF = uint32_t(MB_ACC_FREE); + constexpr uint32_t A_OFF = uint32_t(CD_A_OFF); + constexpr uint32_t B_OFF = uint32_t(CD_B_OFF); + + uint32_t d_parity = in.d_parity; + uint32_t ab_phase = in.ab_phase, ab_slot = in.ab_slot; + uint32_t sf_phase = in.sf_phase, sf_slot = in.sf_slot; + + // The caller has already selected exactly one lane for the whole MMA role. + // Keep the helper predicate constant throughout the tile. + const uint32_t el = 1u; + + // Fixed TMEM targets for the single reused scale-factor region: 9 cp + // tiles of 4 columns each, starting past the accumulator region. + constexpr uint32_t SF = uint32_t(SF_TMEM_BASE); + const uint32_t sfa_t0 = taddr + SF + 0; // SFA word 0 (MMA reads here) + const uint32_t sfa_t1 = taddr + SF + 4; + const uint32_t sfa_t2 = taddr + SF + 8; + const uint32_t sfb_t0 = taddr + SF + 12; // SFB word 0 + const uint32_t sfb_t1 = taddr + SF + 16; + const uint32_t sfb_t2 = taddr + SF + 20; + const uint32_t sfb_t3 = taddr + SF + 24; + const uint32_t sfb_t4 = taddr + SF + 28; + const uint32_t sfb_t5 = taddr + SF + 32; + + // Scale-factor cp source descriptors (low word carries the per-slot + // offset; high word is constant). + const uint64_t sfb_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFB_OFF)); + const uint32_t sfb_src_lo = uint32_t(sfb_src & 0xFFFFFFFFu); + const uint32_t sfb_src_hi = uint32_t(sfb_src >> 32); + const uint64_t sfa_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFA_OFF)); + const uint32_t sfa_src_lo = uint32_t(sfa_src & 0xFFFFFFFFu); + const uint32_t sfa_src_hi = uint32_t(sfa_src >> 32); + + // Operand descriptor HIGH words (low words come from ab_desc_lo per step). + const uint32_t a_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + A_OFF) >> 32); + const uint32_t b_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + B_OFF) >> 32); + + // Each K=96 step needs 6 scale bytes per row, split across two staged + // 32-bit words. Odd-issue steps read word 1: bit 31 of the SF address + // selects the with-id encoding, and the id bits fold into the idesc. + constexpr uint32_t SF_ADDR_ID = 0x80000000u; + const uint32_t sfa_w1 = (taddr + SF + 4) | SF_ADDR_ID; + const uint32_t sfb_w1 = (taddr + SF + 20) | SF_ADDR_ID; + const uint32_t w1_id_bits = + ((sfa_w1 >> 1) & 0x60000000u) | ((sfb_w1 >> 26) & 48u); + + // The tile's first MMA overwrites the accumulator buffer this tile owns; + // the @216 "buffer free" wait runs once, before the first group. + const uint32_t acc_wait_parity = d_parity; + d_parity ^= 1u; + // Single buffer: every tile writes buf0; d_parity still toggles because + // it carries the acc_free handshake parity. + const uint32_t d_tmem = taddr; + + // Instruction descriptor: kind::mxf4nvf4, M=256, N=256, E2M1 x E2M1, + // UE4M3 block-16, K-major, k_dim=1 (K=96). Base value 0x90400480. + constexpr uint32_t IDESC = ptx::mma_inst_desc_mxf4nvf4_block16( + 256u, 256u, 0u, 0u, ptx::Major::K, ptx::Major::K, + false, false, /*k_dim=*/true, ptx::ScaleFormat::E4M3); + static_assert(IDESC == 0x90400480u, "K=96 NVFP4 idesc drifted"); + const uint32_t w0_id_bits = + ((sfa_t0 >> 1) & 0x60000000u) | ((sfb_t0 >> 26) & 48u); + const uint32_t idesc_w0 = w0_id_bits | IDESC; + const uint32_t idesc_w1 = w1_id_bits | IDESC; + + + uint32_t acc_pred = 0; // first MMA of the tile overwrites + bool need_acc_wait = true; // @216 wait fires once per tile + + if (full_groups >= 1) { + const uint32_t num_stages = full_groups * 4u; + uint32_t stage = 0; + do { + // One K=192 scale slot and one K=192 A/B slot feed two K=96 + // MMAs. Return each SMEM source through an ordered commit as soon + // as its final asynchronous reader has been issued. + const uint32_t sf_bar = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar, sf_phase); + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, + sfb_t0, sfb_t1, sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar + SF_FREE_DELTA) - smbase, + commit_mask); + + const uint32_t ab_bar = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar, ab_phase); + mbar_wait_gated_relaxed(smbase + ACC_FREE_OFF, + acc_wait_parity, !need_acc_wait); + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); + acc_pred = 0xFFFFFFFFu; + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); + commit_consumed(el, smbase, (ab_bar + AB_FREE_DELTA) - smbase, + commit_mask); + need_acc_wait = false; + + const uint32_t next_sf = sf_slot + 1; + const bool wrap_sf = (next_sf == SF_RING_SLOTS); + sf_slot = wrap_sf ? 0u : next_sf; + sf_phase ^= wrap_sf ? 1u : 0u; + + const uint32_t next_ab = ab_slot + 1; + const bool wrap_ab = (next_ab == AB_RING_SLOTS); + ab_slot = wrap_ab ? 0u : next_ab; + ab_phase ^= wrap_ab ? 1u : 0u; + + ++stage; + } while (stage != num_stages); + } + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; +} + +} // namespace nvfp4 + +// ============================================================================ +// The kernel. 224 threads = 7 warps, four jobs (blog: the warp table): +// warps 0-3 epilogue: drain TMEM, convert to FP16, store straight to global +// warp 4 the MMA warp (even CTA of each pair only): SF copies + MMAs +// warp 5 TMA producer: A and B tiles into the 6-slot ring +// warp 6 TMA producer: scale factors into the 7-slot ring +// ============================================================================ +namespace nvfp4 { + +// Baseline launch: plain launch bounds and a grid expressed in CTA units. +#define NVFP4_BLOCK_ATTR +constexpr bool GRID_IN_CLUSTERS = false; + +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void nvfp4_gemm_kernel(const __grid_constant__ CUtensorMap A_tmap, + const __grid_constant__ CUtensorMap B_tmap, + const __grid_constant__ CUtensorMap SFA_tmap, + const __grid_constant__ CUtensorMap SFB_tmap, + __half* C_ptr, int M, int N, int K + ) { + extern __shared__ __align__(1024) char smem_buf[]; + SmemCD& smem = *reinterpret_cast(smem_buf); + + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const uint32_t crank = ptx::cluster_rank(); + const uint32_t m_pair = crank & 1u; // 0 = the pair's even CTA + + const int cluster_grid_m = (M + CLUSTER_M - 1) / CLUSTER_M; + const int grid_n = (N + BLOCK_N - 1) / BLOCK_N; + const int total_tiles = cluster_grid_m * grid_n; + + // Baseline rungs execute a whole K=768 group at the edge. Tensor-map OOB + // fill supplies zeros past logical K without changing the input buffers. + const int full_groups = (K + 767) / 768; + const int num_groups = full_groups; + + // Persistent cluster index. The device always sees the grid in CTA units + // (2, 1, clusters) regardless of the host-side launch-unit regime. + const bool grid_z = (gridDim.z > 1); + const int num_clusters = grid_z ? int(gridDim.z) : int(gridDim.x) / 2; + const int cluster_id = grid_z ? int(blockIdx.z) : int(blockIdx.x) / 2; + + + // Which output tile does persistent work item `t` process? + auto tile_of = [&](int t) -> int { + return t; // natural raster: M fast, N slow + }; + + auto producer_tile_of = [&](int t) -> int { + return tile_of(t); + }; + + // ---- mbarrier init (warp 0, lane 0) ---- + if (warp_id == 0 && (threadIdx.x % 32) == 0) { + // Arrival counts are derived from the topology, never tuned: + // "ready" = one producer expect_tx per slot (bit-24 routes both + // peers' bytes to the even CTA's barrier); "free" = one consumer + // commit per slot reuse. + #pragma unroll + for (int s = 0; s < CD_AB_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.ab_ready[s], 1); + ptx::mbar_init(&smem.hdr.ab_free[s], 1); + } + #pragma unroll + for (int s = 0; s < CD_SF_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.sf_ready[s], 1); + ptx::mbar_init(&smem.hdr.sf_free[s], 1); + } + ptx::mbar_init(&smem.hdr.acc_ready, 1); + ptx::mbar_init(&smem.hdr.acc_free, N_EPI_WARPS * CTA_GROUP); // 4 x 2 = 8 + ptx::mbar_init(&smem.hdr.dealloc, 32); // one warp's lanes, cross-peer + ptx::fence_mbarrier_init_release_cluster(); + } + ptx::cluster_sync_rel_acq(); + + // TMEM allocation: epilogue warp 0 allocates for the pair; warps 0-4 + // rendezvous before reading the granted base. The TMA producer warps + // never touch TMEM and proceed straight to their loops. + uint32_t taddr = 0; + if (warp_id <= 4) { + if (warp_id == 0) + ptx::tcgen05_alloc_2sm(ptx::to_shared(&smem.hdr.tmem_addr), + N_ALLOC_COLS); + ptx::bar_sync(2, 160); + taddr = smem.hdr.tmem_addr; + } + + const int wg = warp_id >> 2; // 0 = epilogue warps, 1 = producers + MMA + if (wg == 1) { + const int pw = warp_id - 4; + if (pw == 1 && (threadIdx.x % 32) == 0) { + // ================================================================ + // Warp 5: lane 0 owns the complete A/B TMA producer loop. + // ================================================================ + // free_phase starts at 1: a fresh mbarrier is phase 0, so the + // first ring's worth of "free" waits falls through. + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + // The pair splits the 256x256 tile: even CTA holds A rows + // [0,128) and B cols [0,128); odd CTA the other halves. + const int m_blk = m_row * CTA_GROUP + int(m_pair); + const int n_blk = n_grp * CTA_GROUP + int(m_pair); + #pragma unroll 1 + for (int g = 0; g < num_groups; ++g) { + // Four 96-byte loads cover the packed K=768 group. + #pragma unroll + for (int sub = 0; sub < AB_TMA_SUBS; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.ab_free[slot], free_phase); + const uint32_t a_s = + base + CD_A_OFF + uint32_t(slot) * CD_AB_STRIDE; + const uint32_t b_s = + base + CD_B_OFF + uint32_t(slot) * CD_AB_STRIDE; + uint64_t* bar = &smem.hdr.ab_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, AB_TX); + // A rides the multicast form with a self mask so the + // completion routes to the even peer's barrier. + const uint16_t mask_a = uint16_t(1u << m_pair); + // Read the same row-major FP4 allocation as cuBLAS. + // K192 uses four 96-byte boxes; K256 uses three + // 128-byte boxes. The tensor map zero-fills past K. + const CUtensorMap* a_map = &A_tmap; + const CUtensorMap* b_map = &B_tmap; + const int x = g * 384 + sub * AB_LOAD_BYTES; + const int y_a = m_blk * BLOCK_M; + const int y_b = n_blk * BLOCK_N_PEER; + const int z_a = 0; + const int z_b = 0; + ptx::cp_async_bulk_tensor_3d_load_multicast( + a_s, a_map, x, y_a, z_a, + bar, mask_a + ); + ptx::cp_async_bulk_tensor_3d_load_2sm_bit24( + b_s, b_map, x, y_b, z_b, + bar + ); + if (++slot == CD_AB_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + // Drain: wait for the last-filled slot to be consumed, then hand + // it one clean-shutdown expect_tx. + { + const int t_slot = (slot == 0) ? CD_AB_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.ab_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.ab_ready[t_slot], AB_TX); + } + } else if (pw == 2 && (threadIdx.x % 32) == 0) { + // ================================================================ + // Warp 6: scale-factor TMA producer. SFA is per-peer (each CTA's + // own M half); SFB is the one genuinely multicast operand — the + // cta_group::2 copy reads BOTH peers' smem as one logical + // operand, so both copies must be byte-identical (blog: the scale-factor section). + // ================================================================ + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const int m_blk = m_row * CTA_GROUP + int(m_pair); + for (int g = 0; g < num_groups; ++g) { + // 4 SF ring slots per 768-K group (2 K=96 steps each). + #pragma unroll + for (int sub = 0; sub < 4; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.sf_free[slot], free_phase); + const uint32_t sfa_s = + base + CD_SFA_OFF + uint32_t(slot) * CD_SFA_STRIDE; + const uint32_t sfb_s = + base + CD_SFB_OFF + uint32_t(slot) * CD_SFB_STRIDE; + uint64_t* bar = &smem.hdr.sf_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, SF_TX); + // VEC16 is [outer-block][K-group-of-4][128 rows x 4]. + // Three adjacent groups are the 128 x 12 scale slot. + const int sf_group = (g * 4 + sub) * 3; + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfa_s, &SFA_tmap, 0, 0, sf_group, m_blk, bar, + uint16_t(1u << m_pair) + ); + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfb_s + m_pair * uint32_t(CD_SFB_STRIDE / 2), + &SFB_tmap, 0, 0, sf_group, + n_grp * CTA_GROUP + int(m_pair), bar, + uint16_t(0x3) + ); + if (++slot == CD_SF_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + { + const int t_slot = (slot == 0) ? CD_SF_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.sf_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.sf_ready[t_slot], SF_TX); + } + } else if (pw == 0 && m_pair == 0 && + (threadIdx.x % 32) == 0) { + // ================================================================ + // Warp 4: the MMA warp — issued once per cluster, on the even + // CTA. The odd CTA's warp 4 does nothing all kernel (blog: the warp table). + // The mainloop needs no tile coordinates: operands arrive through + // the rings; this warp only needs its share of the trip count. + // ================================================================ + const uint32_t smbase = ptx::to_shared(&smem); + const uint16_t commit_mask = 0x3; // both CTAs of this pair + MainloopState st{1u, 0u, 0u, 0u, 0u}; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + st = mainloop_tile(taddr, smbase, st, commit_mask, + uint32_t(full_groups) + ); + // The outer role branch already leaves exactly one issuer. + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;" + :: "r"(smbase + uint32_t(MB_ACC_READY)), + "h"(uint16_t(0x3)) : "memory"); + } + // Drain the final tile's buffer-free arrive before exiting. + mbar_wait_relaxed(smbase + uint32_t(MB_ACC_FREE), st.d_parity); + } + } else { + // ==================================================================== + // Warps 0-3: the epilogue (blog: the epilogue). Each warp owns a 32-row band + // and drains a tile as eight 32-column tcgen05.ld loads, converts + // FP32->FP16 in packed pairs, and stores straight to global — no + // shared-memory bounce. + // ==================================================================== + const int epi_warp = warp_id; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + const uint32_t acc_free_rank = crank & ~1u; // the pair's even CTA + int d_db = 0; + int acc_ready_phase = 0; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const uint32_t tmem_base = taddr + uint32_t(d_db) * D_STRIDE; + mbar_wait_relaxed(ptx::to_shared(&smem.hdr.acc_ready), + uint32_t(acc_ready_phase)); + acc_ready_phase ^= 1; + + const int off_m = (m_row * CTA_GROUP + int(m_pair)) * BLOCK_M; + const int off_n = n_grp * BLOCK_N; + const int row = off_m + epi_warp * 32 + lane_id; + const bool row_in = (row < M); + constexpr int EPI_BANDS = BLOCK_N / EPI_CHUNK; // 8 + #pragma unroll 1 + for (int k = 0; k < EPI_BANDS; ++k) { + const int band = k; + const uint32_t taddr_n = + tmem_base + uint32_t(band) * 32 + taddr_lane; + uint32_t w[32]; + ptx::tcgen05_ld_32x32b_x32(taddr_n, w); + // One buffer cannot be reused until all eight bands are out. + if (k == EPI_BANDS - 1) { + ptx::tcgen05_wait_ld_then_arrive_release( + &smem.hdr.acc_free, acc_free_rank, + /*do_arrive=*/(threadIdx.x % 32) == 0); + } + if (!row_in) continue; + const int col_base = off_n + band * 32; + uint32_t h[16]; + #pragma unroll + for (int j = 0; j < 16; ++j) { + h[j] = ptx::cvt_pack_f16x2(__int_as_float(w[2 * j + 1]), + __int_as_float(w[2 * j])); + } + const size_t row_base = size_t(row) * size_t(N); + const bool aligned32 = ((row_base & 15u) == 0u) && + ((uintptr_t(C_ptr) & 31u) == 0u); + if (aligned32 && (col_base & 15) == 0) { + // Precondition here: row_base is 16-element (32-byte) + // aligned — the aligned32 guard owns it. Restating it to + // ptxas via __builtin_assume compiled this file + // SASS-identical (and measured perf-null on the + // production kernel), so the guard stands alone. + #pragma unroll + for (int q = 0; q < 2; ++q) { + const int c0 = col_base + q * 16; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 16 <= N) { + uint4 lo{h[q*8+0], h[q*8+1], h[q*8+2], h[q*8+3]}; + uint4 hi{h[q*8+4], h[q*8+5], h[q*8+6], h[q*8+7]}; + // Bring-up form: twice the store requests, + // twice the L2 store traffic (the ladder's + // tensor-activity 82% -> 42% row). + ptx::st_global_b128( + reinterpret_cast(cdst), lo); + ptx::st_global_b128( + reinterpret_cast(cdst) + 1, hi); + } else { + // N tail: 128-bit then scalar. + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int cq = c0 + p * 8; + if (cq + 8 <= N) { + uint4 v{h[q*8+p*4+0], h[q*8+p*4+1], + h[q*8+p*4+2], h[q*8+p*4+3]}; + ptx::st_global_b128( + reinterpret_cast(cdst + p * 8), v); + } else { + const uint16_t* hb = + reinterpret_cast( + &h[q * 8 + p * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (cq + c < N) cdst[p * 8 + c] = hb[c]; + } + } + } + } + } else { + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int c0 = col_base + q * 8; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 8 <= N) { + uint4 v{h[q*4+0], h[q*4+1], h[q*4+2], h[q*4+3]}; + ptx::st_global_b128(reinterpret_cast(cdst), v); + } else { + const uint16_t* hb = + reinterpret_cast(&h[q * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (c0 + c < N) cdst[c] = hb[c]; + } + } + } + } + if (++d_db == D_DB) d_db = 0; + } + } + + // Teardown: only the epilogue warps rendezvous per CTA; warp 0 then runs + // the cross-peer dealloc handshake. Producer and MMA warps exit freely — + // their drains above already order their work. + if (wg == 0) { + ptx::bar_sync(3, 128); + if (warp_id == 0) { + ptx::tcgen05_relinquish_2sm(); + const uint32_t peer_rank = crank ^ 1u; + // All 32 lanes arrive the PEER's count-32 barrier = one flip. + ptx::mbar_arrive_cluster_release(&smem.hdr.dealloc, peer_rank); + ptx::mbar_wait_parity(&smem.hdr.dealloc, 0); + ptx::tcgen05_dealloc_2sm(taddr, N_ALLOC_COLS); + } + } +} + + +// Keep the host launch convention paired with the kernel's block attribute. +static inline dim3 launch_grid(int clusters) { + return GRID_IN_CLUSTERS ? dim3(1, 1, unsigned(clusters)) + : dim3(2, 1, unsigned(clusters)); +} + +} // namespace nvfp4 + +// ============================================================================ +// Host-side TMA descriptor support. The shared input and benchmark code lives +// in main.cu. +// ============================================================================ +namespace tmap { + +// Thin wrappers over cuTensorMapEncodeTiled. Dim 0 is the innermost +// (stride-1) axis; strides for dim 1+ are in bytes. +inline CUtensorMap encode_tiled_3d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_depth, uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t depth_stride_bytes, + uint32_t box_depth, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[3] = { global_cols, global_rows, global_depth }; + cuuint64_t global_strides[2] = { row_stride_bytes, depth_stride_bytes }; + cuuint32_t box_dim[3] = { box_cols, box_rows, box_depth }; + cuuint32_t element_strides[3] = { 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 3, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +inline CUtensorMap encode_tiled_4d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_slots, uint64_t global_tiles, + uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t tile_stride_bytes, + uint64_t slot_stride_bytes, + uint32_t box_tiles, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[4] = { global_cols, global_rows, global_tiles, + global_slots }; + cuuint64_t global_strides[3] = { row_stride_bytes, tile_stride_bytes, + slot_stride_bytes }; + cuuint32_t box_dim[4] = { box_cols, box_rows, box_tiles, 1u }; + cuuint32_t element_strides[4] = { 1, 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 4, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +} // namespace tmap + +namespace host { + +using namespace nvfp4; + +// ---- implementation-specific tensor maps --------------------------------- + +inline CUtensorMap make_ab_tmap(void* ptr, int rows, int K, + int row_stride) { + return tmap::encode_tiled_3d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*depth=*/1, uint64_t(rows), /*cols=*/uint64_t(K / 2), + /*row_stride=*/uint64_t(row_stride), + /*depth_stride=*/uint64_t(rows) * uint64_t(row_stride), + /*box=*/1, BLOCK_M, AB_LOAD_BYTES, CU_TENSOR_MAP_SWIZZLE_128B); +} + +inline CUtensorMap make_sf_tmap(uint8_t* ptr, int outer, int K) { + const uint64_t sf_inner = (uint64_t(K) + 63) / 64 * 4; + const uint64_t outer_blocks = (uint64_t(outer) + 127) / 128; + return tmap::encode_tiled_4d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*slots=*/outer_blocks, /*tiles=*/sf_inner / 4, + /*rows=*/4, /*cols=*/128, + /*row_stride=*/128, /*tile_stride=*/512, + /*slot_stride=*/sf_inner * 128, + /*box_tiles=*/3, /*box_rows=*/4, /*box_cols=*/128, + CU_TENSOR_MAP_SWIZZLE_NONE); +} +} // namespace host diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/gemm1.cuh b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm1.cuh new file mode 100644 index 000000000..1b985ed4f --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm1.cuh @@ -0,0 +1,1152 @@ +#pragma once +// r1: read %laneid directly for single-lane work. +// This is a complete, directly editable ladder snapshot. +// It intentionally duplicates the implementation in the other gemm*.cuh files. + +// Conditional branches are already materialized; edit this implementation directly. +#define NVFP4_BUILD_NAME "r1" +// Read-only metadata for main.cu; these do not select code in this header. +#define NVFP4_HAS_K64_TAIL 0 +#define NVFP4_HAS_SCHEDULE 0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(expr) do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \ + cudaGetErrorName(_e), __FILE__, __LINE__, \ + cudaGetErrorString(_e)); \ + std::abort(); \ + } \ +} while (0) + +#define CU_CHECK(expr) do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + std::fprintf(stderr, "CU error at %s:%d: %s\n", \ + __FILE__, __LINE__, _s ? _s : "?"); \ + std::abort(); \ + } \ +} while (0) + + +// ============================================================================ +// PTX helpers. Thin zero-cost wrappers over the exact instruction forms the +// kernel needs — nothing generic, nothing speculative. +// ============================================================================ +namespace ptx { + +// Generic pointer -> 32-bit `.shared` window byte address. +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Local `.shared` address -> the same offset in CTA `cta_rank`'s smem +// (`.shared::cluster` instructions take this mapped form). +static __device__ __forceinline__ uint32_t mapa_shared_cluster( + uint32_t local_addr, uint32_t cta_rank) { + uint32_t mapped; + asm("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(mapped) : "r"(local_addr), "r"(cta_rank)); + return mapped; +} + +static __device__ __forceinline__ void bar_sync(uint32_t id, uint32_t count) { + asm volatile("barrier.cta.sync %0, %1;" :: "r"(id), "r"(count)); +} + +// Cluster barrier with release/acquire — the publish boundary between the +// mbarrier inits and any cross-CTA use of those mbarriers. +static __device__ __forceinline__ void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// %cluster_ctarank: 0 or 1 in this kernel's (2,1,1) clusters. +static __device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} + +static __device__ __forceinline__ uint32_t lane_id() { + uint32_t r; + asm("mov.u32 %0, %%laneid;" : "=r"(r)); + return r; +} + +// ---- mbarriers ------------------------------------------------------------- + +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + :: "r"(to_shared(bar)), "r"(count)); +} + +// Publish the inits to the async proxy + the peer CTA before any TMA +// completes on them. +static __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// Cross-CTA arrive with release ordering (teardown handshake). +static __device__ __forceinline__ void mbar_arrive_cluster_release( + uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" + :: "r"(mapped)); +} + +// Fused `tcgen05.wait::ld` + "buffer free" arrive, one asm block. The arrive +// must not become visible while TMEM loads are still in flight (it frees the +// accumulator buffer to the next tile's MMA — the bring-up chapter's +// drain overlap); the fused block +// keeps the arrive pinned behind the wait. (An earlier form also manufactured +// a register dependency on a drained value, `dep * 0` folded into the +// address; verified dead on nvcc 13.1 — ptxas const-folds it, SASS +// byte-identical without it — and retired.) +static __device__ __forceinline__ void tcgen05_wait_ld_then_arrive_release( + uint64_t* bar, uint32_t cta_rank, bool do_arrive) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile( + "{\n\t" + ".reg .pred %%pdo;\n\t" + "setp.ne.s32 %%pdo, %1, 0;\n\t" + "tcgen05.wait::ld.sync.aligned;\n\t" + "@%%pdo mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;\n\t" + "}\n\t" + :: "r"(mapped), "r"((int)do_arrive) : "memory"); +} + + +// Blocking parity wait at the consuming site (producer ring/tail waits + +// teardown). The 10 ms suspend hint is the polite form: a full-rate spin +// competes for issue slots and measurably loses; a nanosleep-backoff spin +// measured the same as this hint across five windows, so the simpler form +// ships. The carried one-stage-ahead try_wait peeks that used to gate these +// waits measured as pure overhead and were removed. +static __device__ __forceinline__ void mbar_wait_parity( + uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(to_shared(bar)), "r"(parity)); +} + +// ---- TMA ------------------------------------------------------------------- + +// Producer expect_tx arrive on the cluster-shared ring barrier: both +// peers' transfer bytes account against the even CTA's barrier, so one +// barrier per ring slot covers the whole cluster-wide gather. (Mechanism: +// `& 0xFEFFFFFF` clears bit 24 of the shared address — the hardware's +// peer-CTA bit — re-pointing the arrive at the even CTA's copy.) +static __device__ __forceinline__ void mbar_arrive_expect_tx_cluster( + uint64_t* bar, uint32_t bytes) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes)); +} + +// 3D tiled load, cta_group::2 + multicast (the A feed; the mask is the +// issuing CTA itself in a (2,1) cluster). Completion routes to the even +// CTA's mbarrier via the bit-24 clear. +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z) + : "memory"); +} + +// 3D tiled load, cta_group::2, no multicast (the B feed — each peer loads +// its own 128-wide N-block). +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_2sm_bit24( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst_smem), "l"(tmap), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) + : "memory"); +} + +// 4D tiled load, cta_group::2 + multicast (the scale-factor feed). The +// 512-byte SF tiles of one ring slot form a 4D box; `w` selects the slot. +static __device__ __forceinline__ void cp_async_bulk_tensor_4d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z2, int32_t w, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.4d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z2), "r"(w) + : "memory"); +} + + +// ---- tcgen05 (tensor memory + MMA plumbing) --------------------------------- + +static __device__ __forceinline__ void tcgen05_alloc_2sm( + uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_dealloc_2sm( + uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" + :: "r"(taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_relinquish_2sm() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); +} + +// The widest single-instruction TMEM read: 32 b32 registers per lane +// (blog: the epilogue — a quarter the load/wait count of the narrow forms). +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t taddr, uint32_t* r) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), + "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), + "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), + "=r"(r[16]), "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), + "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), + "=r"(r[24]), "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), + "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr)); +} + +// ---- stores + conversion ----------------------------------------------------- + +static __device__ __forceinline__ uint32_t cvt_pack_f16x2(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// Baseline output store: use the default cache policy. +static __device__ __forceinline__ void st_global_b128(uint4* ptr, uint4 v) { + asm volatile("st.global.v4.b32 [%0], {%1, %2, %3, %4};" + :: "l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)); +} + + + +// ---- MMA descriptors --------------------------------------------------------- + +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class ScaleFormat : uint8_t { E4M3 = 0, E8M0 = 1 }; + +// Source descriptor for the tcgen05.cp that stages scale factors from smem +// into TMEM: unswizzled, SBO=128, LBO=16. The low word later carries the +// per-slot byte offset (in 16-byte units). +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_sf_cp_src_desc_cd(uint32_t sf_smem_addr) { + uint64_t addr_field = uint64_t((sf_smem_addr >> 4) & 0x3FC0u); + return addr_field | 0x400800010000ull; +} + +// Operand descriptor for a K=96 (48-byte) slice inside a 128-byte-swizzled +// smem K-window. OR-constant 0x4010404000000000: version=1, absolute-address +// LBO mode (bit 52), 128B swizzle, SBO=1024. Both the address field and the +// absolute-LBO field encode the window's own 1024-aligned base; the +// intra-window offset advances the address field only. +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k96_sw128(uint32_t window_base_1024aligned, int intra_window_off = 0) { + uint64_t addr_field = uint64_t(((window_base_1024aligned >> 4) & 0x7FC0u) + + uint32_t(intra_window_off >> 4)); + uint64_t lbo_field = uint64_t((window_base_1024aligned << 12) & 0x7FC00000u); + return addr_field | lbo_field | 0x4010404000000000ull; +} + +// Instruction descriptor for tcgen05.mma.kind::mxf4nvf4.block_scale. +// The K extent is selected by BIT 31, not by the asm mnemonic: k_dim=0 runs +// K=64 (0x10400480), k_dim=1 runs K=96 (0x90400480, sm_103a-only) — the +// cheapest quarter of the gap if your kernel is built on the K=64 form +// (blog: the MMA warp). Bit 13 flips which TMEM scale-factor column the MMA reads. +__host__ __device__ static __forceinline__ constexpr uint32_t +mma_inst_desc_mxf4nvf4_block16( + uint32_t M, uint32_t N, + uint32_t tmem_sfa_addr, uint32_t tmem_sfb_addr, + Major a_major, Major b_major, + bool negate_a, bool negate_b, + bool k_dim, ScaleFormat sf) { + constexpr uint32_t MXF4_E2M1 = 1u; + uint32_t d = 0; + d |= ((tmem_sfb_addr & 0xC0000000u) >> 30) << 4; // b_sf_id + d |= MXF4_E2M1 << 7; // a_format = E2M1 + d |= MXF4_E2M1 << 10; // b_format = E2M1 + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; // N in units of 8 + d |= (static_cast(sf) & 0x1u) << 23; // 0 = UE4M3 + d |= ((M >> 4) & 0x1Fu) << 24; // m_dim at bits 27-28 + d |= ((tmem_sfa_addr & 0xC0000000u) >> 30) << 29; // a_sf_id + if (k_dim) d |= 1u << 31; // 1 = K=96 + return d; +} + +} // namespace ptx + +// ============================================================================ +// Kernel geometry from the blog's opening kernel section. One tile shape, +// one cluster shape: +// 256x256 output per 2-CTA cluster (128 rows per CTA), K advanced in +// 96-element MMA steps staged as 768-element K-groups of 8 steps, persistent +// grid of 76 clusters = 152 CTAs = one per SM. +// ============================================================================ +namespace nvfp4 { + +constexpr int CTA_GROUP = 2; // cta_group::2 — the MMA spans both CTAs +constexpr int BLOCK_M = 128; // A rows per CTA +constexpr int BLOCK_N = 256; // output tile N extent (per cluster) +constexpr int BLOCK_N_PEER = BLOCK_N / CTA_GROUP; // 128 — B cols per CTA +constexpr int CLUSTER_M = CTA_GROUP * BLOCK_M; // 256 — tile M extent +[[maybe_unused]] constexpr int MMA_K = 96; // K per tcgen05.mma (idesc bit 31) +constexpr int WIN_BYTES = 128; // one 128B-swizzled smem K-window +constexpr int AB_LOAD_BYTES = MMA_K; +constexpr int AB_TMA_SUBS = 4; +static_assert(AB_LOAD_BYTES == 96, + "K=192 packed-FP4 stage must be 96 bytes wide"); +constexpr int TB_SIZE = 224; // 7 warps (blog: the warp table) +constexpr int N_EPI_WARPS = 4; // warps 0-3 + +constexpr int SF_TMEM_COLS = 9 * 4; // 3 SFA + 6 SFB cp tiles, 4 cols each +// Bring-up layout: one accumulator followed immediately by the scale factors. +constexpr int D_STRIDE = BLOCK_N; +constexpr int D_DB = 1; +constexpr int N_D_COLS = BLOCK_N; +constexpr int SF_TMEM_BASE = BLOCK_N; +constexpr int N_ALLOC_COLS = 512; +static_assert(N_D_COLS <= SF_TMEM_BASE && + SF_TMEM_BASE + SF_TMEM_COLS <= N_ALLOC_COLS, + "accumulators + scale factors exceed the 512 TMEM columns"); +constexpr int EPI_CHUNK = 32; // cols per regular drain + +// ---- the shared-memory map (blog: the feed) ----------------------------------- +// Every offset is a compile-time expression. +constexpr int CD_HDR = 1024; // mbarrier header pad +constexpr int CD_AB_SLOTS = 6; +constexpr int CD_SF_SLOTS = 7; +constexpr int CD_AB_STRIDE = BLOCK_M * WIN_BYTES; // 16384 per A or B slot +constexpr int CD_SF_TILE = 512; // one (32 lanes x 16B) x4 tile +constexpr int CD_SFA_TILES = 3; // per SFA slot (M = 1 block) +constexpr int CD_SFB_TILES = 6; // per SFB slot (N = 2 blocks) +constexpr int CD_SFA_STRIDE = CD_SFA_TILES * CD_SF_TILE; // 1536 +constexpr int CD_SFB_STRIDE = CD_SFB_TILES * CD_SF_TILE; // 3072 + +// Expected transfer bytes are derived, never tuned: a wrong-by-one +// expectation doesn't crash, it deadlocks (blog: the feed). +constexpr uint32_t AB_TX = uint32_t(CTA_GROUP) * + uint32_t(BLOCK_M + BLOCK_N_PEER) * uint32_t(AB_LOAD_BYTES); +constexpr uint32_t SF_TX = uint32_t(CD_SFA_STRIDE + CD_SFB_STRIDE) * + uint32_t(CTA_GROUP); // 9216 +static_assert(AB_TX == 49152u && SF_TX == 9216u, + "expect_tx derivation drifted"); + +// The map IS the struct: member sizes and alignment produce every byte +// offset; the raw-address constants the mainloop uses read off it with +// offsetof, and the static_asserts pin the published values. +struct SmemCD { + // mbarrier header: four ring regions in handshake order, then the + // singleton accumulator + teardown barriers (alignas pads it to 1024). + struct alignas(1024) Header { + uint64_t ab_ready[CD_AB_SLOTS]; // TMA -> MMA + uint64_t ab_free [CD_AB_SLOTS]; // MMA -> TMA + uint64_t sf_ready[CD_SF_SLOTS]; // TMA -> MMA + uint64_t sf_free [CD_SF_SLOTS]; // MMA -> TMA + uint64_t acc_ready; // MMA -> epi + uint64_t acc_free; // epi -> MMA + uint64_t dealloc; // 2-CTA teardown + uint32_t tmem_addr; // TMEM base addr storage + } hdr; + uint8_t a [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t b [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t sfa[CD_SF_SLOTS][CD_SFA_STRIDE]; + // SFB starts 1024-aligned (the cp source descriptor drops address bits 4-9). + alignas(1024) uint8_t sfb[CD_SF_SLOTS][CD_SFB_STRIDE]; +}; +static_assert(sizeof(SmemCD::Header) == CD_HDR, "mbarrier header overflows the pad"); + +constexpr int CD_A_OFF = offsetof(SmemCD, a); +constexpr int CD_B_OFF = offsetof(SmemCD, b); +constexpr int CD_SFA_OFF = offsetof(SmemCD, sfa); +constexpr int CD_SFB_OFF = offsetof(SmemCD, sfb); +constexpr int CD_TOTAL = sizeof(SmemCD); +static_assert(CD_TOTAL <= 232448, "smem map exceeds the opt-in budget"); +static_assert(CD_AB_SLOTS != 6 || + (CD_B_OFF == 99328 && CD_SFA_OFF == 197632 && + CD_SFB_OFF == 208896 && CD_TOTAL == 230400), + "the 6-slot byte map drifted from its published values"); + +// The mainloop computes raw `.shared` byte addresses from these. +constexpr int MB_AB_READY = offsetof(SmemCD, hdr.ab_ready); +constexpr int MB_AB_FREE = offsetof(SmemCD, hdr.ab_free); +constexpr int MB_SF_READY = offsetof(SmemCD, hdr.sf_ready); +constexpr int MB_SF_FREE = offsetof(SmemCD, hdr.sf_free); +constexpr int MB_ACC_READY = offsetof(SmemCD, hdr.acc_ready); +constexpr int MB_ACC_FREE = offsetof(SmemCD, hdr.acc_free); + + + +} // namespace nvfp4 + +// ============================================================================ +// The consumer mainloop from the blog's pipeline section. One warp — warp 4 +// of the even CTA — +// issues everything: the tcgen05.cp copies that stage scale factors from the +// smem ring into TMEM, immediately followed by the tcgen05.mma stream, in the +// SAME warp. A same-warp tcgen05 stream is one ordered async pipeline, so the +// write-after-read hazard on the single reused scale-factor TMEM region needs +// no fence and no cross-warp barrier. +// +// The baseline stages four K=192 slices. Each slot supplies two K=96 MMAs at +// packed-byte offsets 0 and 48, both inside the same swizzle window. +// ============================================================================ +namespace nvfp4 { + + +// Ring/parity state carried across tiles by the MMA warp. +struct MainloopState { + uint32_t d_parity; // which TMEM accumulator buffer the next tile writes + uint32_t ab_phase, ab_slot; // A/B ring cursor + uint32_t sf_phase, sf_slot; // scale-factor ring cursor +}; + + +// Relaxed blocking wait (no .acquire): the tcgen05 pipeline's own ordering +// covers these ring waits, and the .acquire form perturbs the scoreboard +// cadence on the issuing warp. Waits sit AT the consuming site — the carried +// one-stage-ahead try_wait peeks that used to gate them measured as pure +// overhead and were removed. +static __device__ __forceinline__ void mbar_wait_relaxed( + uint32_t bar_smem, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(bar_smem), "r"(parity) : "memory"); +} + +// Gated form for the @216 buffer-free wait, which fires once per tile: the +// gate branches over the spin for every group past the first. +static __device__ __forceinline__ void mbar_wait_gated_relaxed( + uint32_t bar_smem, uint32_t parity, bool peek) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p bra DONE_%=;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t" + "DONE_%=:\n\t}\n" + :: "r"(bar_smem), "r"(parity), "r"((uint32_t)peek) : "memory"); +} + + +// One scale-factor staging copy, smem -> TMEM. Single-lane semantics behind +// the converged elect predicate: tcgen05.cp is a single-thread instruction — +// issued bare from a converged warp you get 32 initiations and garbage. +static __device__ __forceinline__ void cp_sf( + uint32_t el, uint32_t dst, uint32_t lo, uint32_t hi) { + const uint64_t desc = (uint64_t(hi) << 32) | uint64_t(lo); + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;\n\t}\n" + :: "r"(dst), "l"(desc), "r"(el) : "memory"); +} + +// Ring-slot "consumed" arrive, multicast to both CTAs of the pair. +// arrive::one must fire exactly once — hence the elect predicate. +static __device__ __forceinline__ void commit_consumed( + uint32_t el, uint32_t smbase, uint32_t mb_off, uint16_t mask) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t}\n" + :: "r"(smbase + mb_off), "h"(mask), "r"(el) : "memory"); +} + +// One K=96 NVFP4 MMA. The asm token says scale_vec::4X regardless — K=96 is +// selected by idesc bit 31 (blog: the MMA warp). acc_pred = 0 overwrites the +// accumulator (the tile's first step), nonzero accumulates. +static __device__ __forceinline__ void mma_cell( + uint32_t el, uint32_t d, + uint32_t desc_a_lo, uint32_t desc_a_hi, + uint32_t desc_b_lo, uint32_t desc_b_hi, + uint32_t idesc, uint32_t sf_a, uint32_t sf_b, uint32_t acc_pred) { + const uint64_t desc_a = (uint64_t(desc_a_hi) << 32) | uint64_t(desc_a_lo); + const uint64_t desc_b = (uint64_t(desc_b_hi) << 32) | uint64_t(desc_b_lo); + asm volatile( + "{\n\t.reg .pred p, acc;\n\t" + "setp.ne.u32 p, %7, 0;\n\t" + "setp.ne.b32 acc, %4, 0;\n\t" + "@p tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X" + " [%0], %1, %2, %3, [%5], [%6], acc;\n\t}\n" + :: "r"(d), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(acc_pred), "r"(sf_a), "r"(sf_b), "r"(el) : "memory"); +} + +// One SF ring slot -> TMEM: 9 tcgen05.cp tiles (3 SFA + 6 SFB). Source +// address advances are in 16-byte descriptor units. The SFB tile order +// 0,3,1,4,2,5 interleaves the two 128-wide N-blocks. +static __device__ __forceinline__ void cp_sf_buffer( + uint32_t el, uint32_t slot, + uint32_t sfa0, uint32_t sfa1, uint32_t sfa2, + uint32_t sfb0, uint32_t sfb1, uint32_t sfb2, + uint32_t sfb3, uint32_t sfb4, uint32_t sfb5, + uint32_t sfa_hi, uint32_t sfa_lo, uint32_t sfb_hi, uint32_t sfb_lo) { + constexpr uint32_t SFA_SLOT_U = CD_SFA_STRIDE / 16u; // 96 + constexpr uint32_t SFB_SLOT_U = CD_SFB_STRIDE / 16u; // 192 + constexpr uint32_t TILE_U = CD_SF_TILE / 16u; // 32 + const uint32_t a = slot * SFA_SLOT_U + sfa_lo; + cp_sf(el, sfa0, a, sfa_hi); + cp_sf(el, sfa1, a + 1u*TILE_U, sfa_hi); + cp_sf(el, sfa2, a + 2u*TILE_U, sfa_hi); + const uint32_t b = slot * SFB_SLOT_U + sfb_lo; + cp_sf(el, sfb0, b + 0u*TILE_U, sfb_hi); + cp_sf(el, sfb1, b + 3u*TILE_U, sfb_hi); + cp_sf(el, sfb2, b + 1u*TILE_U, sfb_hi); + cp_sf(el, sfb3, b + 4u*TILE_U, sfb_hi); + cp_sf(el, sfb4, b + 2u*TILE_U, sfb_hi); + cp_sf(el, sfb5, b + 5u*TILE_U, sfb_hi); +} + +// Operand descriptor LOW word for A/B ring slot `slot`, K-step `k_atom` +// (0..7 inside the 768-element group; each step is 48 bytes = +3 in the +// 16-byte address field). +static __device__ __forceinline__ uint32_t ab_desc_lo( + uint32_t smbase, uint32_t ab_base_off, uint32_t slot, uint32_t k_atom) { + const uint64_t d = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + slot * CD_AB_STRIDE, int(k_atom) * 16); + return uint32_t(d & 0xFFFFFFFFu); +} + + +// One output tile: `full_groups` complete 768-element K-groups. Coordinates +// past logical K are zero-filled by TMA from the unchanged canonical inputs. +static __device__ __forceinline__ MainloopState mainloop_tile( + uint32_t taddr, uint32_t smbase, MainloopState in, uint16_t commit_mask, + uint32_t full_groups + ) { + constexpr uint32_t AB_RING_SLOTS = uint32_t(CD_AB_SLOTS); + constexpr uint32_t SF_RING_SLOTS = uint32_t(CD_SF_SLOTS); + constexpr uint32_t AB_FREE_DELTA = uint32_t(MB_AB_FREE - MB_AB_READY); + constexpr uint32_t SF_READY_OFF = uint32_t(MB_SF_READY); + constexpr uint32_t SF_FREE_DELTA = uint32_t(MB_SF_FREE - MB_SF_READY); + constexpr uint32_t ACC_FREE_OFF = uint32_t(MB_ACC_FREE); + constexpr uint32_t A_OFF = uint32_t(CD_A_OFF); + constexpr uint32_t B_OFF = uint32_t(CD_B_OFF); + + uint32_t d_parity = in.d_parity; + uint32_t ab_phase = in.ab_phase, ab_slot = in.ab_slot; + uint32_t sf_phase = in.sf_phase, sf_slot = in.sf_slot; + + // The caller has already selected exactly one lane for the whole MMA role. + // Keep the helper predicate constant throughout the tile. + const uint32_t el = 1u; + + // Fixed TMEM targets for the single reused scale-factor region: 9 cp + // tiles of 4 columns each, starting past the accumulator region. + constexpr uint32_t SF = uint32_t(SF_TMEM_BASE); + const uint32_t sfa_t0 = taddr + SF + 0; // SFA word 0 (MMA reads here) + const uint32_t sfa_t1 = taddr + SF + 4; + const uint32_t sfa_t2 = taddr + SF + 8; + const uint32_t sfb_t0 = taddr + SF + 12; // SFB word 0 + const uint32_t sfb_t1 = taddr + SF + 16; + const uint32_t sfb_t2 = taddr + SF + 20; + const uint32_t sfb_t3 = taddr + SF + 24; + const uint32_t sfb_t4 = taddr + SF + 28; + const uint32_t sfb_t5 = taddr + SF + 32; + + // Scale-factor cp source descriptors (low word carries the per-slot + // offset; high word is constant). + const uint64_t sfb_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFB_OFF)); + const uint32_t sfb_src_lo = uint32_t(sfb_src & 0xFFFFFFFFu); + const uint32_t sfb_src_hi = uint32_t(sfb_src >> 32); + const uint64_t sfa_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFA_OFF)); + const uint32_t sfa_src_lo = uint32_t(sfa_src & 0xFFFFFFFFu); + const uint32_t sfa_src_hi = uint32_t(sfa_src >> 32); + + // Operand descriptor HIGH words (low words come from ab_desc_lo per step). + const uint32_t a_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + A_OFF) >> 32); + const uint32_t b_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + B_OFF) >> 32); + + // Each K=96 step needs 6 scale bytes per row, split across two staged + // 32-bit words. Odd-issue steps read word 1: bit 31 of the SF address + // selects the with-id encoding, and the id bits fold into the idesc. + constexpr uint32_t SF_ADDR_ID = 0x80000000u; + const uint32_t sfa_w1 = (taddr + SF + 4) | SF_ADDR_ID; + const uint32_t sfb_w1 = (taddr + SF + 20) | SF_ADDR_ID; + const uint32_t w1_id_bits = + ((sfa_w1 >> 1) & 0x60000000u) | ((sfb_w1 >> 26) & 48u); + + // The tile's first MMA overwrites the accumulator buffer this tile owns; + // the @216 "buffer free" wait runs once, before the first group. + const uint32_t acc_wait_parity = d_parity; + d_parity ^= 1u; + // Single buffer: every tile writes buf0; d_parity still toggles because + // it carries the acc_free handshake parity. + const uint32_t d_tmem = taddr; + + // Instruction descriptor: kind::mxf4nvf4, M=256, N=256, E2M1 x E2M1, + // UE4M3 block-16, K-major, k_dim=1 (K=96). Base value 0x90400480. + constexpr uint32_t IDESC = ptx::mma_inst_desc_mxf4nvf4_block16( + 256u, 256u, 0u, 0u, ptx::Major::K, ptx::Major::K, + false, false, /*k_dim=*/true, ptx::ScaleFormat::E4M3); + static_assert(IDESC == 0x90400480u, "K=96 NVFP4 idesc drifted"); + const uint32_t w0_id_bits = + ((sfa_t0 >> 1) & 0x60000000u) | ((sfb_t0 >> 26) & 48u); + const uint32_t idesc_w0 = w0_id_bits | IDESC; + const uint32_t idesc_w1 = w1_id_bits | IDESC; + + + uint32_t acc_pred = 0; // first MMA of the tile overwrites + bool need_acc_wait = true; // @216 wait fires once per tile + + if (full_groups >= 1) { + const uint32_t num_stages = full_groups * 4u; + uint32_t stage = 0; + do { + // One K=192 scale slot and one K=192 A/B slot feed two K=96 + // MMAs. Return each SMEM source through an ordered commit as soon + // as its final asynchronous reader has been issued. + const uint32_t sf_bar = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar, sf_phase); + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, + sfb_t0, sfb_t1, sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar + SF_FREE_DELTA) - smbase, + commit_mask); + + const uint32_t ab_bar = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar, ab_phase); + mbar_wait_gated_relaxed(smbase + ACC_FREE_OFF, + acc_wait_parity, !need_acc_wait); + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); + acc_pred = 0xFFFFFFFFu; + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); + commit_consumed(el, smbase, (ab_bar + AB_FREE_DELTA) - smbase, + commit_mask); + need_acc_wait = false; + + const uint32_t next_sf = sf_slot + 1; + const bool wrap_sf = (next_sf == SF_RING_SLOTS); + sf_slot = wrap_sf ? 0u : next_sf; + sf_phase ^= wrap_sf ? 1u : 0u; + + const uint32_t next_ab = ab_slot + 1; + const bool wrap_ab = (next_ab == AB_RING_SLOTS); + ab_slot = wrap_ab ? 0u : next_ab; + ab_phase ^= wrap_ab ? 1u : 0u; + + ++stage; + } while (stage != num_stages); + } + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; +} + +} // namespace nvfp4 + +// ============================================================================ +// The kernel. 224 threads = 7 warps, four jobs (blog: the warp table): +// warps 0-3 epilogue: drain TMEM, convert to FP16, store straight to global +// warp 4 the MMA warp (even CTA of each pair only): SF copies + MMAs +// warp 5 TMA producer: A and B tiles into the 6-slot ring +// warp 6 TMA producer: scale factors into the 7-slot ring +// ============================================================================ +namespace nvfp4 { + +// Baseline launch: plain launch bounds and a grid expressed in CTA units. +#define NVFP4_BLOCK_ATTR +constexpr bool GRID_IN_CLUSTERS = false; + +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void nvfp4_gemm_kernel(const __grid_constant__ CUtensorMap A_tmap, + const __grid_constant__ CUtensorMap B_tmap, + const __grid_constant__ CUtensorMap SFA_tmap, + const __grid_constant__ CUtensorMap SFB_tmap, + __half* C_ptr, int M, int N, int K + ) { + extern __shared__ __align__(1024) char smem_buf[]; + SmemCD& smem = *reinterpret_cast(smem_buf); + + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const uint32_t crank = ptx::cluster_rank(); + const uint32_t m_pair = crank & 1u; // 0 = the pair's even CTA + + const int cluster_grid_m = (M + CLUSTER_M - 1) / CLUSTER_M; + const int grid_n = (N + BLOCK_N - 1) / BLOCK_N; + const int total_tiles = cluster_grid_m * grid_n; + + // Baseline rungs execute a whole K=768 group at the edge. Tensor-map OOB + // fill supplies zeros past logical K without changing the input buffers. + const int full_groups = (K + 767) / 768; + const int num_groups = full_groups; + + // Persistent cluster index. The device always sees the grid in CTA units + // (2, 1, clusters) regardless of the host-side launch-unit regime. + const bool grid_z = (gridDim.z > 1); + const int num_clusters = grid_z ? int(gridDim.z) : int(gridDim.x) / 2; + const int cluster_id = grid_z ? int(blockIdx.z) : int(blockIdx.x) / 2; + + + // Which output tile does persistent work item `t` process? + auto tile_of = [&](int t) -> int { + return t; // natural raster: M fast, N slow + }; + + auto producer_tile_of = [&](int t) -> int { + return tile_of(t); + }; + + // ---- mbarrier init (warp 0, lane 0) ---- + if (warp_id == 0 && ptx::lane_id() == 0) { + // Arrival counts are derived from the topology, never tuned: + // "ready" = one producer expect_tx per slot (bit-24 routes both + // peers' bytes to the even CTA's barrier); "free" = one consumer + // commit per slot reuse. + #pragma unroll + for (int s = 0; s < CD_AB_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.ab_ready[s], 1); + ptx::mbar_init(&smem.hdr.ab_free[s], 1); + } + #pragma unroll + for (int s = 0; s < CD_SF_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.sf_ready[s], 1); + ptx::mbar_init(&smem.hdr.sf_free[s], 1); + } + ptx::mbar_init(&smem.hdr.acc_ready, 1); + ptx::mbar_init(&smem.hdr.acc_free, N_EPI_WARPS * CTA_GROUP); // 4 x 2 = 8 + ptx::mbar_init(&smem.hdr.dealloc, 32); // one warp's lanes, cross-peer + ptx::fence_mbarrier_init_release_cluster(); + } + ptx::cluster_sync_rel_acq(); + + // TMEM allocation: epilogue warp 0 allocates for the pair; warps 0-4 + // rendezvous before reading the granted base. The TMA producer warps + // never touch TMEM and proceed straight to their loops. + uint32_t taddr = 0; + if (warp_id <= 4) { + if (warp_id == 0) + ptx::tcgen05_alloc_2sm(ptx::to_shared(&smem.hdr.tmem_addr), + N_ALLOC_COLS); + ptx::bar_sync(2, 160); + taddr = smem.hdr.tmem_addr; + } + + const int wg = warp_id >> 2; // 0 = epilogue warps, 1 = producers + MMA + if (wg == 1) { + const int pw = warp_id - 4; + if (pw == 1 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 5: lane 0 owns the complete A/B TMA producer loop. + // ================================================================ + // free_phase starts at 1: a fresh mbarrier is phase 0, so the + // first ring's worth of "free" waits falls through. + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + // The pair splits the 256x256 tile: even CTA holds A rows + // [0,128) and B cols [0,128); odd CTA the other halves. + const int m_blk = m_row * CTA_GROUP + int(m_pair); + const int n_blk = n_grp * CTA_GROUP + int(m_pair); + #pragma unroll 1 + for (int g = 0; g < num_groups; ++g) { + // Four 96-byte loads cover the packed K=768 group. + #pragma unroll + for (int sub = 0; sub < AB_TMA_SUBS; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.ab_free[slot], free_phase); + const uint32_t a_s = + base + CD_A_OFF + uint32_t(slot) * CD_AB_STRIDE; + const uint32_t b_s = + base + CD_B_OFF + uint32_t(slot) * CD_AB_STRIDE; + uint64_t* bar = &smem.hdr.ab_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, AB_TX); + // A rides the multicast form with a self mask so the + // completion routes to the even peer's barrier. + const uint16_t mask_a = uint16_t(1u << m_pair); + // Read the same row-major FP4 allocation as cuBLAS. + // K192 uses four 96-byte boxes; K256 uses three + // 128-byte boxes. The tensor map zero-fills past K. + const CUtensorMap* a_map = &A_tmap; + const CUtensorMap* b_map = &B_tmap; + const int x = g * 384 + sub * AB_LOAD_BYTES; + const int y_a = m_blk * BLOCK_M; + const int y_b = n_blk * BLOCK_N_PEER; + const int z_a = 0; + const int z_b = 0; + ptx::cp_async_bulk_tensor_3d_load_multicast( + a_s, a_map, x, y_a, z_a, + bar, mask_a + ); + ptx::cp_async_bulk_tensor_3d_load_2sm_bit24( + b_s, b_map, x, y_b, z_b, + bar + ); + if (++slot == CD_AB_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + // Drain: wait for the last-filled slot to be consumed, then hand + // it one clean-shutdown expect_tx. + { + const int t_slot = (slot == 0) ? CD_AB_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.ab_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.ab_ready[t_slot], AB_TX); + } + } else if (pw == 2 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 6: scale-factor TMA producer. SFA is per-peer (each CTA's + // own M half); SFB is the one genuinely multicast operand — the + // cta_group::2 copy reads BOTH peers' smem as one logical + // operand, so both copies must be byte-identical (blog: the scale-factor section). + // ================================================================ + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const int m_blk = m_row * CTA_GROUP + int(m_pair); + for (int g = 0; g < num_groups; ++g) { + // 4 SF ring slots per 768-K group (2 K=96 steps each). + #pragma unroll + for (int sub = 0; sub < 4; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.sf_free[slot], free_phase); + const uint32_t sfa_s = + base + CD_SFA_OFF + uint32_t(slot) * CD_SFA_STRIDE; + const uint32_t sfb_s = + base + CD_SFB_OFF + uint32_t(slot) * CD_SFB_STRIDE; + uint64_t* bar = &smem.hdr.sf_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, SF_TX); + // VEC16 is [outer-block][K-group-of-4][128 rows x 4]. + // Three adjacent groups are the 128 x 12 scale slot. + const int sf_group = (g * 4 + sub) * 3; + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfa_s, &SFA_tmap, 0, 0, sf_group, m_blk, bar, + uint16_t(1u << m_pair) + ); + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfb_s + m_pair * uint32_t(CD_SFB_STRIDE / 2), + &SFB_tmap, 0, 0, sf_group, + n_grp * CTA_GROUP + int(m_pair), bar, + uint16_t(0x3) + ); + if (++slot == CD_SF_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + { + const int t_slot = (slot == 0) ? CD_SF_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.sf_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.sf_ready[t_slot], SF_TX); + } + } else if (pw == 0 && m_pair == 0 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 4: the MMA warp — issued once per cluster, on the even + // CTA. The odd CTA's warp 4 does nothing all kernel (blog: the warp table). + // The mainloop needs no tile coordinates: operands arrive through + // the rings; this warp only needs its share of the trip count. + // ================================================================ + const uint32_t smbase = ptx::to_shared(&smem); + const uint16_t commit_mask = 0x3; // both CTAs of this pair + MainloopState st{1u, 0u, 0u, 0u, 0u}; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + st = mainloop_tile(taddr, smbase, st, commit_mask, + uint32_t(full_groups) + ); + // The outer role branch already leaves exactly one issuer. + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;" + :: "r"(smbase + uint32_t(MB_ACC_READY)), + "h"(uint16_t(0x3)) : "memory"); + } + // Drain the final tile's buffer-free arrive before exiting. + mbar_wait_relaxed(smbase + uint32_t(MB_ACC_FREE), st.d_parity); + } + } else { + // ==================================================================== + // Warps 0-3: the epilogue (blog: the epilogue). Each warp owns a 32-row band + // and drains a tile as eight 32-column tcgen05.ld loads, converts + // FP32->FP16 in packed pairs, and stores straight to global — no + // shared-memory bounce. + // ==================================================================== + const int epi_warp = warp_id; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + const uint32_t acc_free_rank = crank & ~1u; // the pair's even CTA + int d_db = 0; + int acc_ready_phase = 0; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const uint32_t tmem_base = taddr + uint32_t(d_db) * D_STRIDE; + mbar_wait_relaxed(ptx::to_shared(&smem.hdr.acc_ready), + uint32_t(acc_ready_phase)); + acc_ready_phase ^= 1; + + const int off_m = (m_row * CTA_GROUP + int(m_pair)) * BLOCK_M; + const int off_n = n_grp * BLOCK_N; + const int row = off_m + epi_warp * 32 + lane_id; + const bool row_in = (row < M); + constexpr int EPI_BANDS = BLOCK_N / EPI_CHUNK; // 8 + #pragma unroll 1 + for (int k = 0; k < EPI_BANDS; ++k) { + const int band = k; + const uint32_t taddr_n = + tmem_base + uint32_t(band) * 32 + taddr_lane; + uint32_t w[32]; + ptx::tcgen05_ld_32x32b_x32(taddr_n, w); + // One buffer cannot be reused until all eight bands are out. + if (k == EPI_BANDS - 1) { + ptx::tcgen05_wait_ld_then_arrive_release( + &smem.hdr.acc_free, acc_free_rank, + /*do_arrive=*/ptx::lane_id() == 0); + } + if (!row_in) continue; + const int col_base = off_n + band * 32; + uint32_t h[16]; + #pragma unroll + for (int j = 0; j < 16; ++j) { + h[j] = ptx::cvt_pack_f16x2(__int_as_float(w[2 * j + 1]), + __int_as_float(w[2 * j])); + } + const size_t row_base = size_t(row) * size_t(N); + const bool aligned32 = ((row_base & 15u) == 0u) && + ((uintptr_t(C_ptr) & 31u) == 0u); + if (aligned32 && (col_base & 15) == 0) { + // Precondition here: row_base is 16-element (32-byte) + // aligned — the aligned32 guard owns it. Restating it to + // ptxas via __builtin_assume compiled this file + // SASS-identical (and measured perf-null on the + // production kernel), so the guard stands alone. + #pragma unroll + for (int q = 0; q < 2; ++q) { + const int c0 = col_base + q * 16; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 16 <= N) { + uint4 lo{h[q*8+0], h[q*8+1], h[q*8+2], h[q*8+3]}; + uint4 hi{h[q*8+4], h[q*8+5], h[q*8+6], h[q*8+7]}; + // Bring-up form: twice the store requests, + // twice the L2 store traffic (the ladder's + // tensor-activity 82% -> 42% row). + ptx::st_global_b128( + reinterpret_cast(cdst), lo); + ptx::st_global_b128( + reinterpret_cast(cdst) + 1, hi); + } else { + // N tail: 128-bit then scalar. + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int cq = c0 + p * 8; + if (cq + 8 <= N) { + uint4 v{h[q*8+p*4+0], h[q*8+p*4+1], + h[q*8+p*4+2], h[q*8+p*4+3]}; + ptx::st_global_b128( + reinterpret_cast(cdst + p * 8), v); + } else { + const uint16_t* hb = + reinterpret_cast( + &h[q * 8 + p * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (cq + c < N) cdst[p * 8 + c] = hb[c]; + } + } + } + } + } else { + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int c0 = col_base + q * 8; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 8 <= N) { + uint4 v{h[q*4+0], h[q*4+1], h[q*4+2], h[q*4+3]}; + ptx::st_global_b128(reinterpret_cast(cdst), v); + } else { + const uint16_t* hb = + reinterpret_cast(&h[q * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (c0 + c < N) cdst[c] = hb[c]; + } + } + } + } + if (++d_db == D_DB) d_db = 0; + } + } + + // Teardown: only the epilogue warps rendezvous per CTA; warp 0 then runs + // the cross-peer dealloc handshake. Producer and MMA warps exit freely — + // their drains above already order their work. + if (wg == 0) { + ptx::bar_sync(3, 128); + if (warp_id == 0) { + ptx::tcgen05_relinquish_2sm(); + const uint32_t peer_rank = crank ^ 1u; + // All 32 lanes arrive the PEER's count-32 barrier = one flip. + ptx::mbar_arrive_cluster_release(&smem.hdr.dealloc, peer_rank); + ptx::mbar_wait_parity(&smem.hdr.dealloc, 0); + ptx::tcgen05_dealloc_2sm(taddr, N_ALLOC_COLS); + } + } +} + + +// Keep the host launch convention paired with the kernel's block attribute. +static inline dim3 launch_grid(int clusters) { + return GRID_IN_CLUSTERS ? dim3(1, 1, unsigned(clusters)) + : dim3(2, 1, unsigned(clusters)); +} + +} // namespace nvfp4 + +// ============================================================================ +// Host-side TMA descriptor support. The shared input and benchmark code lives +// in main.cu. +// ============================================================================ +namespace tmap { + +// Thin wrappers over cuTensorMapEncodeTiled. Dim 0 is the innermost +// (stride-1) axis; strides for dim 1+ are in bytes. +inline CUtensorMap encode_tiled_3d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_depth, uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t depth_stride_bytes, + uint32_t box_depth, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[3] = { global_cols, global_rows, global_depth }; + cuuint64_t global_strides[2] = { row_stride_bytes, depth_stride_bytes }; + cuuint32_t box_dim[3] = { box_cols, box_rows, box_depth }; + cuuint32_t element_strides[3] = { 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 3, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +inline CUtensorMap encode_tiled_4d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_slots, uint64_t global_tiles, + uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t tile_stride_bytes, + uint64_t slot_stride_bytes, + uint32_t box_tiles, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[4] = { global_cols, global_rows, global_tiles, + global_slots }; + cuuint64_t global_strides[3] = { row_stride_bytes, tile_stride_bytes, + slot_stride_bytes }; + cuuint32_t box_dim[4] = { box_cols, box_rows, box_tiles, 1u }; + cuuint32_t element_strides[4] = { 1, 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 4, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +} // namespace tmap + +namespace host { + +using namespace nvfp4; + +// ---- implementation-specific tensor maps --------------------------------- + +inline CUtensorMap make_ab_tmap(void* ptr, int rows, int K, + int row_stride) { + return tmap::encode_tiled_3d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*depth=*/1, uint64_t(rows), /*cols=*/uint64_t(K / 2), + /*row_stride=*/uint64_t(row_stride), + /*depth_stride=*/uint64_t(rows) * uint64_t(row_stride), + /*box=*/1, BLOCK_M, AB_LOAD_BYTES, CU_TENSOR_MAP_SWIZZLE_128B); +} + +inline CUtensorMap make_sf_tmap(uint8_t* ptr, int outer, int K) { + const uint64_t sf_inner = (uint64_t(K) + 63) / 64 * 4; + const uint64_t outer_blocks = (uint64_t(outer) + 127) / 128; + return tmap::encode_tiled_4d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*slots=*/outer_blocks, /*tiles=*/sf_inner / 4, + /*rows=*/4, /*cols=*/128, + /*row_stride=*/128, /*tile_stride=*/512, + /*slot_stride=*/sf_inner * 128, + /*box_tiles=*/3, /*box_rows=*/4, /*box_cols=*/128, + CU_TENSOR_MAP_SWIZZLE_NONE); +} +} // namespace host diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/gemm2.cuh b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm2.cuh new file mode 100644 index 000000000..557f3a0d9 --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm2.cuh @@ -0,0 +1,1254 @@ +#pragma once +// r2: stage K256 A/B windows and release each after its last reader. +// This is a complete, directly editable ladder snapshot. +// It intentionally duplicates the implementation in the other gemm*.cuh files. + +// Conditional branches are already materialized; edit this implementation directly. +#define NVFP4_BUILD_NAME "r2" +// Read-only metadata for main.cu; these do not select code in this header. +#define NVFP4_HAS_K64_TAIL 0 +#define NVFP4_HAS_SCHEDULE 0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(expr) do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \ + cudaGetErrorName(_e), __FILE__, __LINE__, \ + cudaGetErrorString(_e)); \ + std::abort(); \ + } \ +} while (0) + +#define CU_CHECK(expr) do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + std::fprintf(stderr, "CU error at %s:%d: %s\n", \ + __FILE__, __LINE__, _s ? _s : "?"); \ + std::abort(); \ + } \ +} while (0) + + +// ============================================================================ +// PTX helpers. Thin zero-cost wrappers over the exact instruction forms the +// kernel needs — nothing generic, nothing speculative. +// ============================================================================ +namespace ptx { + +// Generic pointer -> 32-bit `.shared` window byte address. +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Local `.shared` address -> the same offset in CTA `cta_rank`'s smem +// (`.shared::cluster` instructions take this mapped form). +static __device__ __forceinline__ uint32_t mapa_shared_cluster( + uint32_t local_addr, uint32_t cta_rank) { + uint32_t mapped; + asm("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(mapped) : "r"(local_addr), "r"(cta_rank)); + return mapped; +} + +static __device__ __forceinline__ void bar_sync(uint32_t id, uint32_t count) { + asm volatile("barrier.cta.sync %0, %1;" :: "r"(id), "r"(count)); +} + +// Cluster barrier with release/acquire — the publish boundary between the +// mbarrier inits and any cross-CTA use of those mbarriers. +static __device__ __forceinline__ void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// %cluster_ctarank: 0 or 1 in this kernel's (2,1,1) clusters. +static __device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} + +static __device__ __forceinline__ uint32_t lane_id() { + uint32_t r; + asm("mov.u32 %0, %%laneid;" : "=r"(r)); + return r; +} + +// ---- mbarriers ------------------------------------------------------------- + +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + :: "r"(to_shared(bar)), "r"(count)); +} + +// Publish the inits to the async proxy + the peer CTA before any TMA +// completes on them. +static __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// Cross-CTA arrive with release ordering (teardown handshake). +static __device__ __forceinline__ void mbar_arrive_cluster_release( + uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" + :: "r"(mapped)); +} + +// Fused `tcgen05.wait::ld` + "buffer free" arrive, one asm block. The arrive +// must not become visible while TMEM loads are still in flight (it frees the +// accumulator buffer to the next tile's MMA — the bring-up chapter's +// drain overlap); the fused block +// keeps the arrive pinned behind the wait. (An earlier form also manufactured +// a register dependency on a drained value, `dep * 0` folded into the +// address; verified dead on nvcc 13.1 — ptxas const-folds it, SASS +// byte-identical without it — and retired.) +static __device__ __forceinline__ void tcgen05_wait_ld_then_arrive_release( + uint64_t* bar, uint32_t cta_rank, bool do_arrive) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile( + "{\n\t" + ".reg .pred %%pdo;\n\t" + "setp.ne.s32 %%pdo, %1, 0;\n\t" + "tcgen05.wait::ld.sync.aligned;\n\t" + "@%%pdo mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;\n\t" + "}\n\t" + :: "r"(mapped), "r"((int)do_arrive) : "memory"); +} + + +// Blocking parity wait at the consuming site (producer ring/tail waits + +// teardown). The 10 ms suspend hint is the polite form: a full-rate spin +// competes for issue slots and measurably loses; a nanosleep-backoff spin +// measured the same as this hint across five windows, so the simpler form +// ships. The carried one-stage-ahead try_wait peeks that used to gate these +// waits measured as pure overhead and were removed. +static __device__ __forceinline__ void mbar_wait_parity( + uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(to_shared(bar)), "r"(parity)); +} + +// ---- TMA ------------------------------------------------------------------- + +// Producer expect_tx arrive on the cluster-shared ring barrier: both +// peers' transfer bytes account against the even CTA's barrier, so one +// barrier per ring slot covers the whole cluster-wide gather. (Mechanism: +// `& 0xFEFFFFFF` clears bit 24 of the shared address — the hardware's +// peer-CTA bit — re-pointing the arrive at the even CTA's copy.) +static __device__ __forceinline__ void mbar_arrive_expect_tx_cluster( + uint64_t* bar, uint32_t bytes) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes)); +} + +// 3D tiled load, cta_group::2 + multicast (the A feed; the mask is the +// issuing CTA itself in a (2,1) cluster). Completion routes to the even +// CTA's mbarrier via the bit-24 clear. +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z) + : "memory"); +} + +// 3D tiled load, cta_group::2, no multicast (the B feed — each peer loads +// its own 128-wide N-block). +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_2sm_bit24( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst_smem), "l"(tmap), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) + : "memory"); +} + +// 4D tiled load, cta_group::2 + multicast (the scale-factor feed). The +// 512-byte SF tiles of one ring slot form a 4D box; `w` selects the slot. +static __device__ __forceinline__ void cp_async_bulk_tensor_4d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z2, int32_t w, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.4d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z2), "r"(w) + : "memory"); +} + + +// ---- tcgen05 (tensor memory + MMA plumbing) --------------------------------- + +static __device__ __forceinline__ void tcgen05_alloc_2sm( + uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_dealloc_2sm( + uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" + :: "r"(taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_relinquish_2sm() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); +} + +// The widest single-instruction TMEM read: 32 b32 registers per lane +// (blog: the epilogue — a quarter the load/wait count of the narrow forms). +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t taddr, uint32_t* r) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), + "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), + "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), + "=r"(r[16]), "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), + "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), + "=r"(r[24]), "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), + "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr)); +} + +// ---- stores + conversion ----------------------------------------------------- + +static __device__ __forceinline__ uint32_t cvt_pack_f16x2(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// Baseline output store: use the default cache policy. +static __device__ __forceinline__ void st_global_b128(uint4* ptr, uint4 v) { + asm volatile("st.global.v4.b32 [%0], {%1, %2, %3, %4};" + :: "l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)); +} + + + +// ---- MMA descriptors --------------------------------------------------------- + +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class ScaleFormat : uint8_t { E4M3 = 0, E8M0 = 1 }; + +// Source descriptor for the tcgen05.cp that stages scale factors from smem +// into TMEM: unswizzled, SBO=128, LBO=16. The low word later carries the +// per-slot byte offset (in 16-byte units). +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_sf_cp_src_desc_cd(uint32_t sf_smem_addr) { + uint64_t addr_field = uint64_t((sf_smem_addr >> 4) & 0x3FC0u); + return addr_field | 0x400800010000ull; +} + +// Operand descriptor for a K=96 (48-byte) slice inside a 128-byte-swizzled +// smem K-window. OR-constant 0x4010404000000000: version=1, absolute-address +// LBO mode (bit 52), 128B swizzle, SBO=1024. Both the address field and the +// absolute-LBO field encode the window's own 1024-aligned base; the +// intra-window offset advances the address field only. +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k96_sw128(uint32_t window_base_1024aligned, int intra_window_off = 0) { + uint64_t addr_field = uint64_t(((window_base_1024aligned >> 4) & 0x7FC0u) + + uint32_t(intra_window_off >> 4)); + uint64_t lbo_field = uint64_t((window_base_1024aligned << 12) & 0x7FC00000u); + return addr_field | lbo_field | 0x4010404000000000ull; +} + +// Instruction descriptor for tcgen05.mma.kind::mxf4nvf4.block_scale. +// The K extent is selected by BIT 31, not by the asm mnemonic: k_dim=0 runs +// K=64 (0x10400480), k_dim=1 runs K=96 (0x90400480, sm_103a-only) — the +// cheapest quarter of the gap if your kernel is built on the K=64 form +// (blog: the MMA warp). Bit 13 flips which TMEM scale-factor column the MMA reads. +__host__ __device__ static __forceinline__ constexpr uint32_t +mma_inst_desc_mxf4nvf4_block16( + uint32_t M, uint32_t N, + uint32_t tmem_sfa_addr, uint32_t tmem_sfb_addr, + Major a_major, Major b_major, + bool negate_a, bool negate_b, + bool k_dim, ScaleFormat sf) { + constexpr uint32_t MXF4_E2M1 = 1u; + uint32_t d = 0; + d |= ((tmem_sfb_addr & 0xC0000000u) >> 30) << 4; // b_sf_id + d |= MXF4_E2M1 << 7; // a_format = E2M1 + d |= MXF4_E2M1 << 10; // b_format = E2M1 + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; // N in units of 8 + d |= (static_cast(sf) & 0x1u) << 23; // 0 = UE4M3 + d |= ((M >> 4) & 0x1Fu) << 24; // m_dim at bits 27-28 + d |= ((tmem_sfa_addr & 0xC0000000u) >> 30) << 29; // a_sf_id + if (k_dim) d |= 1u << 31; // 1 = K=96 + return d; +} + +} // namespace ptx + +// ============================================================================ +// Kernel geometry from the blog's opening kernel section. One tile shape, +// one cluster shape: +// 256x256 output per 2-CTA cluster (128 rows per CTA), K advanced in +// 96-element MMA steps staged as 768-element K-groups of 8 steps, persistent +// grid of 76 clusters = 152 CTAs = one per SM. +// ============================================================================ +namespace nvfp4 { + +constexpr int CTA_GROUP = 2; // cta_group::2 — the MMA spans both CTAs +constexpr int BLOCK_M = 128; // A rows per CTA +constexpr int BLOCK_N = 256; // output tile N extent (per cluster) +constexpr int BLOCK_N_PEER = BLOCK_N / CTA_GROUP; // 128 — B cols per CTA +constexpr int CLUSTER_M = CTA_GROUP * BLOCK_M; // 256 — tile M extent +[[maybe_unused]] constexpr int MMA_K = 96; // K per tcgen05.mma (idesc bit 31) +constexpr int WIN_BYTES = 128; // one 128B-swizzled smem K-window +constexpr int AB_LOAD_BYTES = WIN_BYTES; +constexpr int AB_TMA_SUBS = 3; +constexpr int TB_SIZE = 224; // 7 warps (blog: the warp table) +constexpr int N_EPI_WARPS = 4; // warps 0-3 + +constexpr int SF_TMEM_COLS = 9 * 4; // 3 SFA + 6 SFB cp tiles, 4 cols each +// Bring-up layout: one accumulator followed immediately by the scale factors. +constexpr int D_STRIDE = BLOCK_N; +constexpr int D_DB = 1; +constexpr int N_D_COLS = BLOCK_N; +constexpr int SF_TMEM_BASE = BLOCK_N; +constexpr int N_ALLOC_COLS = 512; +static_assert(N_D_COLS <= SF_TMEM_BASE && + SF_TMEM_BASE + SF_TMEM_COLS <= N_ALLOC_COLS, + "accumulators + scale factors exceed the 512 TMEM columns"); +constexpr int EPI_CHUNK = 32; // cols per regular drain + +// ---- the shared-memory map (blog: the feed) ----------------------------------- +// Every offset is a compile-time expression. +constexpr int CD_HDR = 1024; // mbarrier header pad +constexpr int CD_AB_SLOTS = 6; +constexpr int CD_SF_SLOTS = 7; +constexpr int CD_AB_STRIDE = BLOCK_M * WIN_BYTES; // 16384 per A or B slot +constexpr int CD_SF_TILE = 512; // one (32 lanes x 16B) x4 tile +constexpr int CD_SFA_TILES = 3; // per SFA slot (M = 1 block) +constexpr int CD_SFB_TILES = 6; // per SFB slot (N = 2 blocks) +constexpr int CD_SFA_STRIDE = CD_SFA_TILES * CD_SF_TILE; // 1536 +constexpr int CD_SFB_STRIDE = CD_SFB_TILES * CD_SF_TILE; // 3072 + +// Expected transfer bytes are derived, never tuned: a wrong-by-one +// expectation doesn't crash, it deadlocks (blog: the feed). +constexpr uint32_t AB_TX = uint32_t(CTA_GROUP) * + uint32_t(BLOCK_M + BLOCK_N_PEER) * uint32_t(AB_LOAD_BYTES); +constexpr uint32_t SF_TX = uint32_t(CD_SFA_STRIDE + CD_SFB_STRIDE) * + uint32_t(CTA_GROUP); // 9216 +static_assert(AB_TX == 65536u && SF_TX == 9216u, + "expect_tx derivation drifted"); + +// The map IS the struct: member sizes and alignment produce every byte +// offset; the raw-address constants the mainloop uses read off it with +// offsetof, and the static_asserts pin the published values. +struct SmemCD { + // mbarrier header: four ring regions in handshake order, then the + // singleton accumulator + teardown barriers (alignas pads it to 1024). + struct alignas(1024) Header { + uint64_t ab_ready[CD_AB_SLOTS]; // TMA -> MMA + uint64_t ab_free [CD_AB_SLOTS]; // MMA -> TMA + uint64_t sf_ready[CD_SF_SLOTS]; // TMA -> MMA + uint64_t sf_free [CD_SF_SLOTS]; // MMA -> TMA + uint64_t acc_ready; // MMA -> epi + uint64_t acc_free; // epi -> MMA + uint64_t dealloc; // 2-CTA teardown + uint32_t tmem_addr; // TMEM base addr storage + } hdr; + uint8_t a [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t b [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t sfa[CD_SF_SLOTS][CD_SFA_STRIDE]; + // SFB starts 1024-aligned (the cp source descriptor drops address bits 4-9). + alignas(1024) uint8_t sfb[CD_SF_SLOTS][CD_SFB_STRIDE]; +}; +static_assert(sizeof(SmemCD::Header) == CD_HDR, "mbarrier header overflows the pad"); + +constexpr int CD_A_OFF = offsetof(SmemCD, a); +constexpr int CD_B_OFF = offsetof(SmemCD, b); +constexpr int CD_SFA_OFF = offsetof(SmemCD, sfa); +constexpr int CD_SFB_OFF = offsetof(SmemCD, sfb); +constexpr int CD_TOTAL = sizeof(SmemCD); +static_assert(CD_TOTAL <= 232448, "smem map exceeds the opt-in budget"); +static_assert(CD_AB_SLOTS != 6 || + (CD_B_OFF == 99328 && CD_SFA_OFF == 197632 && + CD_SFB_OFF == 208896 && CD_TOTAL == 230400), + "the 6-slot byte map drifted from its published values"); + +// The mainloop computes raw `.shared` byte addresses from these. +constexpr int MB_AB_READY = offsetof(SmemCD, hdr.ab_ready); +constexpr int MB_AB_FREE = offsetof(SmemCD, hdr.ab_free); +constexpr int MB_SF_READY = offsetof(SmemCD, hdr.sf_ready); +constexpr int MB_SF_FREE = offsetof(SmemCD, hdr.sf_free); +constexpr int MB_ACC_READY = offsetof(SmemCD, hdr.acc_ready); +constexpr int MB_ACC_FREE = offsetof(SmemCD, hdr.acc_free); + + + +} // namespace nvfp4 + +// ============================================================================ +// The consumer mainloop from the blog's pipeline section. One warp — warp 4 +// of the even CTA — +// issues everything: the tcgen05.cp copies that stage scale factors from the +// smem ring into TMEM, immediately followed by the tcgen05.mma stream, in the +// SAME warp. A same-warp tcgen05 stream is one ordered async pipeline, so the +// write-after-read hazard on the single reused scale-factor TMEM region needs +// no fence and no cross-warp barrier. +// +// Per 768-element K-group, logical K=96 steps 0..7 issue in K order. Their +// starting offsets, expressed as 16-byte descriptor atoms modulo a 128-byte +// swizzle window, are 0,3,6,1,4,7,2,5. Atoms 6 and 7 straddle windows. +// ============================================================================ +namespace nvfp4 { + + +// Ring/parity state carried across tiles by the MMA warp. +struct MainloopState { + uint32_t d_parity; // which TMEM accumulator buffer the next tile writes + uint32_t ab_phase, ab_slot; // A/B ring cursor + uint32_t sf_phase, sf_slot; // scale-factor ring cursor +}; + + +// Relaxed blocking wait (no .acquire): the tcgen05 pipeline's own ordering +// covers these ring waits, and the .acquire form perturbs the scoreboard +// cadence on the issuing warp. Waits sit AT the consuming site — the carried +// one-stage-ahead try_wait peeks that used to gate them measured as pure +// overhead and were removed. +static __device__ __forceinline__ void mbar_wait_relaxed( + uint32_t bar_smem, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(bar_smem), "r"(parity) : "memory"); +} + +// Gated form for the @216 buffer-free wait, which fires once per tile: the +// gate branches over the spin for every group past the first. +static __device__ __forceinline__ void mbar_wait_gated_relaxed( + uint32_t bar_smem, uint32_t parity, bool peek) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p bra DONE_%=;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t" + "DONE_%=:\n\t}\n" + :: "r"(bar_smem), "r"(parity), "r"((uint32_t)peek) : "memory"); +} + + +// One scale-factor staging copy, smem -> TMEM. Single-lane semantics behind +// the converged elect predicate: tcgen05.cp is a single-thread instruction — +// issued bare from a converged warp you get 32 initiations and garbage. +static __device__ __forceinline__ void cp_sf( + uint32_t el, uint32_t dst, uint32_t lo, uint32_t hi) { + const uint64_t desc = (uint64_t(hi) << 32) | uint64_t(lo); + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;\n\t}\n" + :: "r"(dst), "l"(desc), "r"(el) : "memory"); +} + +// Ring-slot "consumed" arrive, multicast to both CTAs of the pair. +// arrive::one must fire exactly once — hence the elect predicate. +static __device__ __forceinline__ void commit_consumed( + uint32_t el, uint32_t smbase, uint32_t mb_off, uint16_t mask) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t}\n" + :: "r"(smbase + mb_off), "h"(mask), "r"(el) : "memory"); +} + +// One K=96 NVFP4 MMA. The asm token says scale_vec::4X regardless — K=96 is +// selected by idesc bit 31 (blog: the MMA warp). acc_pred = 0 overwrites the +// accumulator (the tile's first step), nonzero accumulates. +static __device__ __forceinline__ void mma_cell( + uint32_t el, uint32_t d, + uint32_t desc_a_lo, uint32_t desc_a_hi, + uint32_t desc_b_lo, uint32_t desc_b_hi, + uint32_t idesc, uint32_t sf_a, uint32_t sf_b, uint32_t acc_pred) { + const uint64_t desc_a = (uint64_t(desc_a_hi) << 32) | uint64_t(desc_a_lo); + const uint64_t desc_b = (uint64_t(desc_b_hi) << 32) | uint64_t(desc_b_lo); + asm volatile( + "{\n\t.reg .pred p, acc;\n\t" + "setp.ne.u32 p, %7, 0;\n\t" + "setp.ne.b32 acc, %4, 0;\n\t" + "@p tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X" + " [%0], %1, %2, %3, [%5], [%6], acc;\n\t}\n" + :: "r"(d), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(acc_pred), "r"(sf_a), "r"(sf_b), "r"(el) : "memory"); +} + +// One SF ring slot -> TMEM: 9 tcgen05.cp tiles (3 SFA + 6 SFB). Source +// address advances are in 16-byte descriptor units. The SFB tile order +// 0,3,1,4,2,5 interleaves the two 128-wide N-blocks. +static __device__ __forceinline__ void cp_sf_buffer( + uint32_t el, uint32_t slot, + uint32_t sfa0, uint32_t sfa1, uint32_t sfa2, + uint32_t sfb0, uint32_t sfb1, uint32_t sfb2, + uint32_t sfb3, uint32_t sfb4, uint32_t sfb5, + uint32_t sfa_hi, uint32_t sfa_lo, uint32_t sfb_hi, uint32_t sfb_lo) { + constexpr uint32_t SFA_SLOT_U = CD_SFA_STRIDE / 16u; // 96 + constexpr uint32_t SFB_SLOT_U = CD_SFB_STRIDE / 16u; // 192 + constexpr uint32_t TILE_U = CD_SF_TILE / 16u; // 32 + const uint32_t a = slot * SFA_SLOT_U + sfa_lo; + cp_sf(el, sfa0, a, sfa_hi); + cp_sf(el, sfa1, a + 1u*TILE_U, sfa_hi); + cp_sf(el, sfa2, a + 2u*TILE_U, sfa_hi); + const uint32_t b = slot * SFB_SLOT_U + sfb_lo; + cp_sf(el, sfb0, b + 0u*TILE_U, sfb_hi); + cp_sf(el, sfb1, b + 3u*TILE_U, sfb_hi); + cp_sf(el, sfb2, b + 1u*TILE_U, sfb_hi); + cp_sf(el, sfb3, b + 4u*TILE_U, sfb_hi); + cp_sf(el, sfb4, b + 2u*TILE_U, sfb_hi); + cp_sf(el, sfb5, b + 5u*TILE_U, sfb_hi); +} + +// Operand descriptor LOW word for A/B ring slot `slot`, K-step `k_atom` +// (0..7 inside the 768-element group; each step is 48 bytes = +3 in the +// 16-byte address field). +static __device__ __forceinline__ uint32_t ab_desc_lo( + uint32_t smbase, uint32_t ab_base_off, uint32_t slot, uint32_t k_atom) { + const uint64_t d = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + slot * CD_AB_STRIDE, int(k_atom) * 16); + return uint32_t(d & 0xFFFFFFFFu); +} + +// K-steps 6 and 7 cross a 128-byte window boundary. Their descriptor address +// comes from `addr_slot`, while the absolute-LBO field comes from the next +// slot whose swizzle window contains the second half. +static __device__ __forceinline__ uint32_t ab_desc_lo_straddle( + uint32_t smbase, uint32_t ab_base_off, + uint32_t addr_slot, uint32_t lbo_slot, uint32_t k_atom) { + const uint64_t d_addr = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + addr_slot * CD_AB_STRIDE, int(k_atom) * 16); + const uint64_t d_lbo = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + lbo_slot * CD_AB_STRIDE, 0); + return uint32_t(d_addr & 0x0000FFFFu) | uint32_t(d_lbo & 0xFFFF0000u); +} + +// One output tile: `full_groups` complete 768-element K-groups. Coordinates +// past logical K are zero-filled by TMA from the unchanged canonical inputs. +static __device__ __forceinline__ MainloopState mainloop_tile( + uint32_t taddr, uint32_t smbase, MainloopState in, uint16_t commit_mask, + uint32_t full_groups + ) { + constexpr uint32_t AB_RING_SLOTS = uint32_t(CD_AB_SLOTS); + constexpr uint32_t SF_RING_SLOTS = uint32_t(CD_SF_SLOTS); + constexpr uint32_t AB_FREE_DELTA = uint32_t(MB_AB_FREE - MB_AB_READY); + constexpr uint32_t SF_READY_OFF = uint32_t(MB_SF_READY); + constexpr uint32_t SF_FREE_DELTA = uint32_t(MB_SF_FREE - MB_SF_READY); + constexpr uint32_t ACC_FREE_OFF = uint32_t(MB_ACC_FREE); + constexpr uint32_t A_OFF = uint32_t(CD_A_OFF); + constexpr uint32_t B_OFF = uint32_t(CD_B_OFF); + + uint32_t d_parity = in.d_parity; + uint32_t ab_phase = in.ab_phase, ab_slot = in.ab_slot; + uint32_t sf_phase = in.sf_phase, sf_slot = in.sf_slot; + + // The caller has already selected exactly one lane for the whole MMA role. + // Keep the helper predicate constant throughout the tile. + const uint32_t el = 1u; + + // Fixed TMEM targets for the single reused scale-factor region: 9 cp + // tiles of 4 columns each, starting past the accumulator region. + constexpr uint32_t SF = uint32_t(SF_TMEM_BASE); + const uint32_t sfa_t0 = taddr + SF + 0; // SFA word 0 (MMA reads here) + const uint32_t sfa_t1 = taddr + SF + 4; + const uint32_t sfa_t2 = taddr + SF + 8; + const uint32_t sfb_t0 = taddr + SF + 12; // SFB word 0 + const uint32_t sfb_t1 = taddr + SF + 16; + const uint32_t sfb_t2 = taddr + SF + 20; + const uint32_t sfb_t3 = taddr + SF + 24; + const uint32_t sfb_t4 = taddr + SF + 28; + const uint32_t sfb_t5 = taddr + SF + 32; + + // Scale-factor cp source descriptors (low word carries the per-slot + // offset; high word is constant). + const uint64_t sfb_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFB_OFF)); + const uint32_t sfb_src_lo = uint32_t(sfb_src & 0xFFFFFFFFu); + const uint32_t sfb_src_hi = uint32_t(sfb_src >> 32); + const uint64_t sfa_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFA_OFF)); + const uint32_t sfa_src_lo = uint32_t(sfa_src & 0xFFFFFFFFu); + const uint32_t sfa_src_hi = uint32_t(sfa_src >> 32); + + // Operand descriptor HIGH words (low words come from ab_desc_lo per step). + const uint32_t a_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + A_OFF) >> 32); + const uint32_t b_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + B_OFF) >> 32); + + // Each K=96 step needs 6 scale bytes per row, split across two staged + // 32-bit words. Odd-issue steps read word 1: bit 31 of the SF address + // selects the with-id encoding, and the id bits fold into the idesc. + constexpr uint32_t SF_ADDR_ID = 0x80000000u; + const uint32_t sfa_w1 = (taddr + SF + 4) | SF_ADDR_ID; + const uint32_t sfb_w1 = (taddr + SF + 20) | SF_ADDR_ID; + const uint32_t w1_id_bits = + ((sfa_w1 >> 1) & 0x60000000u) | ((sfb_w1 >> 26) & 48u); + + // The tile's first MMA overwrites the accumulator buffer this tile owns; + // the @216 "buffer free" wait runs once, before the first group. + const uint32_t acc_wait_parity = d_parity; + d_parity ^= 1u; + // Single buffer: every tile writes buf0; d_parity still toggles because + // it carries the acc_free handshake parity. + const uint32_t d_tmem = taddr; + + // Instruction descriptor: kind::mxf4nvf4, M=256, N=256, E2M1 x E2M1, + // UE4M3 block-16, K-major, k_dim=1 (K=96). Base value 0x90400480. + constexpr uint32_t IDESC = ptx::mma_inst_desc_mxf4nvf4_block16( + 256u, 256u, 0u, 0u, ptx::Major::K, ptx::Major::K, + false, false, /*k_dim=*/true, ptx::ScaleFormat::E4M3); + static_assert(IDESC == 0x90400480u, "K=96 NVFP4 idesc drifted"); + const uint32_t w0_id_bits = + ((sfa_t0 >> 1) & 0x60000000u) | ((sfb_t0 >> 26) & 48u); + const uint32_t idesc_w0 = w0_id_bits | IDESC; + const uint32_t idesc_w1 = w1_id_bits | IDESC; + + + uint32_t acc_pred = 0; // first MMA of the tile overwrites + bool need_acc_wait = true; // @216 wait fires once per tile + + if (full_groups >= 1) { + uint32_t g = 0; + do { + // ---- SF batch 0 -> logical steps 0,1 (atoms 0,3) ---- + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // ---- SF batch 1 -> logical steps 2,3 (atoms 6,1) ---- + mbar_wait_relaxed(sf_bar1, sf_phase1); + uint32_t sf_slot2, sf_phase2; + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + // AB0's last reader is MMA2. + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // ---- SF batch 2 -> logical steps 4,5 (atoms 4,7) ---- + mbar_wait_relaxed(sf_bar2, sf_phase2); + uint32_t sf_slot3, sf_phase3; + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; // ring cursor writeback + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // ---- SF batch 3 -> logical steps 6,7 (atoms 2,5) ---- + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; // ring cursor writeback + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + // AB1's last reader is MMA5; this measured placement leaves the + // commit behind the following scale copy. + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + acc_pred = 0xFFFFFFFFu; // accumulate from the second group on + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + need_acc_wait = false; + // AB2's last reader is MMA7. + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, + commit_mask); + ++g; + } while (g != full_groups); + } + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; +} + +} // namespace nvfp4 + +// ============================================================================ +// The kernel. 224 threads = 7 warps, four jobs (blog: the warp table): +// warps 0-3 epilogue: drain TMEM, convert to FP16, store straight to global +// warp 4 the MMA warp (even CTA of each pair only): SF copies + MMAs +// warp 5 TMA producer: A and B tiles into the 6-slot ring +// warp 6 TMA producer: scale factors into the 7-slot ring +// ============================================================================ +namespace nvfp4 { + +// Baseline launch: plain launch bounds and a grid expressed in CTA units. +#define NVFP4_BLOCK_ATTR +constexpr bool GRID_IN_CLUSTERS = false; + +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void nvfp4_gemm_kernel(const __grid_constant__ CUtensorMap A_tmap, + const __grid_constant__ CUtensorMap B_tmap, + const __grid_constant__ CUtensorMap SFA_tmap, + const __grid_constant__ CUtensorMap SFB_tmap, + __half* C_ptr, int M, int N, int K + ) { + extern __shared__ __align__(1024) char smem_buf[]; + SmemCD& smem = *reinterpret_cast(smem_buf); + + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const uint32_t crank = ptx::cluster_rank(); + const uint32_t m_pair = crank & 1u; // 0 = the pair's even CTA + + const int cluster_grid_m = (M + CLUSTER_M - 1) / CLUSTER_M; + const int grid_n = (N + BLOCK_N - 1) / BLOCK_N; + const int total_tiles = cluster_grid_m * grid_n; + + // Baseline rungs execute a whole K=768 group at the edge. Tensor-map OOB + // fill supplies zeros past logical K without changing the input buffers. + const int full_groups = (K + 767) / 768; + const int num_groups = full_groups; + + // Persistent cluster index. The device always sees the grid in CTA units + // (2, 1, clusters) regardless of the host-side launch-unit regime. + const bool grid_z = (gridDim.z > 1); + const int num_clusters = grid_z ? int(gridDim.z) : int(gridDim.x) / 2; + const int cluster_id = grid_z ? int(blockIdx.z) : int(blockIdx.x) / 2; + + + // Which output tile does persistent work item `t` process? + auto tile_of = [&](int t) -> int { + return t; // natural raster: M fast, N slow + }; + + auto producer_tile_of = [&](int t) -> int { + return tile_of(t); + }; + + // ---- mbarrier init (warp 0, lane 0) ---- + if (warp_id == 0 && ptx::lane_id() == 0) { + // Arrival counts are derived from the topology, never tuned: + // "ready" = one producer expect_tx per slot (bit-24 routes both + // peers' bytes to the even CTA's barrier); "free" = one consumer + // commit per slot reuse. + #pragma unroll + for (int s = 0; s < CD_AB_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.ab_ready[s], 1); + ptx::mbar_init(&smem.hdr.ab_free[s], 1); + } + #pragma unroll + for (int s = 0; s < CD_SF_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.sf_ready[s], 1); + ptx::mbar_init(&smem.hdr.sf_free[s], 1); + } + ptx::mbar_init(&smem.hdr.acc_ready, 1); + ptx::mbar_init(&smem.hdr.acc_free, N_EPI_WARPS * CTA_GROUP); // 4 x 2 = 8 + ptx::mbar_init(&smem.hdr.dealloc, 32); // one warp's lanes, cross-peer + ptx::fence_mbarrier_init_release_cluster(); + } + ptx::cluster_sync_rel_acq(); + + // TMEM allocation: epilogue warp 0 allocates for the pair; warps 0-4 + // rendezvous before reading the granted base. The TMA producer warps + // never touch TMEM and proceed straight to their loops. + uint32_t taddr = 0; + if (warp_id <= 4) { + if (warp_id == 0) + ptx::tcgen05_alloc_2sm(ptx::to_shared(&smem.hdr.tmem_addr), + N_ALLOC_COLS); + ptx::bar_sync(2, 160); + taddr = smem.hdr.tmem_addr; + } + + const int wg = warp_id >> 2; // 0 = epilogue warps, 1 = producers + MMA + if (wg == 1) { + const int pw = warp_id - 4; + if (pw == 1 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 5: lane 0 owns the complete A/B TMA producer loop. + // ================================================================ + // free_phase starts at 1: a fresh mbarrier is phase 0, so the + // first ring's worth of "free" waits falls through. + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + // The pair splits the 256x256 tile: even CTA holds A rows + // [0,128) and B cols [0,128); odd CTA the other halves. + const int m_blk = m_row * CTA_GROUP + int(m_pair); + const int n_blk = n_grp * CTA_GROUP + int(m_pair); + #pragma unroll 1 + for (int g = 0; g < num_groups; ++g) { + // Three full 128-byte windows cover the packed K=768 group. + #pragma unroll + for (int sub = 0; sub < AB_TMA_SUBS; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.ab_free[slot], free_phase); + const uint32_t a_s = + base + CD_A_OFF + uint32_t(slot) * CD_AB_STRIDE; + const uint32_t b_s = + base + CD_B_OFF + uint32_t(slot) * CD_AB_STRIDE; + uint64_t* bar = &smem.hdr.ab_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, AB_TX); + // A rides the multicast form with a self mask so the + // completion routes to the even peer's barrier. + const uint16_t mask_a = uint16_t(1u << m_pair); + // Read the same row-major FP4 allocation as cuBLAS. + // K192 uses four 96-byte boxes; K256 uses three + // 128-byte boxes. The tensor map zero-fills past K. + const CUtensorMap* a_map = &A_tmap; + const CUtensorMap* b_map = &B_tmap; + const int x = g * 384 + sub * AB_LOAD_BYTES; + const int y_a = m_blk * BLOCK_M; + const int y_b = n_blk * BLOCK_N_PEER; + const int z_a = 0; + const int z_b = 0; + ptx::cp_async_bulk_tensor_3d_load_multicast( + a_s, a_map, x, y_a, z_a, + bar, mask_a + ); + ptx::cp_async_bulk_tensor_3d_load_2sm_bit24( + b_s, b_map, x, y_b, z_b, + bar + ); + if (++slot == CD_AB_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + // Drain: wait for the last-filled slot to be consumed, then hand + // it one clean-shutdown expect_tx. + { + const int t_slot = (slot == 0) ? CD_AB_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.ab_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.ab_ready[t_slot], AB_TX); + } + } else if (pw == 2 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 6: scale-factor TMA producer. SFA is per-peer (each CTA's + // own M half); SFB is the one genuinely multicast operand — the + // cta_group::2 copy reads BOTH peers' smem as one logical + // operand, so both copies must be byte-identical (blog: the scale-factor section). + // ================================================================ + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const int m_blk = m_row * CTA_GROUP + int(m_pair); + for (int g = 0; g < num_groups; ++g) { + // 4 SF ring slots per 768-K group (2 K=96 steps each). + #pragma unroll + for (int sub = 0; sub < 4; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.sf_free[slot], free_phase); + const uint32_t sfa_s = + base + CD_SFA_OFF + uint32_t(slot) * CD_SFA_STRIDE; + const uint32_t sfb_s = + base + CD_SFB_OFF + uint32_t(slot) * CD_SFB_STRIDE; + uint64_t* bar = &smem.hdr.sf_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, SF_TX); + // VEC16 is [outer-block][K-group-of-4][128 rows x 4]. + // Three adjacent groups are the 128 x 12 scale slot. + const int sf_group = (g * 4 + sub) * 3; + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfa_s, &SFA_tmap, 0, 0, sf_group, m_blk, bar, + uint16_t(1u << m_pair) + ); + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfb_s + m_pair * uint32_t(CD_SFB_STRIDE / 2), + &SFB_tmap, 0, 0, sf_group, + n_grp * CTA_GROUP + int(m_pair), bar, + uint16_t(0x3) + ); + if (++slot == CD_SF_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + { + const int t_slot = (slot == 0) ? CD_SF_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.sf_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.sf_ready[t_slot], SF_TX); + } + } else if (pw == 0 && m_pair == 0 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 4: the MMA warp — issued once per cluster, on the even + // CTA. The odd CTA's warp 4 does nothing all kernel (blog: the warp table). + // The mainloop needs no tile coordinates: operands arrive through + // the rings; this warp only needs its share of the trip count. + // ================================================================ + const uint32_t smbase = ptx::to_shared(&smem); + const uint16_t commit_mask = 0x3; // both CTAs of this pair + MainloopState st{1u, 0u, 0u, 0u, 0u}; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + st = mainloop_tile(taddr, smbase, st, commit_mask, + uint32_t(full_groups) + ); + // The outer role branch already leaves exactly one issuer. + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;" + :: "r"(smbase + uint32_t(MB_ACC_READY)), + "h"(uint16_t(0x3)) : "memory"); + } + // Drain the final tile's buffer-free arrive before exiting. + mbar_wait_relaxed(smbase + uint32_t(MB_ACC_FREE), st.d_parity); + } + } else { + // ==================================================================== + // Warps 0-3: the epilogue (blog: the epilogue). Each warp owns a 32-row band + // and drains a tile as eight 32-column tcgen05.ld loads, converts + // FP32->FP16 in packed pairs, and stores straight to global — no + // shared-memory bounce. + // ==================================================================== + const int epi_warp = warp_id; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + const uint32_t acc_free_rank = crank & ~1u; // the pair's even CTA + int d_db = 0; + int acc_ready_phase = 0; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const uint32_t tmem_base = taddr + uint32_t(d_db) * D_STRIDE; + mbar_wait_relaxed(ptx::to_shared(&smem.hdr.acc_ready), + uint32_t(acc_ready_phase)); + acc_ready_phase ^= 1; + + const int off_m = (m_row * CTA_GROUP + int(m_pair)) * BLOCK_M; + const int off_n = n_grp * BLOCK_N; + const int row = off_m + epi_warp * 32 + lane_id; + const bool row_in = (row < M); + constexpr int EPI_BANDS = BLOCK_N / EPI_CHUNK; // 8 + #pragma unroll 1 + for (int k = 0; k < EPI_BANDS; ++k) { + const int band = k; + const uint32_t taddr_n = + tmem_base + uint32_t(band) * 32 + taddr_lane; + uint32_t w[32]; + ptx::tcgen05_ld_32x32b_x32(taddr_n, w); + // One buffer cannot be reused until all eight bands are out. + if (k == EPI_BANDS - 1) { + ptx::tcgen05_wait_ld_then_arrive_release( + &smem.hdr.acc_free, acc_free_rank, + /*do_arrive=*/ptx::lane_id() == 0); + } + if (!row_in) continue; + const int col_base = off_n + band * 32; + uint32_t h[16]; + #pragma unroll + for (int j = 0; j < 16; ++j) { + h[j] = ptx::cvt_pack_f16x2(__int_as_float(w[2 * j + 1]), + __int_as_float(w[2 * j])); + } + const size_t row_base = size_t(row) * size_t(N); + const bool aligned32 = ((row_base & 15u) == 0u) && + ((uintptr_t(C_ptr) & 31u) == 0u); + if (aligned32 && (col_base & 15) == 0) { + // Precondition here: row_base is 16-element (32-byte) + // aligned — the aligned32 guard owns it. Restating it to + // ptxas via __builtin_assume compiled this file + // SASS-identical (and measured perf-null on the + // production kernel), so the guard stands alone. + #pragma unroll + for (int q = 0; q < 2; ++q) { + const int c0 = col_base + q * 16; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 16 <= N) { + uint4 lo{h[q*8+0], h[q*8+1], h[q*8+2], h[q*8+3]}; + uint4 hi{h[q*8+4], h[q*8+5], h[q*8+6], h[q*8+7]}; + // Bring-up form: twice the store requests, + // twice the L2 store traffic (the ladder's + // tensor-activity 82% -> 42% row). + ptx::st_global_b128( + reinterpret_cast(cdst), lo); + ptx::st_global_b128( + reinterpret_cast(cdst) + 1, hi); + } else { + // N tail: 128-bit then scalar. + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int cq = c0 + p * 8; + if (cq + 8 <= N) { + uint4 v{h[q*8+p*4+0], h[q*8+p*4+1], + h[q*8+p*4+2], h[q*8+p*4+3]}; + ptx::st_global_b128( + reinterpret_cast(cdst + p * 8), v); + } else { + const uint16_t* hb = + reinterpret_cast( + &h[q * 8 + p * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (cq + c < N) cdst[p * 8 + c] = hb[c]; + } + } + } + } + } else { + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int c0 = col_base + q * 8; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 8 <= N) { + uint4 v{h[q*4+0], h[q*4+1], h[q*4+2], h[q*4+3]}; + ptx::st_global_b128(reinterpret_cast(cdst), v); + } else { + const uint16_t* hb = + reinterpret_cast(&h[q * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (c0 + c < N) cdst[c] = hb[c]; + } + } + } + } + if (++d_db == D_DB) d_db = 0; + } + } + + // Teardown: only the epilogue warps rendezvous per CTA; warp 0 then runs + // the cross-peer dealloc handshake. Producer and MMA warps exit freely — + // their drains above already order their work. + if (wg == 0) { + ptx::bar_sync(3, 128); + if (warp_id == 0) { + ptx::tcgen05_relinquish_2sm(); + const uint32_t peer_rank = crank ^ 1u; + // All 32 lanes arrive the PEER's count-32 barrier = one flip. + ptx::mbar_arrive_cluster_release(&smem.hdr.dealloc, peer_rank); + ptx::mbar_wait_parity(&smem.hdr.dealloc, 0); + ptx::tcgen05_dealloc_2sm(taddr, N_ALLOC_COLS); + } + } +} + + +// Keep the host launch convention paired with the kernel's block attribute. +static inline dim3 launch_grid(int clusters) { + return GRID_IN_CLUSTERS ? dim3(1, 1, unsigned(clusters)) + : dim3(2, 1, unsigned(clusters)); +} + +} // namespace nvfp4 + +// ============================================================================ +// Host-side TMA descriptor support. The shared input and benchmark code lives +// in main.cu. +// ============================================================================ +namespace tmap { + +// Thin wrappers over cuTensorMapEncodeTiled. Dim 0 is the innermost +// (stride-1) axis; strides for dim 1+ are in bytes. +inline CUtensorMap encode_tiled_3d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_depth, uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t depth_stride_bytes, + uint32_t box_depth, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[3] = { global_cols, global_rows, global_depth }; + cuuint64_t global_strides[2] = { row_stride_bytes, depth_stride_bytes }; + cuuint32_t box_dim[3] = { box_cols, box_rows, box_depth }; + cuuint32_t element_strides[3] = { 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 3, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +inline CUtensorMap encode_tiled_4d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_slots, uint64_t global_tiles, + uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t tile_stride_bytes, + uint64_t slot_stride_bytes, + uint32_t box_tiles, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[4] = { global_cols, global_rows, global_tiles, + global_slots }; + cuuint64_t global_strides[3] = { row_stride_bytes, tile_stride_bytes, + slot_stride_bytes }; + cuuint32_t box_dim[4] = { box_cols, box_rows, box_tiles, 1u }; + cuuint32_t element_strides[4] = { 1, 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 4, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +} // namespace tmap + +namespace host { + +using namespace nvfp4; + +// ---- implementation-specific tensor maps --------------------------------- + +inline CUtensorMap make_ab_tmap(void* ptr, int rows, int K, + int row_stride) { + return tmap::encode_tiled_3d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*depth=*/1, uint64_t(rows), /*cols=*/uint64_t(K / 2), + /*row_stride=*/uint64_t(row_stride), + /*depth_stride=*/uint64_t(rows) * uint64_t(row_stride), + /*box=*/1, BLOCK_M, AB_LOAD_BYTES, CU_TENSOR_MAP_SWIZZLE_128B); +} + +inline CUtensorMap make_sf_tmap(uint8_t* ptr, int outer, int K) { + const uint64_t sf_inner = (uint64_t(K) + 63) / 64 * 4; + const uint64_t outer_blocks = (uint64_t(outer) + 127) / 128; + return tmap::encode_tiled_4d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*slots=*/outer_blocks, /*tiles=*/sf_inner / 4, + /*rows=*/4, /*cols=*/128, + /*row_stride=*/128, /*tile_stride=*/512, + /*slot_stride=*/sf_inner * 128, + /*box_tiles=*/3, /*box_rows=*/4, /*box_cols=*/128, + CU_TENSOR_MAP_SWIZZLE_NONE); +} +} // namespace host diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/gemm3.cuh b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm3.cuh new file mode 100644 index 000000000..e9c504047 --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm3.cuh @@ -0,0 +1,1256 @@ +#pragma once +// r3: double-buffer TMEM with overlapping accumulators. +// This is a complete, directly editable ladder snapshot. +// It intentionally duplicates the implementation in the other gemm*.cuh files. + +// Conditional branches are already materialized; edit this implementation directly. +#define NVFP4_BUILD_NAME "r3" +// Read-only metadata for main.cu; these do not select code in this header. +#define NVFP4_HAS_K64_TAIL 0 +#define NVFP4_HAS_SCHEDULE 0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(expr) do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \ + cudaGetErrorName(_e), __FILE__, __LINE__, \ + cudaGetErrorString(_e)); \ + std::abort(); \ + } \ +} while (0) + +#define CU_CHECK(expr) do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + std::fprintf(stderr, "CU error at %s:%d: %s\n", \ + __FILE__, __LINE__, _s ? _s : "?"); \ + std::abort(); \ + } \ +} while (0) + + +// ============================================================================ +// PTX helpers. Thin zero-cost wrappers over the exact instruction forms the +// kernel needs — nothing generic, nothing speculative. +// ============================================================================ +namespace ptx { + +// Generic pointer -> 32-bit `.shared` window byte address. +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Local `.shared` address -> the same offset in CTA `cta_rank`'s smem +// (`.shared::cluster` instructions take this mapped form). +static __device__ __forceinline__ uint32_t mapa_shared_cluster( + uint32_t local_addr, uint32_t cta_rank) { + uint32_t mapped; + asm("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(mapped) : "r"(local_addr), "r"(cta_rank)); + return mapped; +} + +static __device__ __forceinline__ void bar_sync(uint32_t id, uint32_t count) { + asm volatile("barrier.cta.sync %0, %1;" :: "r"(id), "r"(count)); +} + +// Cluster barrier with release/acquire — the publish boundary between the +// mbarrier inits and any cross-CTA use of those mbarriers. +static __device__ __forceinline__ void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// %cluster_ctarank: 0 or 1 in this kernel's (2,1,1) clusters. +static __device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} + +static __device__ __forceinline__ uint32_t lane_id() { + uint32_t r; + asm("mov.u32 %0, %%laneid;" : "=r"(r)); + return r; +} + +// ---- mbarriers ------------------------------------------------------------- + +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + :: "r"(to_shared(bar)), "r"(count)); +} + +// Publish the inits to the async proxy + the peer CTA before any TMA +// completes on them. +static __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// Cross-CTA arrive with release ordering (teardown handshake). +static __device__ __forceinline__ void mbar_arrive_cluster_release( + uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" + :: "r"(mapped)); +} + +// Fused `tcgen05.wait::ld` + "buffer free" arrive, one asm block. The arrive +// must not become visible while TMEM loads are still in flight (it frees the +// accumulator buffer to the next tile's MMA — the bring-up chapter's +// drain overlap); the fused block +// keeps the arrive pinned behind the wait. (An earlier form also manufactured +// a register dependency on a drained value, `dep * 0` folded into the +// address; verified dead on nvcc 13.1 — ptxas const-folds it, SASS +// byte-identical without it — and retired.) +static __device__ __forceinline__ void tcgen05_wait_ld_then_arrive_release( + uint64_t* bar, uint32_t cta_rank, bool do_arrive) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile( + "{\n\t" + ".reg .pred %%pdo;\n\t" + "setp.ne.s32 %%pdo, %1, 0;\n\t" + "tcgen05.wait::ld.sync.aligned;\n\t" + "@%%pdo mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;\n\t" + "}\n\t" + :: "r"(mapped), "r"((int)do_arrive) : "memory"); +} + + +// Blocking parity wait at the consuming site (producer ring/tail waits + +// teardown). The 10 ms suspend hint is the polite form: a full-rate spin +// competes for issue slots and measurably loses; a nanosleep-backoff spin +// measured the same as this hint across five windows, so the simpler form +// ships. The carried one-stage-ahead try_wait peeks that used to gate these +// waits measured as pure overhead and were removed. +static __device__ __forceinline__ void mbar_wait_parity( + uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(to_shared(bar)), "r"(parity)); +} + +// ---- TMA ------------------------------------------------------------------- + +// Producer expect_tx arrive on the cluster-shared ring barrier: both +// peers' transfer bytes account against the even CTA's barrier, so one +// barrier per ring slot covers the whole cluster-wide gather. (Mechanism: +// `& 0xFEFFFFFF` clears bit 24 of the shared address — the hardware's +// peer-CTA bit — re-pointing the arrive at the even CTA's copy.) +static __device__ __forceinline__ void mbar_arrive_expect_tx_cluster( + uint64_t* bar, uint32_t bytes) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes)); +} + +// 3D tiled load, cta_group::2 + multicast (the A feed; the mask is the +// issuing CTA itself in a (2,1) cluster). Completion routes to the even +// CTA's mbarrier via the bit-24 clear. +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z) + : "memory"); +} + +// 3D tiled load, cta_group::2, no multicast (the B feed — each peer loads +// its own 128-wide N-block). +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_2sm_bit24( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst_smem), "l"(tmap), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) + : "memory"); +} + +// 4D tiled load, cta_group::2 + multicast (the scale-factor feed). The +// 512-byte SF tiles of one ring slot form a 4D box; `w` selects the slot. +static __device__ __forceinline__ void cp_async_bulk_tensor_4d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z2, int32_t w, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.4d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z2), "r"(w) + : "memory"); +} + + +// ---- tcgen05 (tensor memory + MMA plumbing) --------------------------------- + +static __device__ __forceinline__ void tcgen05_alloc_2sm( + uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_dealloc_2sm( + uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" + :: "r"(taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_relinquish_2sm() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); +} + +// The widest single-instruction TMEM read: 32 b32 registers per lane +// (blog: the epilogue — a quarter the load/wait count of the narrow forms). +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t taddr, uint32_t* r) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), + "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), + "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), + "=r"(r[16]), "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), + "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), + "=r"(r[24]), "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), + "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr)); +} + +// ---- stores + conversion ----------------------------------------------------- + +static __device__ __forceinline__ uint32_t cvt_pack_f16x2(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// Baseline output store: use the default cache policy. +static __device__ __forceinline__ void st_global_b128(uint4* ptr, uint4 v) { + asm volatile("st.global.v4.b32 [%0], {%1, %2, %3, %4};" + :: "l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)); +} + + + +// ---- MMA descriptors --------------------------------------------------------- + +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class ScaleFormat : uint8_t { E4M3 = 0, E8M0 = 1 }; + +// Source descriptor for the tcgen05.cp that stages scale factors from smem +// into TMEM: unswizzled, SBO=128, LBO=16. The low word later carries the +// per-slot byte offset (in 16-byte units). +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_sf_cp_src_desc_cd(uint32_t sf_smem_addr) { + uint64_t addr_field = uint64_t((sf_smem_addr >> 4) & 0x3FC0u); + return addr_field | 0x400800010000ull; +} + +// Operand descriptor for a K=96 (48-byte) slice inside a 128-byte-swizzled +// smem K-window. OR-constant 0x4010404000000000: version=1, absolute-address +// LBO mode (bit 52), 128B swizzle, SBO=1024. Both the address field and the +// absolute-LBO field encode the window's own 1024-aligned base; the +// intra-window offset advances the address field only. +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k96_sw128(uint32_t window_base_1024aligned, int intra_window_off = 0) { + uint64_t addr_field = uint64_t(((window_base_1024aligned >> 4) & 0x7FC0u) + + uint32_t(intra_window_off >> 4)); + uint64_t lbo_field = uint64_t((window_base_1024aligned << 12) & 0x7FC00000u); + return addr_field | lbo_field | 0x4010404000000000ull; +} + +// Instruction descriptor for tcgen05.mma.kind::mxf4nvf4.block_scale. +// The K extent is selected by BIT 31, not by the asm mnemonic: k_dim=0 runs +// K=64 (0x10400480), k_dim=1 runs K=96 (0x90400480, sm_103a-only) — the +// cheapest quarter of the gap if your kernel is built on the K=64 form +// (blog: the MMA warp). Bit 13 flips which TMEM scale-factor column the MMA reads. +__host__ __device__ static __forceinline__ constexpr uint32_t +mma_inst_desc_mxf4nvf4_block16( + uint32_t M, uint32_t N, + uint32_t tmem_sfa_addr, uint32_t tmem_sfb_addr, + Major a_major, Major b_major, + bool negate_a, bool negate_b, + bool k_dim, ScaleFormat sf) { + constexpr uint32_t MXF4_E2M1 = 1u; + uint32_t d = 0; + d |= ((tmem_sfb_addr & 0xC0000000u) >> 30) << 4; // b_sf_id + d |= MXF4_E2M1 << 7; // a_format = E2M1 + d |= MXF4_E2M1 << 10; // b_format = E2M1 + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; // N in units of 8 + d |= (static_cast(sf) & 0x1u) << 23; // 0 = UE4M3 + d |= ((M >> 4) & 0x1Fu) << 24; // m_dim at bits 27-28 + d |= ((tmem_sfa_addr & 0xC0000000u) >> 30) << 29; // a_sf_id + if (k_dim) d |= 1u << 31; // 1 = K=96 + return d; +} + +} // namespace ptx + +// ============================================================================ +// Kernel geometry from the blog's opening kernel section. One tile shape, +// one cluster shape: +// 256x256 output per 2-CTA cluster (128 rows per CTA), K advanced in +// 96-element MMA steps staged as 768-element K-groups of 8 steps, persistent +// grid of 76 clusters = 152 CTAs = one per SM. +// ============================================================================ +namespace nvfp4 { + +constexpr int CTA_GROUP = 2; // cta_group::2 — the MMA spans both CTAs +constexpr int BLOCK_M = 128; // A rows per CTA +constexpr int BLOCK_N = 256; // output tile N extent (per cluster) +constexpr int BLOCK_N_PEER = BLOCK_N / CTA_GROUP; // 128 — B cols per CTA +constexpr int CLUSTER_M = CTA_GROUP * BLOCK_M; // 256 — tile M extent +[[maybe_unused]] constexpr int MMA_K = 96; // K per tcgen05.mma (idesc bit 31) +constexpr int WIN_BYTES = 128; // one 128B-swizzled smem K-window +constexpr int AB_LOAD_BYTES = WIN_BYTES; +constexpr int AB_TMA_SUBS = 3; +constexpr int TB_SIZE = 224; // 7 warps (blog: the warp table) +constexpr int N_EPI_WARPS = 4; // warps 0-3 + +constexpr int SF_TMEM_COLS = 9 * 4; // 3 SFA + 6 SFB cp tiles, 4 cols each +// Two 256-wide accumulators use a 220-column stride and overlap by the 36 +// columns needed for scale factors. +constexpr int D_STRIDE = 220; +constexpr int D_DB = 2; +constexpr int N_D_COLS = D_STRIDE + BLOCK_N; +constexpr int SF_TMEM_BASE = 476; +constexpr int N_ALLOC_COLS = 512; +static_assert(N_D_COLS <= SF_TMEM_BASE && + SF_TMEM_BASE + SF_TMEM_COLS <= N_ALLOC_COLS, + "accumulators + scale factors exceed the 512 TMEM columns"); +constexpr int EPI_CHUNK = 32; // cols per regular drain + +// ---- the shared-memory map (blog: the feed) ----------------------------------- +// Every offset is a compile-time expression. +constexpr int CD_HDR = 1024; // mbarrier header pad +constexpr int CD_AB_SLOTS = 6; +constexpr int CD_SF_SLOTS = 7; +constexpr int CD_AB_STRIDE = BLOCK_M * WIN_BYTES; // 16384 per A or B slot +constexpr int CD_SF_TILE = 512; // one (32 lanes x 16B) x4 tile +constexpr int CD_SFA_TILES = 3; // per SFA slot (M = 1 block) +constexpr int CD_SFB_TILES = 6; // per SFB slot (N = 2 blocks) +constexpr int CD_SFA_STRIDE = CD_SFA_TILES * CD_SF_TILE; // 1536 +constexpr int CD_SFB_STRIDE = CD_SFB_TILES * CD_SF_TILE; // 3072 + +// Expected transfer bytes are derived, never tuned: a wrong-by-one +// expectation doesn't crash, it deadlocks (blog: the feed). +constexpr uint32_t AB_TX = uint32_t(CTA_GROUP) * + uint32_t(BLOCK_M + BLOCK_N_PEER) * uint32_t(AB_LOAD_BYTES); +constexpr uint32_t SF_TX = uint32_t(CD_SFA_STRIDE + CD_SFB_STRIDE) * + uint32_t(CTA_GROUP); // 9216 +static_assert(AB_TX == 65536u && SF_TX == 9216u, + "expect_tx derivation drifted"); + +// The map IS the struct: member sizes and alignment produce every byte +// offset; the raw-address constants the mainloop uses read off it with +// offsetof, and the static_asserts pin the published values. +struct SmemCD { + // mbarrier header: four ring regions in handshake order, then the + // singleton accumulator + teardown barriers (alignas pads it to 1024). + struct alignas(1024) Header { + uint64_t ab_ready[CD_AB_SLOTS]; // TMA -> MMA + uint64_t ab_free [CD_AB_SLOTS]; // MMA -> TMA + uint64_t sf_ready[CD_SF_SLOTS]; // TMA -> MMA + uint64_t sf_free [CD_SF_SLOTS]; // MMA -> TMA + uint64_t acc_ready; // MMA -> epi + uint64_t acc_free; // epi -> MMA + uint64_t dealloc; // 2-CTA teardown + uint32_t tmem_addr; // TMEM base addr storage + } hdr; + uint8_t a [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t b [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t sfa[CD_SF_SLOTS][CD_SFA_STRIDE]; + // SFB starts 1024-aligned (the cp source descriptor drops address bits 4-9). + alignas(1024) uint8_t sfb[CD_SF_SLOTS][CD_SFB_STRIDE]; +}; +static_assert(sizeof(SmemCD::Header) == CD_HDR, "mbarrier header overflows the pad"); + +constexpr int CD_A_OFF = offsetof(SmemCD, a); +constexpr int CD_B_OFF = offsetof(SmemCD, b); +constexpr int CD_SFA_OFF = offsetof(SmemCD, sfa); +constexpr int CD_SFB_OFF = offsetof(SmemCD, sfb); +constexpr int CD_TOTAL = sizeof(SmemCD); +static_assert(CD_TOTAL <= 232448, "smem map exceeds the opt-in budget"); +static_assert(CD_AB_SLOTS != 6 || + (CD_B_OFF == 99328 && CD_SFA_OFF == 197632 && + CD_SFB_OFF == 208896 && CD_TOTAL == 230400), + "the 6-slot byte map drifted from its published values"); + +// The mainloop computes raw `.shared` byte addresses from these. +constexpr int MB_AB_READY = offsetof(SmemCD, hdr.ab_ready); +constexpr int MB_AB_FREE = offsetof(SmemCD, hdr.ab_free); +constexpr int MB_SF_READY = offsetof(SmemCD, hdr.sf_ready); +constexpr int MB_SF_FREE = offsetof(SmemCD, hdr.sf_free); +constexpr int MB_ACC_READY = offsetof(SmemCD, hdr.acc_ready); +constexpr int MB_ACC_FREE = offsetof(SmemCD, hdr.acc_free); + + + +} // namespace nvfp4 + +// ============================================================================ +// The consumer mainloop from the blog's pipeline section. One warp — warp 4 +// of the even CTA — +// issues everything: the tcgen05.cp copies that stage scale factors from the +// smem ring into TMEM, immediately followed by the tcgen05.mma stream, in the +// SAME warp. A same-warp tcgen05 stream is one ordered async pipeline, so the +// write-after-read hazard on the single reused scale-factor TMEM region needs +// no fence and no cross-warp barrier. +// +// Per 768-element K-group, logical K=96 steps 0..7 issue in K order. Their +// starting offsets, expressed as 16-byte descriptor atoms modulo a 128-byte +// swizzle window, are 0,3,6,1,4,7,2,5. Atoms 6 and 7 straddle windows. +// ============================================================================ +namespace nvfp4 { + + +// Ring/parity state carried across tiles by the MMA warp. +struct MainloopState { + uint32_t d_parity; // which TMEM accumulator buffer the next tile writes + uint32_t ab_phase, ab_slot; // A/B ring cursor + uint32_t sf_phase, sf_slot; // scale-factor ring cursor +}; + + +// Relaxed blocking wait (no .acquire): the tcgen05 pipeline's own ordering +// covers these ring waits, and the .acquire form perturbs the scoreboard +// cadence on the issuing warp. Waits sit AT the consuming site — the carried +// one-stage-ahead try_wait peeks that used to gate them measured as pure +// overhead and were removed. +static __device__ __forceinline__ void mbar_wait_relaxed( + uint32_t bar_smem, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(bar_smem), "r"(parity) : "memory"); +} + +// Gated form for the @216 buffer-free wait, which fires once per tile: the +// gate branches over the spin for every group past the first. +static __device__ __forceinline__ void mbar_wait_gated_relaxed( + uint32_t bar_smem, uint32_t parity, bool peek) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p bra DONE_%=;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t" + "DONE_%=:\n\t}\n" + :: "r"(bar_smem), "r"(parity), "r"((uint32_t)peek) : "memory"); +} + + +// One scale-factor staging copy, smem -> TMEM. Single-lane semantics behind +// the converged elect predicate: tcgen05.cp is a single-thread instruction — +// issued bare from a converged warp you get 32 initiations and garbage. +static __device__ __forceinline__ void cp_sf( + uint32_t el, uint32_t dst, uint32_t lo, uint32_t hi) { + const uint64_t desc = (uint64_t(hi) << 32) | uint64_t(lo); + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;\n\t}\n" + :: "r"(dst), "l"(desc), "r"(el) : "memory"); +} + +// Ring-slot "consumed" arrive, multicast to both CTAs of the pair. +// arrive::one must fire exactly once — hence the elect predicate. +static __device__ __forceinline__ void commit_consumed( + uint32_t el, uint32_t smbase, uint32_t mb_off, uint16_t mask) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t}\n" + :: "r"(smbase + mb_off), "h"(mask), "r"(el) : "memory"); +} + +// One K=96 NVFP4 MMA. The asm token says scale_vec::4X regardless — K=96 is +// selected by idesc bit 31 (blog: the MMA warp). acc_pred = 0 overwrites the +// accumulator (the tile's first step), nonzero accumulates. +static __device__ __forceinline__ void mma_cell( + uint32_t el, uint32_t d, + uint32_t desc_a_lo, uint32_t desc_a_hi, + uint32_t desc_b_lo, uint32_t desc_b_hi, + uint32_t idesc, uint32_t sf_a, uint32_t sf_b, uint32_t acc_pred) { + const uint64_t desc_a = (uint64_t(desc_a_hi) << 32) | uint64_t(desc_a_lo); + const uint64_t desc_b = (uint64_t(desc_b_hi) << 32) | uint64_t(desc_b_lo); + asm volatile( + "{\n\t.reg .pred p, acc;\n\t" + "setp.ne.u32 p, %7, 0;\n\t" + "setp.ne.b32 acc, %4, 0;\n\t" + "@p tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X" + " [%0], %1, %2, %3, [%5], [%6], acc;\n\t}\n" + :: "r"(d), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(acc_pred), "r"(sf_a), "r"(sf_b), "r"(el) : "memory"); +} + +// One SF ring slot -> TMEM: 9 tcgen05.cp tiles (3 SFA + 6 SFB). Source +// address advances are in 16-byte descriptor units. The SFB tile order +// 0,3,1,4,2,5 interleaves the two 128-wide N-blocks. +static __device__ __forceinline__ void cp_sf_buffer( + uint32_t el, uint32_t slot, + uint32_t sfa0, uint32_t sfa1, uint32_t sfa2, + uint32_t sfb0, uint32_t sfb1, uint32_t sfb2, + uint32_t sfb3, uint32_t sfb4, uint32_t sfb5, + uint32_t sfa_hi, uint32_t sfa_lo, uint32_t sfb_hi, uint32_t sfb_lo) { + constexpr uint32_t SFA_SLOT_U = CD_SFA_STRIDE / 16u; // 96 + constexpr uint32_t SFB_SLOT_U = CD_SFB_STRIDE / 16u; // 192 + constexpr uint32_t TILE_U = CD_SF_TILE / 16u; // 32 + const uint32_t a = slot * SFA_SLOT_U + sfa_lo; + cp_sf(el, sfa0, a, sfa_hi); + cp_sf(el, sfa1, a + 1u*TILE_U, sfa_hi); + cp_sf(el, sfa2, a + 2u*TILE_U, sfa_hi); + const uint32_t b = slot * SFB_SLOT_U + sfb_lo; + cp_sf(el, sfb0, b + 0u*TILE_U, sfb_hi); + cp_sf(el, sfb1, b + 3u*TILE_U, sfb_hi); + cp_sf(el, sfb2, b + 1u*TILE_U, sfb_hi); + cp_sf(el, sfb3, b + 4u*TILE_U, sfb_hi); + cp_sf(el, sfb4, b + 2u*TILE_U, sfb_hi); + cp_sf(el, sfb5, b + 5u*TILE_U, sfb_hi); +} + +// Operand descriptor LOW word for A/B ring slot `slot`, K-step `k_atom` +// (0..7 inside the 768-element group; each step is 48 bytes = +3 in the +// 16-byte address field). +static __device__ __forceinline__ uint32_t ab_desc_lo( + uint32_t smbase, uint32_t ab_base_off, uint32_t slot, uint32_t k_atom) { + const uint64_t d = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + slot * CD_AB_STRIDE, int(k_atom) * 16); + return uint32_t(d & 0xFFFFFFFFu); +} + +// K-steps 6 and 7 cross a 128-byte window boundary. Their descriptor address +// comes from `addr_slot`, while the absolute-LBO field comes from the next +// slot whose swizzle window contains the second half. +static __device__ __forceinline__ uint32_t ab_desc_lo_straddle( + uint32_t smbase, uint32_t ab_base_off, + uint32_t addr_slot, uint32_t lbo_slot, uint32_t k_atom) { + const uint64_t d_addr = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + addr_slot * CD_AB_STRIDE, int(k_atom) * 16); + const uint64_t d_lbo = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + lbo_slot * CD_AB_STRIDE, 0); + return uint32_t(d_addr & 0x0000FFFFu) | uint32_t(d_lbo & 0xFFFF0000u); +} + +// One output tile: `full_groups` complete 768-element K-groups. Coordinates +// past logical K are zero-filled by TMA from the unchanged canonical inputs. +static __device__ __forceinline__ MainloopState mainloop_tile( + uint32_t taddr, uint32_t smbase, MainloopState in, uint16_t commit_mask, + uint32_t full_groups + ) { + constexpr uint32_t AB_RING_SLOTS = uint32_t(CD_AB_SLOTS); + constexpr uint32_t SF_RING_SLOTS = uint32_t(CD_SF_SLOTS); + constexpr uint32_t AB_FREE_DELTA = uint32_t(MB_AB_FREE - MB_AB_READY); + constexpr uint32_t SF_READY_OFF = uint32_t(MB_SF_READY); + constexpr uint32_t SF_FREE_DELTA = uint32_t(MB_SF_FREE - MB_SF_READY); + constexpr uint32_t ACC_FREE_OFF = uint32_t(MB_ACC_FREE); + constexpr uint32_t A_OFF = uint32_t(CD_A_OFF); + constexpr uint32_t B_OFF = uint32_t(CD_B_OFF); + + uint32_t d_parity = in.d_parity; + uint32_t ab_phase = in.ab_phase, ab_slot = in.ab_slot; + uint32_t sf_phase = in.sf_phase, sf_slot = in.sf_slot; + + // The caller has already selected exactly one lane for the whole MMA role. + // Keep the helper predicate constant throughout the tile. + const uint32_t el = 1u; + + // Fixed TMEM targets for the single reused scale-factor region: 9 cp + // tiles of 4 columns each, starting past the accumulator region. + constexpr uint32_t SF = uint32_t(SF_TMEM_BASE); + const uint32_t sfa_t0 = taddr + SF + 0; // SFA word 0 (MMA reads here) + const uint32_t sfa_t1 = taddr + SF + 4; + const uint32_t sfa_t2 = taddr + SF + 8; + const uint32_t sfb_t0 = taddr + SF + 12; // SFB word 0 + const uint32_t sfb_t1 = taddr + SF + 16; + const uint32_t sfb_t2 = taddr + SF + 20; + const uint32_t sfb_t3 = taddr + SF + 24; + const uint32_t sfb_t4 = taddr + SF + 28; + const uint32_t sfb_t5 = taddr + SF + 32; + + // Scale-factor cp source descriptors (low word carries the per-slot + // offset; high word is constant). + const uint64_t sfb_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFB_OFF)); + const uint32_t sfb_src_lo = uint32_t(sfb_src & 0xFFFFFFFFu); + const uint32_t sfb_src_hi = uint32_t(sfb_src >> 32); + const uint64_t sfa_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFA_OFF)); + const uint32_t sfa_src_lo = uint32_t(sfa_src & 0xFFFFFFFFu); + const uint32_t sfa_src_hi = uint32_t(sfa_src >> 32); + + // Operand descriptor HIGH words (low words come from ab_desc_lo per step). + const uint32_t a_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + A_OFF) >> 32); + const uint32_t b_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + B_OFF) >> 32); + + // Each K=96 step needs 6 scale bytes per row, split across two staged + // 32-bit words. Odd-issue steps read word 1: bit 31 of the SF address + // selects the with-id encoding, and the id bits fold into the idesc. + constexpr uint32_t SF_ADDR_ID = 0x80000000u; + const uint32_t sfa_w1 = (taddr + SF + 4) | SF_ADDR_ID; + const uint32_t sfb_w1 = (taddr + SF + 20) | SF_ADDR_ID; + const uint32_t w1_id_bits = + ((sfa_w1 >> 1) & 0x60000000u) | ((sfb_w1 >> 26) & 48u); + + // The tile's first MMA overwrites the accumulator buffer this tile owns; + // the @216 "buffer free" wait runs once, before the first group. + const uint32_t acc_wait_parity = d_parity; + d_parity ^= 1u; + const uint32_t d_tmem = d_parity * uint32_t(D_STRIDE) + taddr; + + // Instruction descriptor: kind::mxf4nvf4, M=256, N=256, E2M1 x E2M1, + // UE4M3 block-16, K-major, k_dim=1 (K=96). Base value 0x90400480. + constexpr uint32_t IDESC = ptx::mma_inst_desc_mxf4nvf4_block16( + 256u, 256u, 0u, 0u, ptx::Major::K, ptx::Major::K, + false, false, /*k_dim=*/true, ptx::ScaleFormat::E4M3); + static_assert(IDESC == 0x90400480u, "K=96 NVFP4 idesc drifted"); + const uint32_t w0_id_bits = + ((sfa_t0 >> 1) & 0x60000000u) | ((sfb_t0 >> 26) & 48u); + const uint32_t idesc_w0 = w0_id_bits | IDESC; + const uint32_t idesc_w1 = w1_id_bits | IDESC; + + + uint32_t acc_pred = 0; // first MMA of the tile overwrites + bool need_acc_wait = true; // @216 wait fires once per tile + + if (full_groups >= 1) { + uint32_t g = 0; + do { + // ---- SF batch 0 -> logical steps 0,1 (atoms 0,3) ---- + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // ---- SF batch 1 -> logical steps 2,3 (atoms 6,1) ---- + mbar_wait_relaxed(sf_bar1, sf_phase1); + uint32_t sf_slot2, sf_phase2; + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + // AB0's last reader is MMA2. + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // ---- SF batch 2 -> logical steps 4,5 (atoms 4,7) ---- + mbar_wait_relaxed(sf_bar2, sf_phase2); + uint32_t sf_slot3, sf_phase3; + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; // ring cursor writeback + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // ---- SF batch 3 -> logical steps 6,7 (atoms 2,5) ---- + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; // ring cursor writeback + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + // AB1's last reader is MMA5; this measured placement leaves the + // commit behind the following scale copy. + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + acc_pred = 0xFFFFFFFFu; // accumulate from the second group on + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + need_acc_wait = false; + // AB2's last reader is MMA7. + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, + commit_mask); + ++g; + } while (g != full_groups); + } + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; +} + +} // namespace nvfp4 + +// ============================================================================ +// The kernel. 224 threads = 7 warps, four jobs (blog: the warp table): +// warps 0-3 epilogue: drain TMEM, convert to FP16, store straight to global +// warp 4 the MMA warp (even CTA of each pair only): SF copies + MMAs +// warp 5 TMA producer: A and B tiles into the 6-slot ring +// warp 6 TMA producer: scale factors into the 7-slot ring +// ============================================================================ +namespace nvfp4 { + +// Baseline launch: plain launch bounds and a grid expressed in CTA units. +#define NVFP4_BLOCK_ATTR +constexpr bool GRID_IN_CLUSTERS = false; + +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void nvfp4_gemm_kernel(const __grid_constant__ CUtensorMap A_tmap, + const __grid_constant__ CUtensorMap B_tmap, + const __grid_constant__ CUtensorMap SFA_tmap, + const __grid_constant__ CUtensorMap SFB_tmap, + __half* C_ptr, int M, int N, int K + ) { + extern __shared__ __align__(1024) char smem_buf[]; + SmemCD& smem = *reinterpret_cast(smem_buf); + + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const uint32_t crank = ptx::cluster_rank(); + const uint32_t m_pair = crank & 1u; // 0 = the pair's even CTA + + const int cluster_grid_m = (M + CLUSTER_M - 1) / CLUSTER_M; + const int grid_n = (N + BLOCK_N - 1) / BLOCK_N; + const int total_tiles = cluster_grid_m * grid_n; + + // Baseline rungs execute a whole K=768 group at the edge. Tensor-map OOB + // fill supplies zeros past logical K without changing the input buffers. + const int full_groups = (K + 767) / 768; + const int num_groups = full_groups; + + // Persistent cluster index. The device always sees the grid in CTA units + // (2, 1, clusters) regardless of the host-side launch-unit regime. + const bool grid_z = (gridDim.z > 1); + const int num_clusters = grid_z ? int(gridDim.z) : int(gridDim.x) / 2; + const int cluster_id = grid_z ? int(blockIdx.z) : int(blockIdx.x) / 2; + + + // Which output tile does persistent work item `t` process? + auto tile_of = [&](int t) -> int { + return t; // natural raster: M fast, N slow + }; + + auto producer_tile_of = [&](int t) -> int { + return tile_of(t); + }; + + // ---- mbarrier init (warp 0, lane 0) ---- + if (warp_id == 0 && ptx::lane_id() == 0) { + // Arrival counts are derived from the topology, never tuned: + // "ready" = one producer expect_tx per slot (bit-24 routes both + // peers' bytes to the even CTA's barrier); "free" = one consumer + // commit per slot reuse. + #pragma unroll + for (int s = 0; s < CD_AB_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.ab_ready[s], 1); + ptx::mbar_init(&smem.hdr.ab_free[s], 1); + } + #pragma unroll + for (int s = 0; s < CD_SF_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.sf_ready[s], 1); + ptx::mbar_init(&smem.hdr.sf_free[s], 1); + } + ptx::mbar_init(&smem.hdr.acc_ready, 1); + ptx::mbar_init(&smem.hdr.acc_free, N_EPI_WARPS * CTA_GROUP); // 4 x 2 = 8 + ptx::mbar_init(&smem.hdr.dealloc, 32); // one warp's lanes, cross-peer + ptx::fence_mbarrier_init_release_cluster(); + } + ptx::cluster_sync_rel_acq(); + + // TMEM allocation: epilogue warp 0 allocates for the pair; warps 0-4 + // rendezvous before reading the granted base. The TMA producer warps + // never touch TMEM and proceed straight to their loops. + uint32_t taddr = 0; + if (warp_id <= 4) { + if (warp_id == 0) + ptx::tcgen05_alloc_2sm(ptx::to_shared(&smem.hdr.tmem_addr), + N_ALLOC_COLS); + ptx::bar_sync(2, 160); + taddr = smem.hdr.tmem_addr; + } + + const int wg = warp_id >> 2; // 0 = epilogue warps, 1 = producers + MMA + if (wg == 1) { + const int pw = warp_id - 4; + if (pw == 1 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 5: lane 0 owns the complete A/B TMA producer loop. + // ================================================================ + // free_phase starts at 1: a fresh mbarrier is phase 0, so the + // first ring's worth of "free" waits falls through. + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + // The pair splits the 256x256 tile: even CTA holds A rows + // [0,128) and B cols [0,128); odd CTA the other halves. + const int m_blk = m_row * CTA_GROUP + int(m_pair); + const int n_blk = n_grp * CTA_GROUP + int(m_pair); + #pragma unroll 1 + for (int g = 0; g < num_groups; ++g) { + // Three full 128-byte windows cover the packed K=768 group. + #pragma unroll + for (int sub = 0; sub < AB_TMA_SUBS; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.ab_free[slot], free_phase); + const uint32_t a_s = + base + CD_A_OFF + uint32_t(slot) * CD_AB_STRIDE; + const uint32_t b_s = + base + CD_B_OFF + uint32_t(slot) * CD_AB_STRIDE; + uint64_t* bar = &smem.hdr.ab_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, AB_TX); + // A rides the multicast form with a self mask so the + // completion routes to the even peer's barrier. + const uint16_t mask_a = uint16_t(1u << m_pair); + // Read the same row-major FP4 allocation as cuBLAS. + // K192 uses four 96-byte boxes; K256 uses three + // 128-byte boxes. The tensor map zero-fills past K. + const CUtensorMap* a_map = &A_tmap; + const CUtensorMap* b_map = &B_tmap; + const int x = g * 384 + sub * AB_LOAD_BYTES; + const int y_a = m_blk * BLOCK_M; + const int y_b = n_blk * BLOCK_N_PEER; + const int z_a = 0; + const int z_b = 0; + ptx::cp_async_bulk_tensor_3d_load_multicast( + a_s, a_map, x, y_a, z_a, + bar, mask_a + ); + ptx::cp_async_bulk_tensor_3d_load_2sm_bit24( + b_s, b_map, x, y_b, z_b, + bar + ); + if (++slot == CD_AB_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + // Drain: wait for the last-filled slot to be consumed, then hand + // it one clean-shutdown expect_tx. + { + const int t_slot = (slot == 0) ? CD_AB_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.ab_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.ab_ready[t_slot], AB_TX); + } + } else if (pw == 2 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 6: scale-factor TMA producer. SFA is per-peer (each CTA's + // own M half); SFB is the one genuinely multicast operand — the + // cta_group::2 copy reads BOTH peers' smem as one logical + // operand, so both copies must be byte-identical (blog: the scale-factor section). + // ================================================================ + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const int m_blk = m_row * CTA_GROUP + int(m_pair); + for (int g = 0; g < num_groups; ++g) { + // 4 SF ring slots per 768-K group (2 K=96 steps each). + #pragma unroll + for (int sub = 0; sub < 4; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.sf_free[slot], free_phase); + const uint32_t sfa_s = + base + CD_SFA_OFF + uint32_t(slot) * CD_SFA_STRIDE; + const uint32_t sfb_s = + base + CD_SFB_OFF + uint32_t(slot) * CD_SFB_STRIDE; + uint64_t* bar = &smem.hdr.sf_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, SF_TX); + // VEC16 is [outer-block][K-group-of-4][128 rows x 4]. + // Three adjacent groups are the 128 x 12 scale slot. + const int sf_group = (g * 4 + sub) * 3; + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfa_s, &SFA_tmap, 0, 0, sf_group, m_blk, bar, + uint16_t(1u << m_pair) + ); + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfb_s + m_pair * uint32_t(CD_SFB_STRIDE / 2), + &SFB_tmap, 0, 0, sf_group, + n_grp * CTA_GROUP + int(m_pair), bar, + uint16_t(0x3) + ); + if (++slot == CD_SF_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + { + const int t_slot = (slot == 0) ? CD_SF_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.sf_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.sf_ready[t_slot], SF_TX); + } + } else if (pw == 0 && m_pair == 0 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 4: the MMA warp — issued once per cluster, on the even + // CTA. The odd CTA's warp 4 does nothing all kernel (blog: the warp table). + // The mainloop needs no tile coordinates: operands arrive through + // the rings; this warp only needs its share of the trip count. + // ================================================================ + const uint32_t smbase = ptx::to_shared(&smem); + const uint16_t commit_mask = 0x3; // both CTAs of this pair + MainloopState st{1u, 0u, 0u, 0u, 0u}; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + st = mainloop_tile(taddr, smbase, st, commit_mask, + uint32_t(full_groups) + ); + // The outer role branch already leaves exactly one issuer. + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;" + :: "r"(smbase + uint32_t(MB_ACC_READY)), + "h"(uint16_t(0x3)) : "memory"); + } + // Drain the final tile's buffer-free arrive before exiting. + mbar_wait_relaxed(smbase + uint32_t(MB_ACC_FREE), st.d_parity); + } + } else { + // ==================================================================== + // Warps 0-3: the epilogue (blog: the epilogue). Each warp owns a 32-row band + // and drains a tile as eight 32-column tcgen05.ld loads, converts + // FP32->FP16 in packed pairs, and stores straight to global — no + // shared-memory bounce. + // ==================================================================== + const int epi_warp = warp_id; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + const uint32_t acc_free_rank = crank & ~1u; // the pair's even CTA + int d_db = 0; + int acc_ready_phase = 0; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const uint32_t tmem_base = taddr + uint32_t(d_db) * D_STRIDE; + mbar_wait_relaxed(ptx::to_shared(&smem.hdr.acc_ready), + uint32_t(acc_ready_phase)); + acc_ready_phase ^= 1; + + const int off_m = (m_row * CTA_GROUP + int(m_pair)) * BLOCK_M; + const int off_n = n_grp * BLOCK_N; + const int row = off_m + epi_warp * 32 + lane_id; + const bool row_in = (row < M); + constexpr int EPI_BANDS = BLOCK_N / EPI_CHUNK; // 8 + #pragma unroll 1 + for (int k = 0; k < EPI_BANDS; ++k) { + // Buffer 0 drains high-to-low and buffer 1 low-to-high, so + // their shared columns [220,256) drain first for both. + const int parity = d_db & 1; + const int band = (parity == 0) ? (EPI_BANDS - 1 - k) : k; + const uint32_t taddr_n = + tmem_base + uint32_t(band) * 32 + taddr_lane; + uint32_t w[32]; + ptx::tcgen05_ld_32x32b_x32(taddr_n, w); + // Return the old buffer as soon as its shared edge is out. + if (k == 1) { + ptx::tcgen05_wait_ld_then_arrive_release( + &smem.hdr.acc_free, acc_free_rank, + /*do_arrive=*/ptx::lane_id() == 0); + } + if (!row_in) continue; + const int col_base = off_n + band * 32; + uint32_t h[16]; + #pragma unroll + for (int j = 0; j < 16; ++j) { + h[j] = ptx::cvt_pack_f16x2(__int_as_float(w[2 * j + 1]), + __int_as_float(w[2 * j])); + } + const size_t row_base = size_t(row) * size_t(N); + const bool aligned32 = ((row_base & 15u) == 0u) && + ((uintptr_t(C_ptr) & 31u) == 0u); + if (aligned32 && (col_base & 15) == 0) { + // Precondition here: row_base is 16-element (32-byte) + // aligned — the aligned32 guard owns it. Restating it to + // ptxas via __builtin_assume compiled this file + // SASS-identical (and measured perf-null on the + // production kernel), so the guard stands alone. + #pragma unroll + for (int q = 0; q < 2; ++q) { + const int c0 = col_base + q * 16; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 16 <= N) { + uint4 lo{h[q*8+0], h[q*8+1], h[q*8+2], h[q*8+3]}; + uint4 hi{h[q*8+4], h[q*8+5], h[q*8+6], h[q*8+7]}; + // Bring-up form: twice the store requests, + // twice the L2 store traffic (the ladder's + // tensor-activity 82% -> 42% row). + ptx::st_global_b128( + reinterpret_cast(cdst), lo); + ptx::st_global_b128( + reinterpret_cast(cdst) + 1, hi); + } else { + // N tail: 128-bit then scalar. + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int cq = c0 + p * 8; + if (cq + 8 <= N) { + uint4 v{h[q*8+p*4+0], h[q*8+p*4+1], + h[q*8+p*4+2], h[q*8+p*4+3]}; + ptx::st_global_b128( + reinterpret_cast(cdst + p * 8), v); + } else { + const uint16_t* hb = + reinterpret_cast( + &h[q * 8 + p * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (cq + c < N) cdst[p * 8 + c] = hb[c]; + } + } + } + } + } else { + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int c0 = col_base + q * 8; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 8 <= N) { + uint4 v{h[q*4+0], h[q*4+1], h[q*4+2], h[q*4+3]}; + ptx::st_global_b128(reinterpret_cast(cdst), v); + } else { + const uint16_t* hb = + reinterpret_cast(&h[q * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (c0 + c < N) cdst[c] = hb[c]; + } + } + } + } + if (++d_db == D_DB) d_db = 0; + } + } + + // Teardown: only the epilogue warps rendezvous per CTA; warp 0 then runs + // the cross-peer dealloc handshake. Producer and MMA warps exit freely — + // their drains above already order their work. + if (wg == 0) { + ptx::bar_sync(3, 128); + if (warp_id == 0) { + ptx::tcgen05_relinquish_2sm(); + const uint32_t peer_rank = crank ^ 1u; + // All 32 lanes arrive the PEER's count-32 barrier = one flip. + ptx::mbar_arrive_cluster_release(&smem.hdr.dealloc, peer_rank); + ptx::mbar_wait_parity(&smem.hdr.dealloc, 0); + ptx::tcgen05_dealloc_2sm(taddr, N_ALLOC_COLS); + } + } +} + + +// Keep the host launch convention paired with the kernel's block attribute. +static inline dim3 launch_grid(int clusters) { + return GRID_IN_CLUSTERS ? dim3(1, 1, unsigned(clusters)) + : dim3(2, 1, unsigned(clusters)); +} + +} // namespace nvfp4 + +// ============================================================================ +// Host-side TMA descriptor support. The shared input and benchmark code lives +// in main.cu. +// ============================================================================ +namespace tmap { + +// Thin wrappers over cuTensorMapEncodeTiled. Dim 0 is the innermost +// (stride-1) axis; strides for dim 1+ are in bytes. +inline CUtensorMap encode_tiled_3d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_depth, uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t depth_stride_bytes, + uint32_t box_depth, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[3] = { global_cols, global_rows, global_depth }; + cuuint64_t global_strides[2] = { row_stride_bytes, depth_stride_bytes }; + cuuint32_t box_dim[3] = { box_cols, box_rows, box_depth }; + cuuint32_t element_strides[3] = { 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 3, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +inline CUtensorMap encode_tiled_4d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_slots, uint64_t global_tiles, + uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t tile_stride_bytes, + uint64_t slot_stride_bytes, + uint32_t box_tiles, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[4] = { global_cols, global_rows, global_tiles, + global_slots }; + cuuint64_t global_strides[3] = { row_stride_bytes, tile_stride_bytes, + slot_stride_bytes }; + cuuint32_t box_dim[4] = { box_cols, box_rows, box_tiles, 1u }; + cuuint32_t element_strides[4] = { 1, 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 4, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +} // namespace tmap + +namespace host { + +using namespace nvfp4; + +// ---- implementation-specific tensor maps --------------------------------- + +inline CUtensorMap make_ab_tmap(void* ptr, int rows, int K, + int row_stride) { + return tmap::encode_tiled_3d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*depth=*/1, uint64_t(rows), /*cols=*/uint64_t(K / 2), + /*row_stride=*/uint64_t(row_stride), + /*depth_stride=*/uint64_t(rows) * uint64_t(row_stride), + /*box=*/1, BLOCK_M, AB_LOAD_BYTES, CU_TENSOR_MAP_SWIZZLE_128B); +} + +inline CUtensorMap make_sf_tmap(uint8_t* ptr, int outer, int K) { + const uint64_t sf_inner = (uint64_t(K) + 63) / 64 * 4; + const uint64_t outer_blocks = (uint64_t(outer) + 127) / 128; + return tmap::encode_tiled_4d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*slots=*/outer_blocks, /*tiles=*/sf_inner / 4, + /*rows=*/4, /*cols=*/128, + /*row_stride=*/128, /*tile_stride=*/512, + /*slot_stride=*/sf_inner * 128, + /*box_tiles=*/3, /*box_rows=*/4, /*box_cols=*/128, + CU_TENSOR_MAP_SWIZZLE_NONE); +} +} // namespace host diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/gemm4.cuh b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm4.cuh new file mode 100644 index 000000000..acf78b819 --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm4.cuh @@ -0,0 +1,1258 @@ +#pragma once +// r4: use 256-bit output stores. +// This is a complete, directly editable ladder snapshot. +// It intentionally duplicates the implementation in the other gemm*.cuh files. + +// Conditional branches are already materialized; edit this implementation directly. +#define NVFP4_BUILD_NAME "r4" +// Read-only metadata for main.cu; these do not select code in this header. +#define NVFP4_HAS_K64_TAIL 0 +#define NVFP4_HAS_SCHEDULE 0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(expr) do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \ + cudaGetErrorName(_e), __FILE__, __LINE__, \ + cudaGetErrorString(_e)); \ + std::abort(); \ + } \ +} while (0) + +#define CU_CHECK(expr) do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + std::fprintf(stderr, "CU error at %s:%d: %s\n", \ + __FILE__, __LINE__, _s ? _s : "?"); \ + std::abort(); \ + } \ +} while (0) + + +// ============================================================================ +// PTX helpers. Thin zero-cost wrappers over the exact instruction forms the +// kernel needs — nothing generic, nothing speculative. +// ============================================================================ +namespace ptx { + +// Generic pointer -> 32-bit `.shared` window byte address. +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Local `.shared` address -> the same offset in CTA `cta_rank`'s smem +// (`.shared::cluster` instructions take this mapped form). +static __device__ __forceinline__ uint32_t mapa_shared_cluster( + uint32_t local_addr, uint32_t cta_rank) { + uint32_t mapped; + asm("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(mapped) : "r"(local_addr), "r"(cta_rank)); + return mapped; +} + +static __device__ __forceinline__ void bar_sync(uint32_t id, uint32_t count) { + asm volatile("barrier.cta.sync %0, %1;" :: "r"(id), "r"(count)); +} + +// Cluster barrier with release/acquire — the publish boundary between the +// mbarrier inits and any cross-CTA use of those mbarriers. +static __device__ __forceinline__ void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// %cluster_ctarank: 0 or 1 in this kernel's (2,1,1) clusters. +static __device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} + +static __device__ __forceinline__ uint32_t lane_id() { + uint32_t r; + asm("mov.u32 %0, %%laneid;" : "=r"(r)); + return r; +} + +// ---- mbarriers ------------------------------------------------------------- + +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + :: "r"(to_shared(bar)), "r"(count)); +} + +// Publish the inits to the async proxy + the peer CTA before any TMA +// completes on them. +static __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// Cross-CTA arrive with release ordering (teardown handshake). +static __device__ __forceinline__ void mbar_arrive_cluster_release( + uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" + :: "r"(mapped)); +} + +// Fused `tcgen05.wait::ld` + "buffer free" arrive, one asm block. The arrive +// must not become visible while TMEM loads are still in flight (it frees the +// accumulator buffer to the next tile's MMA — the bring-up chapter's +// drain overlap); the fused block +// keeps the arrive pinned behind the wait. (An earlier form also manufactured +// a register dependency on a drained value, `dep * 0` folded into the +// address; verified dead on nvcc 13.1 — ptxas const-folds it, SASS +// byte-identical without it — and retired.) +static __device__ __forceinline__ void tcgen05_wait_ld_then_arrive_release( + uint64_t* bar, uint32_t cta_rank, bool do_arrive) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile( + "{\n\t" + ".reg .pred %%pdo;\n\t" + "setp.ne.s32 %%pdo, %1, 0;\n\t" + "tcgen05.wait::ld.sync.aligned;\n\t" + "@%%pdo mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;\n\t" + "}\n\t" + :: "r"(mapped), "r"((int)do_arrive) : "memory"); +} + + +// Blocking parity wait at the consuming site (producer ring/tail waits + +// teardown). The 10 ms suspend hint is the polite form: a full-rate spin +// competes for issue slots and measurably loses; a nanosleep-backoff spin +// measured the same as this hint across five windows, so the simpler form +// ships. The carried one-stage-ahead try_wait peeks that used to gate these +// waits measured as pure overhead and were removed. +static __device__ __forceinline__ void mbar_wait_parity( + uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(to_shared(bar)), "r"(parity)); +} + +// ---- TMA ------------------------------------------------------------------- + +// Producer expect_tx arrive on the cluster-shared ring barrier: both +// peers' transfer bytes account against the even CTA's barrier, so one +// barrier per ring slot covers the whole cluster-wide gather. (Mechanism: +// `& 0xFEFFFFFF` clears bit 24 of the shared address — the hardware's +// peer-CTA bit — re-pointing the arrive at the even CTA's copy.) +static __device__ __forceinline__ void mbar_arrive_expect_tx_cluster( + uint64_t* bar, uint32_t bytes) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes)); +} + +// 3D tiled load, cta_group::2 + multicast (the A feed; the mask is the +// issuing CTA itself in a (2,1) cluster). Completion routes to the even +// CTA's mbarrier via the bit-24 clear. +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z) + : "memory"); +} + +// 3D tiled load, cta_group::2, no multicast (the B feed — each peer loads +// its own 128-wide N-block). +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_2sm_bit24( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst_smem), "l"(tmap), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) + : "memory"); +} + +// 4D tiled load, cta_group::2 + multicast (the scale-factor feed). The +// 512-byte SF tiles of one ring slot form a 4D box; `w` selects the slot. +static __device__ __forceinline__ void cp_async_bulk_tensor_4d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z2, int32_t w, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.4d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z2), "r"(w) + : "memory"); +} + + +// ---- tcgen05 (tensor memory + MMA plumbing) --------------------------------- + +static __device__ __forceinline__ void tcgen05_alloc_2sm( + uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_dealloc_2sm( + uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" + :: "r"(taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_relinquish_2sm() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); +} + +// The widest single-instruction TMEM read: 32 b32 registers per lane +// (blog: the epilogue — a quarter the load/wait count of the narrow forms). +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t taddr, uint32_t* r) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), + "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), + "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), + "=r"(r[16]), "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), + "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), + "=r"(r[24]), "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), + "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr)); +} + +// ---- stores + conversion ----------------------------------------------------- + +static __device__ __forceinline__ uint32_t cvt_pack_f16x2(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// Baseline output store: use the default cache policy. +static __device__ __forceinline__ void st_global_b128(uint4* ptr, uint4 v) { + asm volatile("st.global.v4.b32 [%0], {%1, %2, %3, %4};" + :: "l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)); +} + + +// First widen the baseline store while retaining its default cache policy. +static __device__ __forceinline__ void st_global_b256( + void* ptr, uint4 lo, uint4 hi) { + asm volatile("st.global.v8.b32 [%0], " + "{%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"(ptr), "r"(lo.x), "r"(lo.y), "r"(lo.z), "r"(lo.w), + "r"(hi.x), "r"(hi.y), "r"(hi.z), "r"(hi.w)); +} + +// ---- MMA descriptors --------------------------------------------------------- + +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class ScaleFormat : uint8_t { E4M3 = 0, E8M0 = 1 }; + +// Source descriptor for the tcgen05.cp that stages scale factors from smem +// into TMEM: unswizzled, SBO=128, LBO=16. The low word later carries the +// per-slot byte offset (in 16-byte units). +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_sf_cp_src_desc_cd(uint32_t sf_smem_addr) { + uint64_t addr_field = uint64_t((sf_smem_addr >> 4) & 0x3FC0u); + return addr_field | 0x400800010000ull; +} + +// Operand descriptor for a K=96 (48-byte) slice inside a 128-byte-swizzled +// smem K-window. OR-constant 0x4010404000000000: version=1, absolute-address +// LBO mode (bit 52), 128B swizzle, SBO=1024. Both the address field and the +// absolute-LBO field encode the window's own 1024-aligned base; the +// intra-window offset advances the address field only. +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k96_sw128(uint32_t window_base_1024aligned, int intra_window_off = 0) { + uint64_t addr_field = uint64_t(((window_base_1024aligned >> 4) & 0x7FC0u) + + uint32_t(intra_window_off >> 4)); + uint64_t lbo_field = uint64_t((window_base_1024aligned << 12) & 0x7FC00000u); + return addr_field | lbo_field | 0x4010404000000000ull; +} + +// Instruction descriptor for tcgen05.mma.kind::mxf4nvf4.block_scale. +// The K extent is selected by BIT 31, not by the asm mnemonic: k_dim=0 runs +// K=64 (0x10400480), k_dim=1 runs K=96 (0x90400480, sm_103a-only) — the +// cheapest quarter of the gap if your kernel is built on the K=64 form +// (blog: the MMA warp). Bit 13 flips which TMEM scale-factor column the MMA reads. +__host__ __device__ static __forceinline__ constexpr uint32_t +mma_inst_desc_mxf4nvf4_block16( + uint32_t M, uint32_t N, + uint32_t tmem_sfa_addr, uint32_t tmem_sfb_addr, + Major a_major, Major b_major, + bool negate_a, bool negate_b, + bool k_dim, ScaleFormat sf) { + constexpr uint32_t MXF4_E2M1 = 1u; + uint32_t d = 0; + d |= ((tmem_sfb_addr & 0xC0000000u) >> 30) << 4; // b_sf_id + d |= MXF4_E2M1 << 7; // a_format = E2M1 + d |= MXF4_E2M1 << 10; // b_format = E2M1 + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; // N in units of 8 + d |= (static_cast(sf) & 0x1u) << 23; // 0 = UE4M3 + d |= ((M >> 4) & 0x1Fu) << 24; // m_dim at bits 27-28 + d |= ((tmem_sfa_addr & 0xC0000000u) >> 30) << 29; // a_sf_id + if (k_dim) d |= 1u << 31; // 1 = K=96 + return d; +} + +} // namespace ptx + +// ============================================================================ +// Kernel geometry from the blog's opening kernel section. One tile shape, +// one cluster shape: +// 256x256 output per 2-CTA cluster (128 rows per CTA), K advanced in +// 96-element MMA steps staged as 768-element K-groups of 8 steps, persistent +// grid of 76 clusters = 152 CTAs = one per SM. +// ============================================================================ +namespace nvfp4 { + +constexpr int CTA_GROUP = 2; // cta_group::2 — the MMA spans both CTAs +constexpr int BLOCK_M = 128; // A rows per CTA +constexpr int BLOCK_N = 256; // output tile N extent (per cluster) +constexpr int BLOCK_N_PEER = BLOCK_N / CTA_GROUP; // 128 — B cols per CTA +constexpr int CLUSTER_M = CTA_GROUP * BLOCK_M; // 256 — tile M extent +[[maybe_unused]] constexpr int MMA_K = 96; // K per tcgen05.mma (idesc bit 31) +constexpr int WIN_BYTES = 128; // one 128B-swizzled smem K-window +constexpr int AB_LOAD_BYTES = WIN_BYTES; +constexpr int AB_TMA_SUBS = 3; +constexpr int TB_SIZE = 224; // 7 warps (blog: the warp table) +constexpr int N_EPI_WARPS = 4; // warps 0-3 + +constexpr int SF_TMEM_COLS = 9 * 4; // 3 SFA + 6 SFB cp tiles, 4 cols each +// Two 256-wide accumulators use a 220-column stride and overlap by the 36 +// columns needed for scale factors. +constexpr int D_STRIDE = 220; +constexpr int D_DB = 2; +constexpr int N_D_COLS = D_STRIDE + BLOCK_N; +constexpr int SF_TMEM_BASE = 476; +constexpr int N_ALLOC_COLS = 512; +static_assert(N_D_COLS <= SF_TMEM_BASE && + SF_TMEM_BASE + SF_TMEM_COLS <= N_ALLOC_COLS, + "accumulators + scale factors exceed the 512 TMEM columns"); +constexpr int EPI_CHUNK = 32; // cols per regular drain + +// ---- the shared-memory map (blog: the feed) ----------------------------------- +// Every offset is a compile-time expression. +constexpr int CD_HDR = 1024; // mbarrier header pad +constexpr int CD_AB_SLOTS = 6; +constexpr int CD_SF_SLOTS = 7; +constexpr int CD_AB_STRIDE = BLOCK_M * WIN_BYTES; // 16384 per A or B slot +constexpr int CD_SF_TILE = 512; // one (32 lanes x 16B) x4 tile +constexpr int CD_SFA_TILES = 3; // per SFA slot (M = 1 block) +constexpr int CD_SFB_TILES = 6; // per SFB slot (N = 2 blocks) +constexpr int CD_SFA_STRIDE = CD_SFA_TILES * CD_SF_TILE; // 1536 +constexpr int CD_SFB_STRIDE = CD_SFB_TILES * CD_SF_TILE; // 3072 + +// Expected transfer bytes are derived, never tuned: a wrong-by-one +// expectation doesn't crash, it deadlocks (blog: the feed). +constexpr uint32_t AB_TX = uint32_t(CTA_GROUP) * + uint32_t(BLOCK_M + BLOCK_N_PEER) * uint32_t(AB_LOAD_BYTES); +constexpr uint32_t SF_TX = uint32_t(CD_SFA_STRIDE + CD_SFB_STRIDE) * + uint32_t(CTA_GROUP); // 9216 +static_assert(AB_TX == 65536u && SF_TX == 9216u, + "expect_tx derivation drifted"); + +// The map IS the struct: member sizes and alignment produce every byte +// offset; the raw-address constants the mainloop uses read off it with +// offsetof, and the static_asserts pin the published values. +struct SmemCD { + // mbarrier header: four ring regions in handshake order, then the + // singleton accumulator + teardown barriers (alignas pads it to 1024). + struct alignas(1024) Header { + uint64_t ab_ready[CD_AB_SLOTS]; // TMA -> MMA + uint64_t ab_free [CD_AB_SLOTS]; // MMA -> TMA + uint64_t sf_ready[CD_SF_SLOTS]; // TMA -> MMA + uint64_t sf_free [CD_SF_SLOTS]; // MMA -> TMA + uint64_t acc_ready; // MMA -> epi + uint64_t acc_free; // epi -> MMA + uint64_t dealloc; // 2-CTA teardown + uint32_t tmem_addr; // TMEM base addr storage + } hdr; + uint8_t a [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t b [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t sfa[CD_SF_SLOTS][CD_SFA_STRIDE]; + // SFB starts 1024-aligned (the cp source descriptor drops address bits 4-9). + alignas(1024) uint8_t sfb[CD_SF_SLOTS][CD_SFB_STRIDE]; +}; +static_assert(sizeof(SmemCD::Header) == CD_HDR, "mbarrier header overflows the pad"); + +constexpr int CD_A_OFF = offsetof(SmemCD, a); +constexpr int CD_B_OFF = offsetof(SmemCD, b); +constexpr int CD_SFA_OFF = offsetof(SmemCD, sfa); +constexpr int CD_SFB_OFF = offsetof(SmemCD, sfb); +constexpr int CD_TOTAL = sizeof(SmemCD); +static_assert(CD_TOTAL <= 232448, "smem map exceeds the opt-in budget"); +static_assert(CD_AB_SLOTS != 6 || + (CD_B_OFF == 99328 && CD_SFA_OFF == 197632 && + CD_SFB_OFF == 208896 && CD_TOTAL == 230400), + "the 6-slot byte map drifted from its published values"); + +// The mainloop computes raw `.shared` byte addresses from these. +constexpr int MB_AB_READY = offsetof(SmemCD, hdr.ab_ready); +constexpr int MB_AB_FREE = offsetof(SmemCD, hdr.ab_free); +constexpr int MB_SF_READY = offsetof(SmemCD, hdr.sf_ready); +constexpr int MB_SF_FREE = offsetof(SmemCD, hdr.sf_free); +constexpr int MB_ACC_READY = offsetof(SmemCD, hdr.acc_ready); +constexpr int MB_ACC_FREE = offsetof(SmemCD, hdr.acc_free); + + + +} // namespace nvfp4 + +// ============================================================================ +// The consumer mainloop from the blog's pipeline section. One warp — warp 4 +// of the even CTA — +// issues everything: the tcgen05.cp copies that stage scale factors from the +// smem ring into TMEM, immediately followed by the tcgen05.mma stream, in the +// SAME warp. A same-warp tcgen05 stream is one ordered async pipeline, so the +// write-after-read hazard on the single reused scale-factor TMEM region needs +// no fence and no cross-warp barrier. +// +// Per 768-element K-group, logical K=96 steps 0..7 issue in K order. Their +// starting offsets, expressed as 16-byte descriptor atoms modulo a 128-byte +// swizzle window, are 0,3,6,1,4,7,2,5. Atoms 6 and 7 straddle windows. +// ============================================================================ +namespace nvfp4 { + + +// Ring/parity state carried across tiles by the MMA warp. +struct MainloopState { + uint32_t d_parity; // which TMEM accumulator buffer the next tile writes + uint32_t ab_phase, ab_slot; // A/B ring cursor + uint32_t sf_phase, sf_slot; // scale-factor ring cursor +}; + + +// Relaxed blocking wait (no .acquire): the tcgen05 pipeline's own ordering +// covers these ring waits, and the .acquire form perturbs the scoreboard +// cadence on the issuing warp. Waits sit AT the consuming site — the carried +// one-stage-ahead try_wait peeks that used to gate them measured as pure +// overhead and were removed. +static __device__ __forceinline__ void mbar_wait_relaxed( + uint32_t bar_smem, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(bar_smem), "r"(parity) : "memory"); +} + +// Gated form for the @216 buffer-free wait, which fires once per tile: the +// gate branches over the spin for every group past the first. +static __device__ __forceinline__ void mbar_wait_gated_relaxed( + uint32_t bar_smem, uint32_t parity, bool peek) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p bra DONE_%=;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t" + "DONE_%=:\n\t}\n" + :: "r"(bar_smem), "r"(parity), "r"((uint32_t)peek) : "memory"); +} + + +// One scale-factor staging copy, smem -> TMEM. Single-lane semantics behind +// the converged elect predicate: tcgen05.cp is a single-thread instruction — +// issued bare from a converged warp you get 32 initiations and garbage. +static __device__ __forceinline__ void cp_sf( + uint32_t el, uint32_t dst, uint32_t lo, uint32_t hi) { + const uint64_t desc = (uint64_t(hi) << 32) | uint64_t(lo); + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;\n\t}\n" + :: "r"(dst), "l"(desc), "r"(el) : "memory"); +} + +// Ring-slot "consumed" arrive, multicast to both CTAs of the pair. +// arrive::one must fire exactly once — hence the elect predicate. +static __device__ __forceinline__ void commit_consumed( + uint32_t el, uint32_t smbase, uint32_t mb_off, uint16_t mask) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t}\n" + :: "r"(smbase + mb_off), "h"(mask), "r"(el) : "memory"); +} + +// One K=96 NVFP4 MMA. The asm token says scale_vec::4X regardless — K=96 is +// selected by idesc bit 31 (blog: the MMA warp). acc_pred = 0 overwrites the +// accumulator (the tile's first step), nonzero accumulates. +static __device__ __forceinline__ void mma_cell( + uint32_t el, uint32_t d, + uint32_t desc_a_lo, uint32_t desc_a_hi, + uint32_t desc_b_lo, uint32_t desc_b_hi, + uint32_t idesc, uint32_t sf_a, uint32_t sf_b, uint32_t acc_pred) { + const uint64_t desc_a = (uint64_t(desc_a_hi) << 32) | uint64_t(desc_a_lo); + const uint64_t desc_b = (uint64_t(desc_b_hi) << 32) | uint64_t(desc_b_lo); + asm volatile( + "{\n\t.reg .pred p, acc;\n\t" + "setp.ne.u32 p, %7, 0;\n\t" + "setp.ne.b32 acc, %4, 0;\n\t" + "@p tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X" + " [%0], %1, %2, %3, [%5], [%6], acc;\n\t}\n" + :: "r"(d), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(acc_pred), "r"(sf_a), "r"(sf_b), "r"(el) : "memory"); +} + +// One SF ring slot -> TMEM: 9 tcgen05.cp tiles (3 SFA + 6 SFB). Source +// address advances are in 16-byte descriptor units. The SFB tile order +// 0,3,1,4,2,5 interleaves the two 128-wide N-blocks. +static __device__ __forceinline__ void cp_sf_buffer( + uint32_t el, uint32_t slot, + uint32_t sfa0, uint32_t sfa1, uint32_t sfa2, + uint32_t sfb0, uint32_t sfb1, uint32_t sfb2, + uint32_t sfb3, uint32_t sfb4, uint32_t sfb5, + uint32_t sfa_hi, uint32_t sfa_lo, uint32_t sfb_hi, uint32_t sfb_lo) { + constexpr uint32_t SFA_SLOT_U = CD_SFA_STRIDE / 16u; // 96 + constexpr uint32_t SFB_SLOT_U = CD_SFB_STRIDE / 16u; // 192 + constexpr uint32_t TILE_U = CD_SF_TILE / 16u; // 32 + const uint32_t a = slot * SFA_SLOT_U + sfa_lo; + cp_sf(el, sfa0, a, sfa_hi); + cp_sf(el, sfa1, a + 1u*TILE_U, sfa_hi); + cp_sf(el, sfa2, a + 2u*TILE_U, sfa_hi); + const uint32_t b = slot * SFB_SLOT_U + sfb_lo; + cp_sf(el, sfb0, b + 0u*TILE_U, sfb_hi); + cp_sf(el, sfb1, b + 3u*TILE_U, sfb_hi); + cp_sf(el, sfb2, b + 1u*TILE_U, sfb_hi); + cp_sf(el, sfb3, b + 4u*TILE_U, sfb_hi); + cp_sf(el, sfb4, b + 2u*TILE_U, sfb_hi); + cp_sf(el, sfb5, b + 5u*TILE_U, sfb_hi); +} + +// Operand descriptor LOW word for A/B ring slot `slot`, K-step `k_atom` +// (0..7 inside the 768-element group; each step is 48 bytes = +3 in the +// 16-byte address field). +static __device__ __forceinline__ uint32_t ab_desc_lo( + uint32_t smbase, uint32_t ab_base_off, uint32_t slot, uint32_t k_atom) { + const uint64_t d = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + slot * CD_AB_STRIDE, int(k_atom) * 16); + return uint32_t(d & 0xFFFFFFFFu); +} + +// K-steps 6 and 7 cross a 128-byte window boundary. Their descriptor address +// comes from `addr_slot`, while the absolute-LBO field comes from the next +// slot whose swizzle window contains the second half. +static __device__ __forceinline__ uint32_t ab_desc_lo_straddle( + uint32_t smbase, uint32_t ab_base_off, + uint32_t addr_slot, uint32_t lbo_slot, uint32_t k_atom) { + const uint64_t d_addr = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + addr_slot * CD_AB_STRIDE, int(k_atom) * 16); + const uint64_t d_lbo = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + lbo_slot * CD_AB_STRIDE, 0); + return uint32_t(d_addr & 0x0000FFFFu) | uint32_t(d_lbo & 0xFFFF0000u); +} + +// One output tile: `full_groups` complete 768-element K-groups. Coordinates +// past logical K are zero-filled by TMA from the unchanged canonical inputs. +static __device__ __forceinline__ MainloopState mainloop_tile( + uint32_t taddr, uint32_t smbase, MainloopState in, uint16_t commit_mask, + uint32_t full_groups + ) { + constexpr uint32_t AB_RING_SLOTS = uint32_t(CD_AB_SLOTS); + constexpr uint32_t SF_RING_SLOTS = uint32_t(CD_SF_SLOTS); + constexpr uint32_t AB_FREE_DELTA = uint32_t(MB_AB_FREE - MB_AB_READY); + constexpr uint32_t SF_READY_OFF = uint32_t(MB_SF_READY); + constexpr uint32_t SF_FREE_DELTA = uint32_t(MB_SF_FREE - MB_SF_READY); + constexpr uint32_t ACC_FREE_OFF = uint32_t(MB_ACC_FREE); + constexpr uint32_t A_OFF = uint32_t(CD_A_OFF); + constexpr uint32_t B_OFF = uint32_t(CD_B_OFF); + + uint32_t d_parity = in.d_parity; + uint32_t ab_phase = in.ab_phase, ab_slot = in.ab_slot; + uint32_t sf_phase = in.sf_phase, sf_slot = in.sf_slot; + + // The caller has already selected exactly one lane for the whole MMA role. + // Keep the helper predicate constant throughout the tile. + const uint32_t el = 1u; + + // Fixed TMEM targets for the single reused scale-factor region: 9 cp + // tiles of 4 columns each, starting past the accumulator region. + constexpr uint32_t SF = uint32_t(SF_TMEM_BASE); + const uint32_t sfa_t0 = taddr + SF + 0; // SFA word 0 (MMA reads here) + const uint32_t sfa_t1 = taddr + SF + 4; + const uint32_t sfa_t2 = taddr + SF + 8; + const uint32_t sfb_t0 = taddr + SF + 12; // SFB word 0 + const uint32_t sfb_t1 = taddr + SF + 16; + const uint32_t sfb_t2 = taddr + SF + 20; + const uint32_t sfb_t3 = taddr + SF + 24; + const uint32_t sfb_t4 = taddr + SF + 28; + const uint32_t sfb_t5 = taddr + SF + 32; + + // Scale-factor cp source descriptors (low word carries the per-slot + // offset; high word is constant). + const uint64_t sfb_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFB_OFF)); + const uint32_t sfb_src_lo = uint32_t(sfb_src & 0xFFFFFFFFu); + const uint32_t sfb_src_hi = uint32_t(sfb_src >> 32); + const uint64_t sfa_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFA_OFF)); + const uint32_t sfa_src_lo = uint32_t(sfa_src & 0xFFFFFFFFu); + const uint32_t sfa_src_hi = uint32_t(sfa_src >> 32); + + // Operand descriptor HIGH words (low words come from ab_desc_lo per step). + const uint32_t a_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + A_OFF) >> 32); + const uint32_t b_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + B_OFF) >> 32); + + // Each K=96 step needs 6 scale bytes per row, split across two staged + // 32-bit words. Odd-issue steps read word 1: bit 31 of the SF address + // selects the with-id encoding, and the id bits fold into the idesc. + constexpr uint32_t SF_ADDR_ID = 0x80000000u; + const uint32_t sfa_w1 = (taddr + SF + 4) | SF_ADDR_ID; + const uint32_t sfb_w1 = (taddr + SF + 20) | SF_ADDR_ID; + const uint32_t w1_id_bits = + ((sfa_w1 >> 1) & 0x60000000u) | ((sfb_w1 >> 26) & 48u); + + // The tile's first MMA overwrites the accumulator buffer this tile owns; + // the @216 "buffer free" wait runs once, before the first group. + const uint32_t acc_wait_parity = d_parity; + d_parity ^= 1u; + const uint32_t d_tmem = d_parity * uint32_t(D_STRIDE) + taddr; + + // Instruction descriptor: kind::mxf4nvf4, M=256, N=256, E2M1 x E2M1, + // UE4M3 block-16, K-major, k_dim=1 (K=96). Base value 0x90400480. + constexpr uint32_t IDESC = ptx::mma_inst_desc_mxf4nvf4_block16( + 256u, 256u, 0u, 0u, ptx::Major::K, ptx::Major::K, + false, false, /*k_dim=*/true, ptx::ScaleFormat::E4M3); + static_assert(IDESC == 0x90400480u, "K=96 NVFP4 idesc drifted"); + const uint32_t w0_id_bits = + ((sfa_t0 >> 1) & 0x60000000u) | ((sfb_t0 >> 26) & 48u); + const uint32_t idesc_w0 = w0_id_bits | IDESC; + const uint32_t idesc_w1 = w1_id_bits | IDESC; + + + uint32_t acc_pred = 0; // first MMA of the tile overwrites + bool need_acc_wait = true; // @216 wait fires once per tile + + if (full_groups >= 1) { + uint32_t g = 0; + do { + // ---- SF batch 0 -> logical steps 0,1 (atoms 0,3) ---- + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // ---- SF batch 1 -> logical steps 2,3 (atoms 6,1) ---- + mbar_wait_relaxed(sf_bar1, sf_phase1); + uint32_t sf_slot2, sf_phase2; + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + // AB0's last reader is MMA2. + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // ---- SF batch 2 -> logical steps 4,5 (atoms 4,7) ---- + mbar_wait_relaxed(sf_bar2, sf_phase2); + uint32_t sf_slot3, sf_phase3; + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; // ring cursor writeback + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // ---- SF batch 3 -> logical steps 6,7 (atoms 2,5) ---- + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; // ring cursor writeback + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + // AB1's last reader is MMA5; this measured placement leaves the + // commit behind the following scale copy. + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + acc_pred = 0xFFFFFFFFu; // accumulate from the second group on + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + need_acc_wait = false; + // AB2's last reader is MMA7. + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, + commit_mask); + ++g; + } while (g != full_groups); + } + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; +} + +} // namespace nvfp4 + +// ============================================================================ +// The kernel. 224 threads = 7 warps, four jobs (blog: the warp table): +// warps 0-3 epilogue: drain TMEM, convert to FP16, store straight to global +// warp 4 the MMA warp (even CTA of each pair only): SF copies + MMAs +// warp 5 TMA producer: A and B tiles into the 6-slot ring +// warp 6 TMA producer: scale factors into the 7-slot ring +// ============================================================================ +namespace nvfp4 { + +// Baseline launch: plain launch bounds and a grid expressed in CTA units. +#define NVFP4_BLOCK_ATTR +constexpr bool GRID_IN_CLUSTERS = false; + +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void nvfp4_gemm_kernel(const __grid_constant__ CUtensorMap A_tmap, + const __grid_constant__ CUtensorMap B_tmap, + const __grid_constant__ CUtensorMap SFA_tmap, + const __grid_constant__ CUtensorMap SFB_tmap, + __half* C_ptr, int M, int N, int K + ) { + extern __shared__ __align__(1024) char smem_buf[]; + SmemCD& smem = *reinterpret_cast(smem_buf); + + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const uint32_t crank = ptx::cluster_rank(); + const uint32_t m_pair = crank & 1u; // 0 = the pair's even CTA + + const int cluster_grid_m = (M + CLUSTER_M - 1) / CLUSTER_M; + const int grid_n = (N + BLOCK_N - 1) / BLOCK_N; + const int total_tiles = cluster_grid_m * grid_n; + + // Baseline rungs execute a whole K=768 group at the edge. Tensor-map OOB + // fill supplies zeros past logical K without changing the input buffers. + const int full_groups = (K + 767) / 768; + const int num_groups = full_groups; + + // Persistent cluster index. The device always sees the grid in CTA units + // (2, 1, clusters) regardless of the host-side launch-unit regime. + const bool grid_z = (gridDim.z > 1); + const int num_clusters = grid_z ? int(gridDim.z) : int(gridDim.x) / 2; + const int cluster_id = grid_z ? int(blockIdx.z) : int(blockIdx.x) / 2; + + + // Which output tile does persistent work item `t` process? + auto tile_of = [&](int t) -> int { + return t; // natural raster: M fast, N slow + }; + + auto producer_tile_of = [&](int t) -> int { + return tile_of(t); + }; + + // ---- mbarrier init (warp 0, lane 0) ---- + if (warp_id == 0 && ptx::lane_id() == 0) { + // Arrival counts are derived from the topology, never tuned: + // "ready" = one producer expect_tx per slot (bit-24 routes both + // peers' bytes to the even CTA's barrier); "free" = one consumer + // commit per slot reuse. + #pragma unroll + for (int s = 0; s < CD_AB_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.ab_ready[s], 1); + ptx::mbar_init(&smem.hdr.ab_free[s], 1); + } + #pragma unroll + for (int s = 0; s < CD_SF_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.sf_ready[s], 1); + ptx::mbar_init(&smem.hdr.sf_free[s], 1); + } + ptx::mbar_init(&smem.hdr.acc_ready, 1); + ptx::mbar_init(&smem.hdr.acc_free, N_EPI_WARPS * CTA_GROUP); // 4 x 2 = 8 + ptx::mbar_init(&smem.hdr.dealloc, 32); // one warp's lanes, cross-peer + ptx::fence_mbarrier_init_release_cluster(); + } + ptx::cluster_sync_rel_acq(); + + // TMEM allocation: epilogue warp 0 allocates for the pair; warps 0-4 + // rendezvous before reading the granted base. The TMA producer warps + // never touch TMEM and proceed straight to their loops. + uint32_t taddr = 0; + if (warp_id <= 4) { + if (warp_id == 0) + ptx::tcgen05_alloc_2sm(ptx::to_shared(&smem.hdr.tmem_addr), + N_ALLOC_COLS); + ptx::bar_sync(2, 160); + taddr = smem.hdr.tmem_addr; + } + + const int wg = warp_id >> 2; // 0 = epilogue warps, 1 = producers + MMA + if (wg == 1) { + const int pw = warp_id - 4; + if (pw == 1 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 5: lane 0 owns the complete A/B TMA producer loop. + // ================================================================ + // free_phase starts at 1: a fresh mbarrier is phase 0, so the + // first ring's worth of "free" waits falls through. + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + // The pair splits the 256x256 tile: even CTA holds A rows + // [0,128) and B cols [0,128); odd CTA the other halves. + const int m_blk = m_row * CTA_GROUP + int(m_pair); + const int n_blk = n_grp * CTA_GROUP + int(m_pair); + #pragma unroll 1 + for (int g = 0; g < num_groups; ++g) { + // Three full 128-byte windows cover the packed K=768 group. + #pragma unroll + for (int sub = 0; sub < AB_TMA_SUBS; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.ab_free[slot], free_phase); + const uint32_t a_s = + base + CD_A_OFF + uint32_t(slot) * CD_AB_STRIDE; + const uint32_t b_s = + base + CD_B_OFF + uint32_t(slot) * CD_AB_STRIDE; + uint64_t* bar = &smem.hdr.ab_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, AB_TX); + // A rides the multicast form with a self mask so the + // completion routes to the even peer's barrier. + const uint16_t mask_a = uint16_t(1u << m_pair); + // Read the same row-major FP4 allocation as cuBLAS. + // K192 uses four 96-byte boxes; K256 uses three + // 128-byte boxes. The tensor map zero-fills past K. + const CUtensorMap* a_map = &A_tmap; + const CUtensorMap* b_map = &B_tmap; + const int x = g * 384 + sub * AB_LOAD_BYTES; + const int y_a = m_blk * BLOCK_M; + const int y_b = n_blk * BLOCK_N_PEER; + const int z_a = 0; + const int z_b = 0; + ptx::cp_async_bulk_tensor_3d_load_multicast( + a_s, a_map, x, y_a, z_a, + bar, mask_a + ); + ptx::cp_async_bulk_tensor_3d_load_2sm_bit24( + b_s, b_map, x, y_b, z_b, + bar + ); + if (++slot == CD_AB_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + // Drain: wait for the last-filled slot to be consumed, then hand + // it one clean-shutdown expect_tx. + { + const int t_slot = (slot == 0) ? CD_AB_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.ab_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.ab_ready[t_slot], AB_TX); + } + } else if (pw == 2 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 6: scale-factor TMA producer. SFA is per-peer (each CTA's + // own M half); SFB is the one genuinely multicast operand — the + // cta_group::2 copy reads BOTH peers' smem as one logical + // operand, so both copies must be byte-identical (blog: the scale-factor section). + // ================================================================ + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const int m_blk = m_row * CTA_GROUP + int(m_pair); + for (int g = 0; g < num_groups; ++g) { + // 4 SF ring slots per 768-K group (2 K=96 steps each). + #pragma unroll + for (int sub = 0; sub < 4; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.sf_free[slot], free_phase); + const uint32_t sfa_s = + base + CD_SFA_OFF + uint32_t(slot) * CD_SFA_STRIDE; + const uint32_t sfb_s = + base + CD_SFB_OFF + uint32_t(slot) * CD_SFB_STRIDE; + uint64_t* bar = &smem.hdr.sf_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, SF_TX); + // VEC16 is [outer-block][K-group-of-4][128 rows x 4]. + // Three adjacent groups are the 128 x 12 scale slot. + const int sf_group = (g * 4 + sub) * 3; + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfa_s, &SFA_tmap, 0, 0, sf_group, m_blk, bar, + uint16_t(1u << m_pair) + ); + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfb_s + m_pair * uint32_t(CD_SFB_STRIDE / 2), + &SFB_tmap, 0, 0, sf_group, + n_grp * CTA_GROUP + int(m_pair), bar, + uint16_t(0x3) + ); + if (++slot == CD_SF_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + { + const int t_slot = (slot == 0) ? CD_SF_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.sf_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.sf_ready[t_slot], SF_TX); + } + } else if (pw == 0 && m_pair == 0 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 4: the MMA warp — issued once per cluster, on the even + // CTA. The odd CTA's warp 4 does nothing all kernel (blog: the warp table). + // The mainloop needs no tile coordinates: operands arrive through + // the rings; this warp only needs its share of the trip count. + // ================================================================ + const uint32_t smbase = ptx::to_shared(&smem); + const uint16_t commit_mask = 0x3; // both CTAs of this pair + MainloopState st{1u, 0u, 0u, 0u, 0u}; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + st = mainloop_tile(taddr, smbase, st, commit_mask, + uint32_t(full_groups) + ); + // The outer role branch already leaves exactly one issuer. + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;" + :: "r"(smbase + uint32_t(MB_ACC_READY)), + "h"(uint16_t(0x3)) : "memory"); + } + // Drain the final tile's buffer-free arrive before exiting. + mbar_wait_relaxed(smbase + uint32_t(MB_ACC_FREE), st.d_parity); + } + } else { + // ==================================================================== + // Warps 0-3: the epilogue (blog: the epilogue). Each warp owns a 32-row band + // and drains a tile as eight 32-column tcgen05.ld loads, converts + // FP32->FP16 in packed pairs, and stores straight to global — no + // shared-memory bounce. + // ==================================================================== + const int epi_warp = warp_id; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + const uint32_t acc_free_rank = crank & ~1u; // the pair's even CTA + int d_db = 0; + int acc_ready_phase = 0; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const uint32_t tmem_base = taddr + uint32_t(d_db) * D_STRIDE; + mbar_wait_relaxed(ptx::to_shared(&smem.hdr.acc_ready), + uint32_t(acc_ready_phase)); + acc_ready_phase ^= 1; + + const int off_m = (m_row * CTA_GROUP + int(m_pair)) * BLOCK_M; + const int off_n = n_grp * BLOCK_N; + const int row = off_m + epi_warp * 32 + lane_id; + const bool row_in = (row < M); + constexpr int EPI_BANDS = BLOCK_N / EPI_CHUNK; // 8 + #pragma unroll 1 + for (int k = 0; k < EPI_BANDS; ++k) { + // Buffer 0 drains high-to-low and buffer 1 low-to-high, so + // their shared columns [220,256) drain first for both. + const int parity = d_db & 1; + const int band = (parity == 0) ? (EPI_BANDS - 1 - k) : k; + const uint32_t taddr_n = + tmem_base + uint32_t(band) * 32 + taddr_lane; + uint32_t w[32]; + ptx::tcgen05_ld_32x32b_x32(taddr_n, w); + // Return the old buffer as soon as its shared edge is out. + if (k == 1) { + ptx::tcgen05_wait_ld_then_arrive_release( + &smem.hdr.acc_free, acc_free_rank, + /*do_arrive=*/ptx::lane_id() == 0); + } + if (!row_in) continue; + const int col_base = off_n + band * 32; + uint32_t h[16]; + #pragma unroll + for (int j = 0; j < 16; ++j) { + h[j] = ptx::cvt_pack_f16x2(__int_as_float(w[2 * j + 1]), + __int_as_float(w[2 * j])); + } + const size_t row_base = size_t(row) * size_t(N); + const bool aligned32 = ((row_base & 15u) == 0u) && + ((uintptr_t(C_ptr) & 31u) == 0u); + if (aligned32 && (col_base & 15) == 0) { + // Precondition here: row_base is 16-element (32-byte) + // aligned — the aligned32 guard owns it. Restating it to + // ptxas via __builtin_assume compiled this file + // SASS-identical (and measured perf-null on the + // production kernel), so the guard stands alone. + #pragma unroll + for (int q = 0; q < 2; ++q) { + const int c0 = col_base + q * 16; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 16 <= N) { + uint4 lo{h[q*8+0], h[q*8+1], h[q*8+2], h[q*8+3]}; + uint4 hi{h[q*8+4], h[q*8+5], h[q*8+6], h[q*8+7]}; + ptx::st_global_b256(cdst, lo, hi); + } else { + // N tail: 128-bit then scalar. + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int cq = c0 + p * 8; + if (cq + 8 <= N) { + uint4 v{h[q*8+p*4+0], h[q*8+p*4+1], + h[q*8+p*4+2], h[q*8+p*4+3]}; + ptx::st_global_b128( + reinterpret_cast(cdst + p * 8), v); + } else { + const uint16_t* hb = + reinterpret_cast( + &h[q * 8 + p * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (cq + c < N) cdst[p * 8 + c] = hb[c]; + } + } + } + } + } else { + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int c0 = col_base + q * 8; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 8 <= N) { + uint4 v{h[q*4+0], h[q*4+1], h[q*4+2], h[q*4+3]}; + ptx::st_global_b128(reinterpret_cast(cdst), v); + } else { + const uint16_t* hb = + reinterpret_cast(&h[q * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (c0 + c < N) cdst[c] = hb[c]; + } + } + } + } + if (++d_db == D_DB) d_db = 0; + } + } + + // Teardown: only the epilogue warps rendezvous per CTA; warp 0 then runs + // the cross-peer dealloc handshake. Producer and MMA warps exit freely — + // their drains above already order their work. + if (wg == 0) { + ptx::bar_sync(3, 128); + if (warp_id == 0) { + ptx::tcgen05_relinquish_2sm(); + const uint32_t peer_rank = crank ^ 1u; + // All 32 lanes arrive the PEER's count-32 barrier = one flip. + ptx::mbar_arrive_cluster_release(&smem.hdr.dealloc, peer_rank); + ptx::mbar_wait_parity(&smem.hdr.dealloc, 0); + ptx::tcgen05_dealloc_2sm(taddr, N_ALLOC_COLS); + } + } +} + + +// Keep the host launch convention paired with the kernel's block attribute. +static inline dim3 launch_grid(int clusters) { + return GRID_IN_CLUSTERS ? dim3(1, 1, unsigned(clusters)) + : dim3(2, 1, unsigned(clusters)); +} + +} // namespace nvfp4 + +// ============================================================================ +// Host-side TMA descriptor support. The shared input and benchmark code lives +// in main.cu. +// ============================================================================ +namespace tmap { + +// Thin wrappers over cuTensorMapEncodeTiled. Dim 0 is the innermost +// (stride-1) axis; strides for dim 1+ are in bytes. +inline CUtensorMap encode_tiled_3d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_depth, uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t depth_stride_bytes, + uint32_t box_depth, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[3] = { global_cols, global_rows, global_depth }; + cuuint64_t global_strides[2] = { row_stride_bytes, depth_stride_bytes }; + cuuint32_t box_dim[3] = { box_cols, box_rows, box_depth }; + cuuint32_t element_strides[3] = { 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 3, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +inline CUtensorMap encode_tiled_4d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_slots, uint64_t global_tiles, + uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t tile_stride_bytes, + uint64_t slot_stride_bytes, + uint32_t box_tiles, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[4] = { global_cols, global_rows, global_tiles, + global_slots }; + cuuint64_t global_strides[3] = { row_stride_bytes, tile_stride_bytes, + slot_stride_bytes }; + cuuint32_t box_dim[4] = { box_cols, box_rows, box_tiles, 1u }; + cuuint32_t element_strides[4] = { 1, 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 4, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +} // namespace tmap + +namespace host { + +using namespace nvfp4; + +// ---- implementation-specific tensor maps --------------------------------- + +inline CUtensorMap make_ab_tmap(void* ptr, int rows, int K, + int row_stride) { + return tmap::encode_tiled_3d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*depth=*/1, uint64_t(rows), /*cols=*/uint64_t(K / 2), + /*row_stride=*/uint64_t(row_stride), + /*depth_stride=*/uint64_t(rows) * uint64_t(row_stride), + /*box=*/1, BLOCK_M, AB_LOAD_BYTES, CU_TENSOR_MAP_SWIZZLE_128B); +} + +inline CUtensorMap make_sf_tmap(uint8_t* ptr, int outer, int K) { + const uint64_t sf_inner = (uint64_t(K) + 63) / 64 * 4; + const uint64_t outer_blocks = (uint64_t(outer) + 127) / 128; + return tmap::encode_tiled_4d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*slots=*/outer_blocks, /*tiles=*/sf_inner / 4, + /*rows=*/4, /*cols=*/128, + /*row_stride=*/128, /*tile_stride=*/512, + /*slot_stride=*/sf_inner * 128, + /*box_tiles=*/3, /*box_rows=*/4, /*box_cols=*/128, + CU_TENSOR_MAP_SWIZZLE_NONE); +} +} // namespace host diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/gemm5.cuh b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm5.cuh new file mode 100644 index 000000000..7faa1e892 --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm5.cuh @@ -0,0 +1,1263 @@ +#pragma once +// r5: add the output cache policy. +// This is a complete, directly editable ladder snapshot. +// It intentionally duplicates the implementation in the other gemm*.cuh files. + +// Conditional branches are already materialized; edit this implementation directly. +#define NVFP4_BUILD_NAME "r5" +// Read-only metadata for main.cu; these do not select code in this header. +#define NVFP4_HAS_K64_TAIL 0 +#define NVFP4_HAS_SCHEDULE 0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(expr) do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \ + cudaGetErrorName(_e), __FILE__, __LINE__, \ + cudaGetErrorString(_e)); \ + std::abort(); \ + } \ +} while (0) + +#define CU_CHECK(expr) do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + std::fprintf(stderr, "CU error at %s:%d: %s\n", \ + __FILE__, __LINE__, _s ? _s : "?"); \ + std::abort(); \ + } \ +} while (0) + + +// ============================================================================ +// PTX helpers. Thin zero-cost wrappers over the exact instruction forms the +// kernel needs — nothing generic, nothing speculative. +// ============================================================================ +namespace ptx { + +// Generic pointer -> 32-bit `.shared` window byte address. +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Local `.shared` address -> the same offset in CTA `cta_rank`'s smem +// (`.shared::cluster` instructions take this mapped form). +static __device__ __forceinline__ uint32_t mapa_shared_cluster( + uint32_t local_addr, uint32_t cta_rank) { + uint32_t mapped; + asm("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(mapped) : "r"(local_addr), "r"(cta_rank)); + return mapped; +} + +static __device__ __forceinline__ void bar_sync(uint32_t id, uint32_t count) { + asm volatile("barrier.cta.sync %0, %1;" :: "r"(id), "r"(count)); +} + +// Cluster barrier with release/acquire — the publish boundary between the +// mbarrier inits and any cross-CTA use of those mbarriers. +static __device__ __forceinline__ void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// %cluster_ctarank: 0 or 1 in this kernel's (2,1,1) clusters. +static __device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} + +static __device__ __forceinline__ uint32_t lane_id() { + uint32_t r; + asm("mov.u32 %0, %%laneid;" : "=r"(r)); + return r; +} + +// ---- mbarriers ------------------------------------------------------------- + +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + :: "r"(to_shared(bar)), "r"(count)); +} + +// Publish the inits to the async proxy + the peer CTA before any TMA +// completes on them. +static __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// Cross-CTA arrive with release ordering (teardown handshake). +static __device__ __forceinline__ void mbar_arrive_cluster_release( + uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" + :: "r"(mapped)); +} + +// Fused `tcgen05.wait::ld` + "buffer free" arrive, one asm block. The arrive +// must not become visible while TMEM loads are still in flight (it frees the +// accumulator buffer to the next tile's MMA — the bring-up chapter's +// drain overlap); the fused block +// keeps the arrive pinned behind the wait. (An earlier form also manufactured +// a register dependency on a drained value, `dep * 0` folded into the +// address; verified dead on nvcc 13.1 — ptxas const-folds it, SASS +// byte-identical without it — and retired.) +static __device__ __forceinline__ void tcgen05_wait_ld_then_arrive_release( + uint64_t* bar, uint32_t cta_rank, bool do_arrive) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile( + "{\n\t" + ".reg .pred %%pdo;\n\t" + "setp.ne.s32 %%pdo, %1, 0;\n\t" + "tcgen05.wait::ld.sync.aligned;\n\t" + "@%%pdo mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;\n\t" + "}\n\t" + :: "r"(mapped), "r"((int)do_arrive) : "memory"); +} + + +// Blocking parity wait at the consuming site (producer ring/tail waits + +// teardown). The 10 ms suspend hint is the polite form: a full-rate spin +// competes for issue slots and measurably loses; a nanosleep-backoff spin +// measured the same as this hint across five windows, so the simpler form +// ships. The carried one-stage-ahead try_wait peeks that used to gate these +// waits measured as pure overhead and were removed. +static __device__ __forceinline__ void mbar_wait_parity( + uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(to_shared(bar)), "r"(parity)); +} + +// ---- TMA ------------------------------------------------------------------- + +// Producer expect_tx arrive on the cluster-shared ring barrier: both +// peers' transfer bytes account against the even CTA's barrier, so one +// barrier per ring slot covers the whole cluster-wide gather. (Mechanism: +// `& 0xFEFFFFFF` clears bit 24 of the shared address — the hardware's +// peer-CTA bit — re-pointing the arrive at the even CTA's copy.) +static __device__ __forceinline__ void mbar_arrive_expect_tx_cluster( + uint64_t* bar, uint32_t bytes) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes)); +} + +// 3D tiled load, cta_group::2 + multicast (the A feed; the mask is the +// issuing CTA itself in a (2,1) cluster). Completion routes to the even +// CTA's mbarrier via the bit-24 clear. +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z) + : "memory"); +} + +// 3D tiled load, cta_group::2, no multicast (the B feed — each peer loads +// its own 128-wide N-block). +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_2sm_bit24( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst_smem), "l"(tmap), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) + : "memory"); +} + +// 4D tiled load, cta_group::2 + multicast (the scale-factor feed). The +// 512-byte SF tiles of one ring slot form a 4D box; `w` selects the slot. +static __device__ __forceinline__ void cp_async_bulk_tensor_4d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z2, int32_t w, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.4d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z2), "r"(w) + : "memory"); +} + + +// ---- tcgen05 (tensor memory + MMA plumbing) --------------------------------- + +static __device__ __forceinline__ void tcgen05_alloc_2sm( + uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_dealloc_2sm( + uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" + :: "r"(taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_relinquish_2sm() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); +} + +// The widest single-instruction TMEM read: 32 b32 registers per lane +// (blog: the epilogue — a quarter the load/wait count of the narrow forms). +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t taddr, uint32_t* r) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), + "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), + "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), + "=r"(r[16]), "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), + "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), + "=r"(r[24]), "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), + "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr)); +} + +// ---- stores + conversion ----------------------------------------------------- + +static __device__ __forceinline__ uint32_t cvt_pack_f16x2(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// L2::evict_first requires a 256-bit store, but the 128-bit tail can still +// adopt the no-L1-allocation part of the policy at this rung. +static __device__ __forceinline__ void st_global_na_b128( + uint4* ptr, uint4 v) { + asm volatile("st.global.L1::no_allocate.v4.b32 " + "[%0], {%1, %2, %3, %4};" + :: "l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)); +} + +// Add the output cache policy to the 256-bit store: bypass L1 and evict the +// write-once line first from L2 so it does not displace reusable operands. +static __device__ __forceinline__ void st_global_na_ef_b256( + void* ptr, uint4 lo, uint4 hi) { + asm volatile("st.global.L1::no_allocate.L2::evict_first.v8.b32 [%0], " + "{%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"(ptr), "r"(lo.x), "r"(lo.y), "r"(lo.z), "r"(lo.w), + "r"(hi.x), "r"(hi.y), "r"(hi.z), "r"(hi.w)); +} + + +// ---- MMA descriptors --------------------------------------------------------- + +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class ScaleFormat : uint8_t { E4M3 = 0, E8M0 = 1 }; + +// Source descriptor for the tcgen05.cp that stages scale factors from smem +// into TMEM: unswizzled, SBO=128, LBO=16. The low word later carries the +// per-slot byte offset (in 16-byte units). +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_sf_cp_src_desc_cd(uint32_t sf_smem_addr) { + uint64_t addr_field = uint64_t((sf_smem_addr >> 4) & 0x3FC0u); + return addr_field | 0x400800010000ull; +} + +// Operand descriptor for a K=96 (48-byte) slice inside a 128-byte-swizzled +// smem K-window. OR-constant 0x4010404000000000: version=1, absolute-address +// LBO mode (bit 52), 128B swizzle, SBO=1024. Both the address field and the +// absolute-LBO field encode the window's own 1024-aligned base; the +// intra-window offset advances the address field only. +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k96_sw128(uint32_t window_base_1024aligned, int intra_window_off = 0) { + uint64_t addr_field = uint64_t(((window_base_1024aligned >> 4) & 0x7FC0u) + + uint32_t(intra_window_off >> 4)); + uint64_t lbo_field = uint64_t((window_base_1024aligned << 12) & 0x7FC00000u); + return addr_field | lbo_field | 0x4010404000000000ull; +} + +// Instruction descriptor for tcgen05.mma.kind::mxf4nvf4.block_scale. +// The K extent is selected by BIT 31, not by the asm mnemonic: k_dim=0 runs +// K=64 (0x10400480), k_dim=1 runs K=96 (0x90400480, sm_103a-only) — the +// cheapest quarter of the gap if your kernel is built on the K=64 form +// (blog: the MMA warp). Bit 13 flips which TMEM scale-factor column the MMA reads. +__host__ __device__ static __forceinline__ constexpr uint32_t +mma_inst_desc_mxf4nvf4_block16( + uint32_t M, uint32_t N, + uint32_t tmem_sfa_addr, uint32_t tmem_sfb_addr, + Major a_major, Major b_major, + bool negate_a, bool negate_b, + bool k_dim, ScaleFormat sf) { + constexpr uint32_t MXF4_E2M1 = 1u; + uint32_t d = 0; + d |= ((tmem_sfb_addr & 0xC0000000u) >> 30) << 4; // b_sf_id + d |= MXF4_E2M1 << 7; // a_format = E2M1 + d |= MXF4_E2M1 << 10; // b_format = E2M1 + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; // N in units of 8 + d |= (static_cast(sf) & 0x1u) << 23; // 0 = UE4M3 + d |= ((M >> 4) & 0x1Fu) << 24; // m_dim at bits 27-28 + d |= ((tmem_sfa_addr & 0xC0000000u) >> 30) << 29; // a_sf_id + if (k_dim) d |= 1u << 31; // 1 = K=96 + return d; +} + +} // namespace ptx + +// ============================================================================ +// Kernel geometry from the blog's opening kernel section. One tile shape, +// one cluster shape: +// 256x256 output per 2-CTA cluster (128 rows per CTA), K advanced in +// 96-element MMA steps staged as 768-element K-groups of 8 steps, persistent +// grid of 76 clusters = 152 CTAs = one per SM. +// ============================================================================ +namespace nvfp4 { + +constexpr int CTA_GROUP = 2; // cta_group::2 — the MMA spans both CTAs +constexpr int BLOCK_M = 128; // A rows per CTA +constexpr int BLOCK_N = 256; // output tile N extent (per cluster) +constexpr int BLOCK_N_PEER = BLOCK_N / CTA_GROUP; // 128 — B cols per CTA +constexpr int CLUSTER_M = CTA_GROUP * BLOCK_M; // 256 — tile M extent +[[maybe_unused]] constexpr int MMA_K = 96; // K per tcgen05.mma (idesc bit 31) +constexpr int WIN_BYTES = 128; // one 128B-swizzled smem K-window +constexpr int AB_LOAD_BYTES = WIN_BYTES; +constexpr int AB_TMA_SUBS = 3; +constexpr int TB_SIZE = 224; // 7 warps (blog: the warp table) +constexpr int N_EPI_WARPS = 4; // warps 0-3 + +constexpr int SF_TMEM_COLS = 9 * 4; // 3 SFA + 6 SFB cp tiles, 4 cols each +// Two 256-wide accumulators use a 220-column stride and overlap by the 36 +// columns needed for scale factors. +constexpr int D_STRIDE = 220; +constexpr int D_DB = 2; +constexpr int N_D_COLS = D_STRIDE + BLOCK_N; +constexpr int SF_TMEM_BASE = 476; +constexpr int N_ALLOC_COLS = 512; +static_assert(N_D_COLS <= SF_TMEM_BASE && + SF_TMEM_BASE + SF_TMEM_COLS <= N_ALLOC_COLS, + "accumulators + scale factors exceed the 512 TMEM columns"); +constexpr int EPI_CHUNK = 32; // cols per regular drain + +// ---- the shared-memory map (blog: the feed) ----------------------------------- +// Every offset is a compile-time expression. +constexpr int CD_HDR = 1024; // mbarrier header pad +constexpr int CD_AB_SLOTS = 6; +constexpr int CD_SF_SLOTS = 7; +constexpr int CD_AB_STRIDE = BLOCK_M * WIN_BYTES; // 16384 per A or B slot +constexpr int CD_SF_TILE = 512; // one (32 lanes x 16B) x4 tile +constexpr int CD_SFA_TILES = 3; // per SFA slot (M = 1 block) +constexpr int CD_SFB_TILES = 6; // per SFB slot (N = 2 blocks) +constexpr int CD_SFA_STRIDE = CD_SFA_TILES * CD_SF_TILE; // 1536 +constexpr int CD_SFB_STRIDE = CD_SFB_TILES * CD_SF_TILE; // 3072 + +// Expected transfer bytes are derived, never tuned: a wrong-by-one +// expectation doesn't crash, it deadlocks (blog: the feed). +constexpr uint32_t AB_TX = uint32_t(CTA_GROUP) * + uint32_t(BLOCK_M + BLOCK_N_PEER) * uint32_t(AB_LOAD_BYTES); +constexpr uint32_t SF_TX = uint32_t(CD_SFA_STRIDE + CD_SFB_STRIDE) * + uint32_t(CTA_GROUP); // 9216 +static_assert(AB_TX == 65536u && SF_TX == 9216u, + "expect_tx derivation drifted"); + +// The map IS the struct: member sizes and alignment produce every byte +// offset; the raw-address constants the mainloop uses read off it with +// offsetof, and the static_asserts pin the published values. +struct SmemCD { + // mbarrier header: four ring regions in handshake order, then the + // singleton accumulator + teardown barriers (alignas pads it to 1024). + struct alignas(1024) Header { + uint64_t ab_ready[CD_AB_SLOTS]; // TMA -> MMA + uint64_t ab_free [CD_AB_SLOTS]; // MMA -> TMA + uint64_t sf_ready[CD_SF_SLOTS]; // TMA -> MMA + uint64_t sf_free [CD_SF_SLOTS]; // MMA -> TMA + uint64_t acc_ready; // MMA -> epi + uint64_t acc_free; // epi -> MMA + uint64_t dealloc; // 2-CTA teardown + uint32_t tmem_addr; // TMEM base addr storage + } hdr; + uint8_t a [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t b [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t sfa[CD_SF_SLOTS][CD_SFA_STRIDE]; + // SFB starts 1024-aligned (the cp source descriptor drops address bits 4-9). + alignas(1024) uint8_t sfb[CD_SF_SLOTS][CD_SFB_STRIDE]; +}; +static_assert(sizeof(SmemCD::Header) == CD_HDR, "mbarrier header overflows the pad"); + +constexpr int CD_A_OFF = offsetof(SmemCD, a); +constexpr int CD_B_OFF = offsetof(SmemCD, b); +constexpr int CD_SFA_OFF = offsetof(SmemCD, sfa); +constexpr int CD_SFB_OFF = offsetof(SmemCD, sfb); +constexpr int CD_TOTAL = sizeof(SmemCD); +static_assert(CD_TOTAL <= 232448, "smem map exceeds the opt-in budget"); +static_assert(CD_AB_SLOTS != 6 || + (CD_B_OFF == 99328 && CD_SFA_OFF == 197632 && + CD_SFB_OFF == 208896 && CD_TOTAL == 230400), + "the 6-slot byte map drifted from its published values"); + +// The mainloop computes raw `.shared` byte addresses from these. +constexpr int MB_AB_READY = offsetof(SmemCD, hdr.ab_ready); +constexpr int MB_AB_FREE = offsetof(SmemCD, hdr.ab_free); +constexpr int MB_SF_READY = offsetof(SmemCD, hdr.sf_ready); +constexpr int MB_SF_FREE = offsetof(SmemCD, hdr.sf_free); +constexpr int MB_ACC_READY = offsetof(SmemCD, hdr.acc_ready); +constexpr int MB_ACC_FREE = offsetof(SmemCD, hdr.acc_free); + + + +} // namespace nvfp4 + +// ============================================================================ +// The consumer mainloop from the blog's pipeline section. One warp — warp 4 +// of the even CTA — +// issues everything: the tcgen05.cp copies that stage scale factors from the +// smem ring into TMEM, immediately followed by the tcgen05.mma stream, in the +// SAME warp. A same-warp tcgen05 stream is one ordered async pipeline, so the +// write-after-read hazard on the single reused scale-factor TMEM region needs +// no fence and no cross-warp barrier. +// +// Per 768-element K-group, logical K=96 steps 0..7 issue in K order. Their +// starting offsets, expressed as 16-byte descriptor atoms modulo a 128-byte +// swizzle window, are 0,3,6,1,4,7,2,5. Atoms 6 and 7 straddle windows. +// ============================================================================ +namespace nvfp4 { + + +// Ring/parity state carried across tiles by the MMA warp. +struct MainloopState { + uint32_t d_parity; // which TMEM accumulator buffer the next tile writes + uint32_t ab_phase, ab_slot; // A/B ring cursor + uint32_t sf_phase, sf_slot; // scale-factor ring cursor +}; + + +// Relaxed blocking wait (no .acquire): the tcgen05 pipeline's own ordering +// covers these ring waits, and the .acquire form perturbs the scoreboard +// cadence on the issuing warp. Waits sit AT the consuming site — the carried +// one-stage-ahead try_wait peeks that used to gate them measured as pure +// overhead and were removed. +static __device__ __forceinline__ void mbar_wait_relaxed( + uint32_t bar_smem, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(bar_smem), "r"(parity) : "memory"); +} + +// Gated form for the @216 buffer-free wait, which fires once per tile: the +// gate branches over the spin for every group past the first. +static __device__ __forceinline__ void mbar_wait_gated_relaxed( + uint32_t bar_smem, uint32_t parity, bool peek) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p bra DONE_%=;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t" + "DONE_%=:\n\t}\n" + :: "r"(bar_smem), "r"(parity), "r"((uint32_t)peek) : "memory"); +} + + +// One scale-factor staging copy, smem -> TMEM. Single-lane semantics behind +// the converged elect predicate: tcgen05.cp is a single-thread instruction — +// issued bare from a converged warp you get 32 initiations and garbage. +static __device__ __forceinline__ void cp_sf( + uint32_t el, uint32_t dst, uint32_t lo, uint32_t hi) { + const uint64_t desc = (uint64_t(hi) << 32) | uint64_t(lo); + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;\n\t}\n" + :: "r"(dst), "l"(desc), "r"(el) : "memory"); +} + +// Ring-slot "consumed" arrive, multicast to both CTAs of the pair. +// arrive::one must fire exactly once — hence the elect predicate. +static __device__ __forceinline__ void commit_consumed( + uint32_t el, uint32_t smbase, uint32_t mb_off, uint16_t mask) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t}\n" + :: "r"(smbase + mb_off), "h"(mask), "r"(el) : "memory"); +} + +// One K=96 NVFP4 MMA. The asm token says scale_vec::4X regardless — K=96 is +// selected by idesc bit 31 (blog: the MMA warp). acc_pred = 0 overwrites the +// accumulator (the tile's first step), nonzero accumulates. +static __device__ __forceinline__ void mma_cell( + uint32_t el, uint32_t d, + uint32_t desc_a_lo, uint32_t desc_a_hi, + uint32_t desc_b_lo, uint32_t desc_b_hi, + uint32_t idesc, uint32_t sf_a, uint32_t sf_b, uint32_t acc_pred) { + const uint64_t desc_a = (uint64_t(desc_a_hi) << 32) | uint64_t(desc_a_lo); + const uint64_t desc_b = (uint64_t(desc_b_hi) << 32) | uint64_t(desc_b_lo); + asm volatile( + "{\n\t.reg .pred p, acc;\n\t" + "setp.ne.u32 p, %7, 0;\n\t" + "setp.ne.b32 acc, %4, 0;\n\t" + "@p tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X" + " [%0], %1, %2, %3, [%5], [%6], acc;\n\t}\n" + :: "r"(d), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(acc_pred), "r"(sf_a), "r"(sf_b), "r"(el) : "memory"); +} + +// One SF ring slot -> TMEM: 9 tcgen05.cp tiles (3 SFA + 6 SFB). Source +// address advances are in 16-byte descriptor units. The SFB tile order +// 0,3,1,4,2,5 interleaves the two 128-wide N-blocks. +static __device__ __forceinline__ void cp_sf_buffer( + uint32_t el, uint32_t slot, + uint32_t sfa0, uint32_t sfa1, uint32_t sfa2, + uint32_t sfb0, uint32_t sfb1, uint32_t sfb2, + uint32_t sfb3, uint32_t sfb4, uint32_t sfb5, + uint32_t sfa_hi, uint32_t sfa_lo, uint32_t sfb_hi, uint32_t sfb_lo) { + constexpr uint32_t SFA_SLOT_U = CD_SFA_STRIDE / 16u; // 96 + constexpr uint32_t SFB_SLOT_U = CD_SFB_STRIDE / 16u; // 192 + constexpr uint32_t TILE_U = CD_SF_TILE / 16u; // 32 + const uint32_t a = slot * SFA_SLOT_U + sfa_lo; + cp_sf(el, sfa0, a, sfa_hi); + cp_sf(el, sfa1, a + 1u*TILE_U, sfa_hi); + cp_sf(el, sfa2, a + 2u*TILE_U, sfa_hi); + const uint32_t b = slot * SFB_SLOT_U + sfb_lo; + cp_sf(el, sfb0, b + 0u*TILE_U, sfb_hi); + cp_sf(el, sfb1, b + 3u*TILE_U, sfb_hi); + cp_sf(el, sfb2, b + 1u*TILE_U, sfb_hi); + cp_sf(el, sfb3, b + 4u*TILE_U, sfb_hi); + cp_sf(el, sfb4, b + 2u*TILE_U, sfb_hi); + cp_sf(el, sfb5, b + 5u*TILE_U, sfb_hi); +} + +// Operand descriptor LOW word for A/B ring slot `slot`, K-step `k_atom` +// (0..7 inside the 768-element group; each step is 48 bytes = +3 in the +// 16-byte address field). +static __device__ __forceinline__ uint32_t ab_desc_lo( + uint32_t smbase, uint32_t ab_base_off, uint32_t slot, uint32_t k_atom) { + const uint64_t d = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + slot * CD_AB_STRIDE, int(k_atom) * 16); + return uint32_t(d & 0xFFFFFFFFu); +} + +// K-steps 6 and 7 cross a 128-byte window boundary. Their descriptor address +// comes from `addr_slot`, while the absolute-LBO field comes from the next +// slot whose swizzle window contains the second half. +static __device__ __forceinline__ uint32_t ab_desc_lo_straddle( + uint32_t smbase, uint32_t ab_base_off, + uint32_t addr_slot, uint32_t lbo_slot, uint32_t k_atom) { + const uint64_t d_addr = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + addr_slot * CD_AB_STRIDE, int(k_atom) * 16); + const uint64_t d_lbo = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + lbo_slot * CD_AB_STRIDE, 0); + return uint32_t(d_addr & 0x0000FFFFu) | uint32_t(d_lbo & 0xFFFF0000u); +} + +// One output tile: `full_groups` complete 768-element K-groups. Coordinates +// past logical K are zero-filled by TMA from the unchanged canonical inputs. +static __device__ __forceinline__ MainloopState mainloop_tile( + uint32_t taddr, uint32_t smbase, MainloopState in, uint16_t commit_mask, + uint32_t full_groups + ) { + constexpr uint32_t AB_RING_SLOTS = uint32_t(CD_AB_SLOTS); + constexpr uint32_t SF_RING_SLOTS = uint32_t(CD_SF_SLOTS); + constexpr uint32_t AB_FREE_DELTA = uint32_t(MB_AB_FREE - MB_AB_READY); + constexpr uint32_t SF_READY_OFF = uint32_t(MB_SF_READY); + constexpr uint32_t SF_FREE_DELTA = uint32_t(MB_SF_FREE - MB_SF_READY); + constexpr uint32_t ACC_FREE_OFF = uint32_t(MB_ACC_FREE); + constexpr uint32_t A_OFF = uint32_t(CD_A_OFF); + constexpr uint32_t B_OFF = uint32_t(CD_B_OFF); + + uint32_t d_parity = in.d_parity; + uint32_t ab_phase = in.ab_phase, ab_slot = in.ab_slot; + uint32_t sf_phase = in.sf_phase, sf_slot = in.sf_slot; + + // The caller has already selected exactly one lane for the whole MMA role. + // Keep the helper predicate constant throughout the tile. + const uint32_t el = 1u; + + // Fixed TMEM targets for the single reused scale-factor region: 9 cp + // tiles of 4 columns each, starting past the accumulator region. + constexpr uint32_t SF = uint32_t(SF_TMEM_BASE); + const uint32_t sfa_t0 = taddr + SF + 0; // SFA word 0 (MMA reads here) + const uint32_t sfa_t1 = taddr + SF + 4; + const uint32_t sfa_t2 = taddr + SF + 8; + const uint32_t sfb_t0 = taddr + SF + 12; // SFB word 0 + const uint32_t sfb_t1 = taddr + SF + 16; + const uint32_t sfb_t2 = taddr + SF + 20; + const uint32_t sfb_t3 = taddr + SF + 24; + const uint32_t sfb_t4 = taddr + SF + 28; + const uint32_t sfb_t5 = taddr + SF + 32; + + // Scale-factor cp source descriptors (low word carries the per-slot + // offset; high word is constant). + const uint64_t sfb_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFB_OFF)); + const uint32_t sfb_src_lo = uint32_t(sfb_src & 0xFFFFFFFFu); + const uint32_t sfb_src_hi = uint32_t(sfb_src >> 32); + const uint64_t sfa_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFA_OFF)); + const uint32_t sfa_src_lo = uint32_t(sfa_src & 0xFFFFFFFFu); + const uint32_t sfa_src_hi = uint32_t(sfa_src >> 32); + + // Operand descriptor HIGH words (low words come from ab_desc_lo per step). + const uint32_t a_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + A_OFF) >> 32); + const uint32_t b_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + B_OFF) >> 32); + + // Each K=96 step needs 6 scale bytes per row, split across two staged + // 32-bit words. Odd-issue steps read word 1: bit 31 of the SF address + // selects the with-id encoding, and the id bits fold into the idesc. + constexpr uint32_t SF_ADDR_ID = 0x80000000u; + const uint32_t sfa_w1 = (taddr + SF + 4) | SF_ADDR_ID; + const uint32_t sfb_w1 = (taddr + SF + 20) | SF_ADDR_ID; + const uint32_t w1_id_bits = + ((sfa_w1 >> 1) & 0x60000000u) | ((sfb_w1 >> 26) & 48u); + + // The tile's first MMA overwrites the accumulator buffer this tile owns; + // the @216 "buffer free" wait runs once, before the first group. + const uint32_t acc_wait_parity = d_parity; + d_parity ^= 1u; + const uint32_t d_tmem = d_parity * uint32_t(D_STRIDE) + taddr; + + // Instruction descriptor: kind::mxf4nvf4, M=256, N=256, E2M1 x E2M1, + // UE4M3 block-16, K-major, k_dim=1 (K=96). Base value 0x90400480. + constexpr uint32_t IDESC = ptx::mma_inst_desc_mxf4nvf4_block16( + 256u, 256u, 0u, 0u, ptx::Major::K, ptx::Major::K, + false, false, /*k_dim=*/true, ptx::ScaleFormat::E4M3); + static_assert(IDESC == 0x90400480u, "K=96 NVFP4 idesc drifted"); + const uint32_t w0_id_bits = + ((sfa_t0 >> 1) & 0x60000000u) | ((sfb_t0 >> 26) & 48u); + const uint32_t idesc_w0 = w0_id_bits | IDESC; + const uint32_t idesc_w1 = w1_id_bits | IDESC; + + + uint32_t acc_pred = 0; // first MMA of the tile overwrites + bool need_acc_wait = true; // @216 wait fires once per tile + + if (full_groups >= 1) { + uint32_t g = 0; + do { + // ---- SF batch 0 -> logical steps 0,1 (atoms 0,3) ---- + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // ---- SF batch 1 -> logical steps 2,3 (atoms 6,1) ---- + mbar_wait_relaxed(sf_bar1, sf_phase1); + uint32_t sf_slot2, sf_phase2; + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + // AB0's last reader is MMA2. + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // ---- SF batch 2 -> logical steps 4,5 (atoms 4,7) ---- + mbar_wait_relaxed(sf_bar2, sf_phase2); + uint32_t sf_slot3, sf_phase3; + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; // ring cursor writeback + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // ---- SF batch 3 -> logical steps 6,7 (atoms 2,5) ---- + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; // ring cursor writeback + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + // AB1's last reader is MMA5; this measured placement leaves the + // commit behind the following scale copy. + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + acc_pred = 0xFFFFFFFFu; // accumulate from the second group on + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + need_acc_wait = false; + // AB2's last reader is MMA7. + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, + commit_mask); + ++g; + } while (g != full_groups); + } + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; +} + +} // namespace nvfp4 + +// ============================================================================ +// The kernel. 224 threads = 7 warps, four jobs (blog: the warp table): +// warps 0-3 epilogue: drain TMEM, convert to FP16, store straight to global +// warp 4 the MMA warp (even CTA of each pair only): SF copies + MMAs +// warp 5 TMA producer: A and B tiles into the 6-slot ring +// warp 6 TMA producer: scale factors into the 7-slot ring +// ============================================================================ +namespace nvfp4 { + +// Baseline launch: plain launch bounds and a grid expressed in CTA units. +#define NVFP4_BLOCK_ATTR +constexpr bool GRID_IN_CLUSTERS = false; + +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void nvfp4_gemm_kernel(const __grid_constant__ CUtensorMap A_tmap, + const __grid_constant__ CUtensorMap B_tmap, + const __grid_constant__ CUtensorMap SFA_tmap, + const __grid_constant__ CUtensorMap SFB_tmap, + __half* C_ptr, int M, int N, int K + ) { + extern __shared__ __align__(1024) char smem_buf[]; + SmemCD& smem = *reinterpret_cast(smem_buf); + + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const uint32_t crank = ptx::cluster_rank(); + const uint32_t m_pair = crank & 1u; // 0 = the pair's even CTA + + const int cluster_grid_m = (M + CLUSTER_M - 1) / CLUSTER_M; + const int grid_n = (N + BLOCK_N - 1) / BLOCK_N; + const int total_tiles = cluster_grid_m * grid_n; + + // Baseline rungs execute a whole K=768 group at the edge. Tensor-map OOB + // fill supplies zeros past logical K without changing the input buffers. + const int full_groups = (K + 767) / 768; + const int num_groups = full_groups; + + // Persistent cluster index. The device always sees the grid in CTA units + // (2, 1, clusters) regardless of the host-side launch-unit regime. + const bool grid_z = (gridDim.z > 1); + const int num_clusters = grid_z ? int(gridDim.z) : int(gridDim.x) / 2; + const int cluster_id = grid_z ? int(blockIdx.z) : int(blockIdx.x) / 2; + + + // Which output tile does persistent work item `t` process? + auto tile_of = [&](int t) -> int { + return t; // natural raster: M fast, N slow + }; + + auto producer_tile_of = [&](int t) -> int { + return tile_of(t); + }; + + // ---- mbarrier init (warp 0, lane 0) ---- + if (warp_id == 0 && ptx::lane_id() == 0) { + // Arrival counts are derived from the topology, never tuned: + // "ready" = one producer expect_tx per slot (bit-24 routes both + // peers' bytes to the even CTA's barrier); "free" = one consumer + // commit per slot reuse. + #pragma unroll + for (int s = 0; s < CD_AB_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.ab_ready[s], 1); + ptx::mbar_init(&smem.hdr.ab_free[s], 1); + } + #pragma unroll + for (int s = 0; s < CD_SF_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.sf_ready[s], 1); + ptx::mbar_init(&smem.hdr.sf_free[s], 1); + } + ptx::mbar_init(&smem.hdr.acc_ready, 1); + ptx::mbar_init(&smem.hdr.acc_free, N_EPI_WARPS * CTA_GROUP); // 4 x 2 = 8 + ptx::mbar_init(&smem.hdr.dealloc, 32); // one warp's lanes, cross-peer + ptx::fence_mbarrier_init_release_cluster(); + } + ptx::cluster_sync_rel_acq(); + + // TMEM allocation: epilogue warp 0 allocates for the pair; warps 0-4 + // rendezvous before reading the granted base. The TMA producer warps + // never touch TMEM and proceed straight to their loops. + uint32_t taddr = 0; + if (warp_id <= 4) { + if (warp_id == 0) + ptx::tcgen05_alloc_2sm(ptx::to_shared(&smem.hdr.tmem_addr), + N_ALLOC_COLS); + ptx::bar_sync(2, 160); + taddr = smem.hdr.tmem_addr; + } + + const int wg = warp_id >> 2; // 0 = epilogue warps, 1 = producers + MMA + if (wg == 1) { + const int pw = warp_id - 4; + if (pw == 1 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 5: lane 0 owns the complete A/B TMA producer loop. + // ================================================================ + // free_phase starts at 1: a fresh mbarrier is phase 0, so the + // first ring's worth of "free" waits falls through. + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + // The pair splits the 256x256 tile: even CTA holds A rows + // [0,128) and B cols [0,128); odd CTA the other halves. + const int m_blk = m_row * CTA_GROUP + int(m_pair); + const int n_blk = n_grp * CTA_GROUP + int(m_pair); + #pragma unroll 1 + for (int g = 0; g < num_groups; ++g) { + // Three full 128-byte windows cover the packed K=768 group. + #pragma unroll + for (int sub = 0; sub < AB_TMA_SUBS; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.ab_free[slot], free_phase); + const uint32_t a_s = + base + CD_A_OFF + uint32_t(slot) * CD_AB_STRIDE; + const uint32_t b_s = + base + CD_B_OFF + uint32_t(slot) * CD_AB_STRIDE; + uint64_t* bar = &smem.hdr.ab_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, AB_TX); + // A rides the multicast form with a self mask so the + // completion routes to the even peer's barrier. + const uint16_t mask_a = uint16_t(1u << m_pair); + // Read the same row-major FP4 allocation as cuBLAS. + // K192 uses four 96-byte boxes; K256 uses three + // 128-byte boxes. The tensor map zero-fills past K. + const CUtensorMap* a_map = &A_tmap; + const CUtensorMap* b_map = &B_tmap; + const int x = g * 384 + sub * AB_LOAD_BYTES; + const int y_a = m_blk * BLOCK_M; + const int y_b = n_blk * BLOCK_N_PEER; + const int z_a = 0; + const int z_b = 0; + ptx::cp_async_bulk_tensor_3d_load_multicast( + a_s, a_map, x, y_a, z_a, + bar, mask_a + ); + ptx::cp_async_bulk_tensor_3d_load_2sm_bit24( + b_s, b_map, x, y_b, z_b, + bar + ); + if (++slot == CD_AB_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + // Drain: wait for the last-filled slot to be consumed, then hand + // it one clean-shutdown expect_tx. + { + const int t_slot = (slot == 0) ? CD_AB_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.ab_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.ab_ready[t_slot], AB_TX); + } + } else if (pw == 2 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 6: scale-factor TMA producer. SFA is per-peer (each CTA's + // own M half); SFB is the one genuinely multicast operand — the + // cta_group::2 copy reads BOTH peers' smem as one logical + // operand, so both copies must be byte-identical (blog: the scale-factor section). + // ================================================================ + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const int m_blk = m_row * CTA_GROUP + int(m_pair); + for (int g = 0; g < num_groups; ++g) { + // 4 SF ring slots per 768-K group (2 K=96 steps each). + #pragma unroll + for (int sub = 0; sub < 4; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.sf_free[slot], free_phase); + const uint32_t sfa_s = + base + CD_SFA_OFF + uint32_t(slot) * CD_SFA_STRIDE; + const uint32_t sfb_s = + base + CD_SFB_OFF + uint32_t(slot) * CD_SFB_STRIDE; + uint64_t* bar = &smem.hdr.sf_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, SF_TX); + // VEC16 is [outer-block][K-group-of-4][128 rows x 4]. + // Three adjacent groups are the 128 x 12 scale slot. + const int sf_group = (g * 4 + sub) * 3; + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfa_s, &SFA_tmap, 0, 0, sf_group, m_blk, bar, + uint16_t(1u << m_pair) + ); + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfb_s + m_pair * uint32_t(CD_SFB_STRIDE / 2), + &SFB_tmap, 0, 0, sf_group, + n_grp * CTA_GROUP + int(m_pair), bar, + uint16_t(0x3) + ); + if (++slot == CD_SF_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + { + const int t_slot = (slot == 0) ? CD_SF_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.sf_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.sf_ready[t_slot], SF_TX); + } + } else if (pw == 0 && m_pair == 0 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 4: the MMA warp — issued once per cluster, on the even + // CTA. The odd CTA's warp 4 does nothing all kernel (blog: the warp table). + // The mainloop needs no tile coordinates: operands arrive through + // the rings; this warp only needs its share of the trip count. + // ================================================================ + const uint32_t smbase = ptx::to_shared(&smem); + const uint16_t commit_mask = 0x3; // both CTAs of this pair + MainloopState st{1u, 0u, 0u, 0u, 0u}; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + st = mainloop_tile(taddr, smbase, st, commit_mask, + uint32_t(full_groups) + ); + // The outer role branch already leaves exactly one issuer. + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;" + :: "r"(smbase + uint32_t(MB_ACC_READY)), + "h"(uint16_t(0x3)) : "memory"); + } + // Drain the final tile's buffer-free arrive before exiting. + mbar_wait_relaxed(smbase + uint32_t(MB_ACC_FREE), st.d_parity); + } + } else { + // ==================================================================== + // Warps 0-3: the epilogue (blog: the epilogue). Each warp owns a 32-row band + // and drains a tile as eight 32-column tcgen05.ld loads, converts + // FP32->FP16 in packed pairs, and stores straight to global — no + // shared-memory bounce. + // ==================================================================== + const int epi_warp = warp_id; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + const uint32_t acc_free_rank = crank & ~1u; // the pair's even CTA + int d_db = 0; + int acc_ready_phase = 0; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const uint32_t tmem_base = taddr + uint32_t(d_db) * D_STRIDE; + mbar_wait_relaxed(ptx::to_shared(&smem.hdr.acc_ready), + uint32_t(acc_ready_phase)); + acc_ready_phase ^= 1; + + const int off_m = (m_row * CTA_GROUP + int(m_pair)) * BLOCK_M; + const int off_n = n_grp * BLOCK_N; + const int row = off_m + epi_warp * 32 + lane_id; + const bool row_in = (row < M); + constexpr int EPI_BANDS = BLOCK_N / EPI_CHUNK; // 8 + #pragma unroll 1 + for (int k = 0; k < EPI_BANDS; ++k) { + // Buffer 0 drains high-to-low and buffer 1 low-to-high, so + // their shared columns [220,256) drain first for both. + const int parity = d_db & 1; + const int band = (parity == 0) ? (EPI_BANDS - 1 - k) : k; + const uint32_t taddr_n = + tmem_base + uint32_t(band) * 32 + taddr_lane; + uint32_t w[32]; + ptx::tcgen05_ld_32x32b_x32(taddr_n, w); + // Return the old buffer as soon as its shared edge is out. + if (k == 1) { + ptx::tcgen05_wait_ld_then_arrive_release( + &smem.hdr.acc_free, acc_free_rank, + /*do_arrive=*/ptx::lane_id() == 0); + } + if (!row_in) continue; + const int col_base = off_n + band * 32; + uint32_t h[16]; + #pragma unroll + for (int j = 0; j < 16; ++j) { + h[j] = ptx::cvt_pack_f16x2(__int_as_float(w[2 * j + 1]), + __int_as_float(w[2 * j])); + } + const size_t row_base = size_t(row) * size_t(N); + const bool aligned32 = ((row_base & 15u) == 0u) && + ((uintptr_t(C_ptr) & 31u) == 0u); + if (aligned32 && (col_base & 15) == 0) { + // Precondition here: row_base is 16-element (32-byte) + // aligned — the aligned32 guard owns it. Restating it to + // ptxas via __builtin_assume compiled this file + // SASS-identical (and measured perf-null on the + // production kernel), so the guard stands alone. + #pragma unroll + for (int q = 0; q < 2; ++q) { + const int c0 = col_base + q * 16; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 16 <= N) { + uint4 lo{h[q*8+0], h[q*8+1], h[q*8+2], h[q*8+3]}; + uint4 hi{h[q*8+4], h[q*8+5], h[q*8+6], h[q*8+7]}; + ptx::st_global_na_ef_b256(cdst, lo, hi); + } else { + // N tail: 128-bit then scalar. + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int cq = c0 + p * 8; + if (cq + 8 <= N) { + uint4 v{h[q*8+p*4+0], h[q*8+p*4+1], + h[q*8+p*4+2], h[q*8+p*4+3]}; + ptx::st_global_na_b128( + reinterpret_cast(cdst + p * 8), v); + } else { + const uint16_t* hb = + reinterpret_cast( + &h[q * 8 + p * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (cq + c < N) cdst[p * 8 + c] = hb[c]; + } + } + } + } + } else { + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int c0 = col_base + q * 8; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 8 <= N) { + uint4 v{h[q*4+0], h[q*4+1], h[q*4+2], h[q*4+3]}; + ptx::st_global_na_b128( + reinterpret_cast(cdst), v); + } else { + const uint16_t* hb = + reinterpret_cast(&h[q * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (c0 + c < N) cdst[c] = hb[c]; + } + } + } + } + if (++d_db == D_DB) d_db = 0; + } + } + + // Teardown: only the epilogue warps rendezvous per CTA; warp 0 then runs + // the cross-peer dealloc handshake. Producer and MMA warps exit freely — + // their drains above already order their work. + if (wg == 0) { + ptx::bar_sync(3, 128); + if (warp_id == 0) { + ptx::tcgen05_relinquish_2sm(); + const uint32_t peer_rank = crank ^ 1u; + // All 32 lanes arrive the PEER's count-32 barrier = one flip. + ptx::mbar_arrive_cluster_release(&smem.hdr.dealloc, peer_rank); + ptx::mbar_wait_parity(&smem.hdr.dealloc, 0); + ptx::tcgen05_dealloc_2sm(taddr, N_ALLOC_COLS); + } + } +} + + +// Keep the host launch convention paired with the kernel's block attribute. +static inline dim3 launch_grid(int clusters) { + return GRID_IN_CLUSTERS ? dim3(1, 1, unsigned(clusters)) + : dim3(2, 1, unsigned(clusters)); +} + +} // namespace nvfp4 + +// ============================================================================ +// Host-side TMA descriptor support. The shared input and benchmark code lives +// in main.cu. +// ============================================================================ +namespace tmap { + +// Thin wrappers over cuTensorMapEncodeTiled. Dim 0 is the innermost +// (stride-1) axis; strides for dim 1+ are in bytes. +inline CUtensorMap encode_tiled_3d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_depth, uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t depth_stride_bytes, + uint32_t box_depth, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[3] = { global_cols, global_rows, global_depth }; + cuuint64_t global_strides[2] = { row_stride_bytes, depth_stride_bytes }; + cuuint32_t box_dim[3] = { box_cols, box_rows, box_depth }; + cuuint32_t element_strides[3] = { 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 3, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +inline CUtensorMap encode_tiled_4d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_slots, uint64_t global_tiles, + uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t tile_stride_bytes, + uint64_t slot_stride_bytes, + uint32_t box_tiles, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[4] = { global_cols, global_rows, global_tiles, + global_slots }; + cuuint64_t global_strides[3] = { row_stride_bytes, tile_stride_bytes, + slot_stride_bytes }; + cuuint32_t box_dim[4] = { box_cols, box_rows, box_tiles, 1u }; + cuuint32_t element_strides[4] = { 1, 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 4, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +} // namespace tmap + +namespace host { + +using namespace nvfp4; + +// ---- implementation-specific tensor maps --------------------------------- + +inline CUtensorMap make_ab_tmap(void* ptr, int rows, int K, + int row_stride) { + return tmap::encode_tiled_3d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*depth=*/1, uint64_t(rows), /*cols=*/uint64_t(K / 2), + /*row_stride=*/uint64_t(row_stride), + /*depth_stride=*/uint64_t(rows) * uint64_t(row_stride), + /*box=*/1, BLOCK_M, AB_LOAD_BYTES, CU_TENSOR_MAP_SWIZZLE_128B); +} + +inline CUtensorMap make_sf_tmap(uint8_t* ptr, int outer, int K) { + const uint64_t sf_inner = (uint64_t(K) + 63) / 64 * 4; + const uint64_t outer_blocks = (uint64_t(outer) + 127) / 128; + return tmap::encode_tiled_4d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*slots=*/outer_blocks, /*tiles=*/sf_inner / 4, + /*rows=*/4, /*cols=*/128, + /*row_stride=*/128, /*tile_stride=*/512, + /*slot_stride=*/sf_inner * 128, + /*box_tiles=*/3, /*box_rows=*/4, /*box_cols=*/128, + CU_TENSOR_MAP_SWIZZLE_NONE); +} +} // namespace host diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/gemm6.cuh b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm6.cuh new file mode 100644 index 000000000..b13205322 --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm6.cuh @@ -0,0 +1,1264 @@ +#pragma once +// r6: promise the exact 224-thread block size. +// This is a complete, directly editable ladder snapshot. +// It intentionally duplicates the implementation in the other gemm*.cuh files. + +// Conditional branches are already materialized; edit this implementation directly. +#define NVFP4_BUILD_NAME "r6" +// Read-only metadata for main.cu; these do not select code in this header. +#define NVFP4_HAS_K64_TAIL 0 +#define NVFP4_HAS_SCHEDULE 0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(expr) do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \ + cudaGetErrorName(_e), __FILE__, __LINE__, \ + cudaGetErrorString(_e)); \ + std::abort(); \ + } \ +} while (0) + +#define CU_CHECK(expr) do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + std::fprintf(stderr, "CU error at %s:%d: %s\n", \ + __FILE__, __LINE__, _s ? _s : "?"); \ + std::abort(); \ + } \ +} while (0) + + +// ============================================================================ +// PTX helpers. Thin zero-cost wrappers over the exact instruction forms the +// kernel needs — nothing generic, nothing speculative. +// ============================================================================ +namespace ptx { + +// Generic pointer -> 32-bit `.shared` window byte address. +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Local `.shared` address -> the same offset in CTA `cta_rank`'s smem +// (`.shared::cluster` instructions take this mapped form). +static __device__ __forceinline__ uint32_t mapa_shared_cluster( + uint32_t local_addr, uint32_t cta_rank) { + uint32_t mapped; + asm("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(mapped) : "r"(local_addr), "r"(cta_rank)); + return mapped; +} + +static __device__ __forceinline__ void bar_sync(uint32_t id, uint32_t count) { + asm volatile("barrier.cta.sync %0, %1;" :: "r"(id), "r"(count)); +} + +// Cluster barrier with release/acquire — the publish boundary between the +// mbarrier inits and any cross-CTA use of those mbarriers. +static __device__ __forceinline__ void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// %cluster_ctarank: 0 or 1 in this kernel's (2,1,1) clusters. +static __device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} + +static __device__ __forceinline__ uint32_t lane_id() { + uint32_t r; + asm("mov.u32 %0, %%laneid;" : "=r"(r)); + return r; +} + +// ---- mbarriers ------------------------------------------------------------- + +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + :: "r"(to_shared(bar)), "r"(count)); +} + +// Publish the inits to the async proxy + the peer CTA before any TMA +// completes on them. +static __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// Cross-CTA arrive with release ordering (teardown handshake). +static __device__ __forceinline__ void mbar_arrive_cluster_release( + uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" + :: "r"(mapped)); +} + +// Fused `tcgen05.wait::ld` + "buffer free" arrive, one asm block. The arrive +// must not become visible while TMEM loads are still in flight (it frees the +// accumulator buffer to the next tile's MMA — the bring-up chapter's +// drain overlap); the fused block +// keeps the arrive pinned behind the wait. (An earlier form also manufactured +// a register dependency on a drained value, `dep * 0` folded into the +// address; verified dead on nvcc 13.1 — ptxas const-folds it, SASS +// byte-identical without it — and retired.) +static __device__ __forceinline__ void tcgen05_wait_ld_then_arrive_release( + uint64_t* bar, uint32_t cta_rank, bool do_arrive) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile( + "{\n\t" + ".reg .pred %%pdo;\n\t" + "setp.ne.s32 %%pdo, %1, 0;\n\t" + "tcgen05.wait::ld.sync.aligned;\n\t" + "@%%pdo mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;\n\t" + "}\n\t" + :: "r"(mapped), "r"((int)do_arrive) : "memory"); +} + + +// Blocking parity wait at the consuming site (producer ring/tail waits + +// teardown). The 10 ms suspend hint is the polite form: a full-rate spin +// competes for issue slots and measurably loses; a nanosleep-backoff spin +// measured the same as this hint across five windows, so the simpler form +// ships. The carried one-stage-ahead try_wait peeks that used to gate these +// waits measured as pure overhead and were removed. +static __device__ __forceinline__ void mbar_wait_parity( + uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(to_shared(bar)), "r"(parity)); +} + +// ---- TMA ------------------------------------------------------------------- + +// Producer expect_tx arrive on the cluster-shared ring barrier: both +// peers' transfer bytes account against the even CTA's barrier, so one +// barrier per ring slot covers the whole cluster-wide gather. (Mechanism: +// `& 0xFEFFFFFF` clears bit 24 of the shared address — the hardware's +// peer-CTA bit — re-pointing the arrive at the even CTA's copy.) +static __device__ __forceinline__ void mbar_arrive_expect_tx_cluster( + uint64_t* bar, uint32_t bytes) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes)); +} + +// 3D tiled load, cta_group::2 + multicast (the A feed; the mask is the +// issuing CTA itself in a (2,1) cluster). Completion routes to the even +// CTA's mbarrier via the bit-24 clear. +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z) + : "memory"); +} + +// 3D tiled load, cta_group::2, no multicast (the B feed — each peer loads +// its own 128-wide N-block). +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_2sm_bit24( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst_smem), "l"(tmap), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) + : "memory"); +} + +// 4D tiled load, cta_group::2 + multicast (the scale-factor feed). The +// 512-byte SF tiles of one ring slot form a 4D box; `w` selects the slot. +static __device__ __forceinline__ void cp_async_bulk_tensor_4d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z2, int32_t w, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.4d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z2), "r"(w) + : "memory"); +} + + +// ---- tcgen05 (tensor memory + MMA plumbing) --------------------------------- + +static __device__ __forceinline__ void tcgen05_alloc_2sm( + uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_dealloc_2sm( + uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" + :: "r"(taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_relinquish_2sm() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); +} + +// The widest single-instruction TMEM read: 32 b32 registers per lane +// (blog: the epilogue — a quarter the load/wait count of the narrow forms). +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t taddr, uint32_t* r) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), + "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), + "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), + "=r"(r[16]), "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), + "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), + "=r"(r[24]), "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), + "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr)); +} + +// ---- stores + conversion ----------------------------------------------------- + +static __device__ __forceinline__ uint32_t cvt_pack_f16x2(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// L2::evict_first requires a 256-bit store, but the 128-bit tail can still +// adopt the no-L1-allocation part of the policy at this rung. +static __device__ __forceinline__ void st_global_na_b128( + uint4* ptr, uint4 v) { + asm volatile("st.global.L1::no_allocate.v4.b32 " + "[%0], {%1, %2, %3, %4};" + :: "l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)); +} + +// Add the output cache policy to the 256-bit store: bypass L1 and evict the +// write-once line first from L2 so it does not displace reusable operands. +static __device__ __forceinline__ void st_global_na_ef_b256( + void* ptr, uint4 lo, uint4 hi) { + asm volatile("st.global.L1::no_allocate.L2::evict_first.v8.b32 [%0], " + "{%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"(ptr), "r"(lo.x), "r"(lo.y), "r"(lo.z), "r"(lo.w), + "r"(hi.x), "r"(hi.y), "r"(hi.z), "r"(hi.w)); +} + + +// ---- MMA descriptors --------------------------------------------------------- + +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class ScaleFormat : uint8_t { E4M3 = 0, E8M0 = 1 }; + +// Source descriptor for the tcgen05.cp that stages scale factors from smem +// into TMEM: unswizzled, SBO=128, LBO=16. The low word later carries the +// per-slot byte offset (in 16-byte units). +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_sf_cp_src_desc_cd(uint32_t sf_smem_addr) { + uint64_t addr_field = uint64_t((sf_smem_addr >> 4) & 0x3FC0u); + return addr_field | 0x400800010000ull; +} + +// Operand descriptor for a K=96 (48-byte) slice inside a 128-byte-swizzled +// smem K-window. OR-constant 0x4010404000000000: version=1, absolute-address +// LBO mode (bit 52), 128B swizzle, SBO=1024. Both the address field and the +// absolute-LBO field encode the window's own 1024-aligned base; the +// intra-window offset advances the address field only. +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k96_sw128(uint32_t window_base_1024aligned, int intra_window_off = 0) { + uint64_t addr_field = uint64_t(((window_base_1024aligned >> 4) & 0x7FC0u) + + uint32_t(intra_window_off >> 4)); + uint64_t lbo_field = uint64_t((window_base_1024aligned << 12) & 0x7FC00000u); + return addr_field | lbo_field | 0x4010404000000000ull; +} + +// Instruction descriptor for tcgen05.mma.kind::mxf4nvf4.block_scale. +// The K extent is selected by BIT 31, not by the asm mnemonic: k_dim=0 runs +// K=64 (0x10400480), k_dim=1 runs K=96 (0x90400480, sm_103a-only) — the +// cheapest quarter of the gap if your kernel is built on the K=64 form +// (blog: the MMA warp). Bit 13 flips which TMEM scale-factor column the MMA reads. +__host__ __device__ static __forceinline__ constexpr uint32_t +mma_inst_desc_mxf4nvf4_block16( + uint32_t M, uint32_t N, + uint32_t tmem_sfa_addr, uint32_t tmem_sfb_addr, + Major a_major, Major b_major, + bool negate_a, bool negate_b, + bool k_dim, ScaleFormat sf) { + constexpr uint32_t MXF4_E2M1 = 1u; + uint32_t d = 0; + d |= ((tmem_sfb_addr & 0xC0000000u) >> 30) << 4; // b_sf_id + d |= MXF4_E2M1 << 7; // a_format = E2M1 + d |= MXF4_E2M1 << 10; // b_format = E2M1 + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; // N in units of 8 + d |= (static_cast(sf) & 0x1u) << 23; // 0 = UE4M3 + d |= ((M >> 4) & 0x1Fu) << 24; // m_dim at bits 27-28 + d |= ((tmem_sfa_addr & 0xC0000000u) >> 30) << 29; // a_sf_id + if (k_dim) d |= 1u << 31; // 1 = K=96 + return d; +} + +} // namespace ptx + +// ============================================================================ +// Kernel geometry from the blog's opening kernel section. One tile shape, +// one cluster shape: +// 256x256 output per 2-CTA cluster (128 rows per CTA), K advanced in +// 96-element MMA steps staged as 768-element K-groups of 8 steps, persistent +// grid of 76 clusters = 152 CTAs = one per SM. +// ============================================================================ +namespace nvfp4 { + +constexpr int CTA_GROUP = 2; // cta_group::2 — the MMA spans both CTAs +constexpr int BLOCK_M = 128; // A rows per CTA +constexpr int BLOCK_N = 256; // output tile N extent (per cluster) +constexpr int BLOCK_N_PEER = BLOCK_N / CTA_GROUP; // 128 — B cols per CTA +constexpr int CLUSTER_M = CTA_GROUP * BLOCK_M; // 256 — tile M extent +[[maybe_unused]] constexpr int MMA_K = 96; // K per tcgen05.mma (idesc bit 31) +constexpr int WIN_BYTES = 128; // one 128B-swizzled smem K-window +constexpr int AB_LOAD_BYTES = WIN_BYTES; +constexpr int AB_TMA_SUBS = 3; +constexpr int TB_SIZE = 224; // 7 warps (blog: the warp table) +constexpr int N_EPI_WARPS = 4; // warps 0-3 + +constexpr int SF_TMEM_COLS = 9 * 4; // 3 SFA + 6 SFB cp tiles, 4 cols each +// Two 256-wide accumulators use a 220-column stride and overlap by the 36 +// columns needed for scale factors. +constexpr int D_STRIDE = 220; +constexpr int D_DB = 2; +constexpr int N_D_COLS = D_STRIDE + BLOCK_N; +constexpr int SF_TMEM_BASE = 476; +constexpr int N_ALLOC_COLS = 512; +static_assert(N_D_COLS <= SF_TMEM_BASE && + SF_TMEM_BASE + SF_TMEM_COLS <= N_ALLOC_COLS, + "accumulators + scale factors exceed the 512 TMEM columns"); +constexpr int EPI_CHUNK = 32; // cols per regular drain + +// ---- the shared-memory map (blog: the feed) ----------------------------------- +// Every offset is a compile-time expression. +constexpr int CD_HDR = 1024; // mbarrier header pad +constexpr int CD_AB_SLOTS = 6; +constexpr int CD_SF_SLOTS = 7; +constexpr int CD_AB_STRIDE = BLOCK_M * WIN_BYTES; // 16384 per A or B slot +constexpr int CD_SF_TILE = 512; // one (32 lanes x 16B) x4 tile +constexpr int CD_SFA_TILES = 3; // per SFA slot (M = 1 block) +constexpr int CD_SFB_TILES = 6; // per SFB slot (N = 2 blocks) +constexpr int CD_SFA_STRIDE = CD_SFA_TILES * CD_SF_TILE; // 1536 +constexpr int CD_SFB_STRIDE = CD_SFB_TILES * CD_SF_TILE; // 3072 + +// Expected transfer bytes are derived, never tuned: a wrong-by-one +// expectation doesn't crash, it deadlocks (blog: the feed). +constexpr uint32_t AB_TX = uint32_t(CTA_GROUP) * + uint32_t(BLOCK_M + BLOCK_N_PEER) * uint32_t(AB_LOAD_BYTES); +constexpr uint32_t SF_TX = uint32_t(CD_SFA_STRIDE + CD_SFB_STRIDE) * + uint32_t(CTA_GROUP); // 9216 +static_assert(AB_TX == 65536u && SF_TX == 9216u, + "expect_tx derivation drifted"); + +// The map IS the struct: member sizes and alignment produce every byte +// offset; the raw-address constants the mainloop uses read off it with +// offsetof, and the static_asserts pin the published values. +struct SmemCD { + // mbarrier header: four ring regions in handshake order, then the + // singleton accumulator + teardown barriers (alignas pads it to 1024). + struct alignas(1024) Header { + uint64_t ab_ready[CD_AB_SLOTS]; // TMA -> MMA + uint64_t ab_free [CD_AB_SLOTS]; // MMA -> TMA + uint64_t sf_ready[CD_SF_SLOTS]; // TMA -> MMA + uint64_t sf_free [CD_SF_SLOTS]; // MMA -> TMA + uint64_t acc_ready; // MMA -> epi + uint64_t acc_free; // epi -> MMA + uint64_t dealloc; // 2-CTA teardown + uint32_t tmem_addr; // TMEM base addr storage + } hdr; + uint8_t a [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t b [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t sfa[CD_SF_SLOTS][CD_SFA_STRIDE]; + // SFB starts 1024-aligned (the cp source descriptor drops address bits 4-9). + alignas(1024) uint8_t sfb[CD_SF_SLOTS][CD_SFB_STRIDE]; +}; +static_assert(sizeof(SmemCD::Header) == CD_HDR, "mbarrier header overflows the pad"); + +constexpr int CD_A_OFF = offsetof(SmemCD, a); +constexpr int CD_B_OFF = offsetof(SmemCD, b); +constexpr int CD_SFA_OFF = offsetof(SmemCD, sfa); +constexpr int CD_SFB_OFF = offsetof(SmemCD, sfb); +constexpr int CD_TOTAL = sizeof(SmemCD); +static_assert(CD_TOTAL <= 232448, "smem map exceeds the opt-in budget"); +static_assert(CD_AB_SLOTS != 6 || + (CD_B_OFF == 99328 && CD_SFA_OFF == 197632 && + CD_SFB_OFF == 208896 && CD_TOTAL == 230400), + "the 6-slot byte map drifted from its published values"); + +// The mainloop computes raw `.shared` byte addresses from these. +constexpr int MB_AB_READY = offsetof(SmemCD, hdr.ab_ready); +constexpr int MB_AB_FREE = offsetof(SmemCD, hdr.ab_free); +constexpr int MB_SF_READY = offsetof(SmemCD, hdr.sf_ready); +constexpr int MB_SF_FREE = offsetof(SmemCD, hdr.sf_free); +constexpr int MB_ACC_READY = offsetof(SmemCD, hdr.acc_ready); +constexpr int MB_ACC_FREE = offsetof(SmemCD, hdr.acc_free); + + + +} // namespace nvfp4 + +// ============================================================================ +// The consumer mainloop from the blog's pipeline section. One warp — warp 4 +// of the even CTA — +// issues everything: the tcgen05.cp copies that stage scale factors from the +// smem ring into TMEM, immediately followed by the tcgen05.mma stream, in the +// SAME warp. A same-warp tcgen05 stream is one ordered async pipeline, so the +// write-after-read hazard on the single reused scale-factor TMEM region needs +// no fence and no cross-warp barrier. +// +// Per 768-element K-group, logical K=96 steps 0..7 issue in K order. Their +// starting offsets, expressed as 16-byte descriptor atoms modulo a 128-byte +// swizzle window, are 0,3,6,1,4,7,2,5. Atoms 6 and 7 straddle windows. +// ============================================================================ +namespace nvfp4 { + + +// Ring/parity state carried across tiles by the MMA warp. +struct MainloopState { + uint32_t d_parity; // which TMEM accumulator buffer the next tile writes + uint32_t ab_phase, ab_slot; // A/B ring cursor + uint32_t sf_phase, sf_slot; // scale-factor ring cursor +}; + + +// Relaxed blocking wait (no .acquire): the tcgen05 pipeline's own ordering +// covers these ring waits, and the .acquire form perturbs the scoreboard +// cadence on the issuing warp. Waits sit AT the consuming site — the carried +// one-stage-ahead try_wait peeks that used to gate them measured as pure +// overhead and were removed. +static __device__ __forceinline__ void mbar_wait_relaxed( + uint32_t bar_smem, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(bar_smem), "r"(parity) : "memory"); +} + +// Gated form for the @216 buffer-free wait, which fires once per tile: the +// gate branches over the spin for every group past the first. +static __device__ __forceinline__ void mbar_wait_gated_relaxed( + uint32_t bar_smem, uint32_t parity, bool peek) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p bra DONE_%=;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t" + "DONE_%=:\n\t}\n" + :: "r"(bar_smem), "r"(parity), "r"((uint32_t)peek) : "memory"); +} + + +// One scale-factor staging copy, smem -> TMEM. Single-lane semantics behind +// the converged elect predicate: tcgen05.cp is a single-thread instruction — +// issued bare from a converged warp you get 32 initiations and garbage. +static __device__ __forceinline__ void cp_sf( + uint32_t el, uint32_t dst, uint32_t lo, uint32_t hi) { + const uint64_t desc = (uint64_t(hi) << 32) | uint64_t(lo); + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;\n\t}\n" + :: "r"(dst), "l"(desc), "r"(el) : "memory"); +} + +// Ring-slot "consumed" arrive, multicast to both CTAs of the pair. +// arrive::one must fire exactly once — hence the elect predicate. +static __device__ __forceinline__ void commit_consumed( + uint32_t el, uint32_t smbase, uint32_t mb_off, uint16_t mask) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t}\n" + :: "r"(smbase + mb_off), "h"(mask), "r"(el) : "memory"); +} + +// One K=96 NVFP4 MMA. The asm token says scale_vec::4X regardless — K=96 is +// selected by idesc bit 31 (blog: the MMA warp). acc_pred = 0 overwrites the +// accumulator (the tile's first step), nonzero accumulates. +static __device__ __forceinline__ void mma_cell( + uint32_t el, uint32_t d, + uint32_t desc_a_lo, uint32_t desc_a_hi, + uint32_t desc_b_lo, uint32_t desc_b_hi, + uint32_t idesc, uint32_t sf_a, uint32_t sf_b, uint32_t acc_pred) { + const uint64_t desc_a = (uint64_t(desc_a_hi) << 32) | uint64_t(desc_a_lo); + const uint64_t desc_b = (uint64_t(desc_b_hi) << 32) | uint64_t(desc_b_lo); + asm volatile( + "{\n\t.reg .pred p, acc;\n\t" + "setp.ne.u32 p, %7, 0;\n\t" + "setp.ne.b32 acc, %4, 0;\n\t" + "@p tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X" + " [%0], %1, %2, %3, [%5], [%6], acc;\n\t}\n" + :: "r"(d), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(acc_pred), "r"(sf_a), "r"(sf_b), "r"(el) : "memory"); +} + +// One SF ring slot -> TMEM: 9 tcgen05.cp tiles (3 SFA + 6 SFB). Source +// address advances are in 16-byte descriptor units. The SFB tile order +// 0,3,1,4,2,5 interleaves the two 128-wide N-blocks. +static __device__ __forceinline__ void cp_sf_buffer( + uint32_t el, uint32_t slot, + uint32_t sfa0, uint32_t sfa1, uint32_t sfa2, + uint32_t sfb0, uint32_t sfb1, uint32_t sfb2, + uint32_t sfb3, uint32_t sfb4, uint32_t sfb5, + uint32_t sfa_hi, uint32_t sfa_lo, uint32_t sfb_hi, uint32_t sfb_lo) { + constexpr uint32_t SFA_SLOT_U = CD_SFA_STRIDE / 16u; // 96 + constexpr uint32_t SFB_SLOT_U = CD_SFB_STRIDE / 16u; // 192 + constexpr uint32_t TILE_U = CD_SF_TILE / 16u; // 32 + const uint32_t a = slot * SFA_SLOT_U + sfa_lo; + cp_sf(el, sfa0, a, sfa_hi); + cp_sf(el, sfa1, a + 1u*TILE_U, sfa_hi); + cp_sf(el, sfa2, a + 2u*TILE_U, sfa_hi); + const uint32_t b = slot * SFB_SLOT_U + sfb_lo; + cp_sf(el, sfb0, b + 0u*TILE_U, sfb_hi); + cp_sf(el, sfb1, b + 3u*TILE_U, sfb_hi); + cp_sf(el, sfb2, b + 1u*TILE_U, sfb_hi); + cp_sf(el, sfb3, b + 4u*TILE_U, sfb_hi); + cp_sf(el, sfb4, b + 2u*TILE_U, sfb_hi); + cp_sf(el, sfb5, b + 5u*TILE_U, sfb_hi); +} + +// Operand descriptor LOW word for A/B ring slot `slot`, K-step `k_atom` +// (0..7 inside the 768-element group; each step is 48 bytes = +3 in the +// 16-byte address field). +static __device__ __forceinline__ uint32_t ab_desc_lo( + uint32_t smbase, uint32_t ab_base_off, uint32_t slot, uint32_t k_atom) { + const uint64_t d = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + slot * CD_AB_STRIDE, int(k_atom) * 16); + return uint32_t(d & 0xFFFFFFFFu); +} + +// K-steps 6 and 7 cross a 128-byte window boundary. Their descriptor address +// comes from `addr_slot`, while the absolute-LBO field comes from the next +// slot whose swizzle window contains the second half. +static __device__ __forceinline__ uint32_t ab_desc_lo_straddle( + uint32_t smbase, uint32_t ab_base_off, + uint32_t addr_slot, uint32_t lbo_slot, uint32_t k_atom) { + const uint64_t d_addr = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + addr_slot * CD_AB_STRIDE, int(k_atom) * 16); + const uint64_t d_lbo = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + lbo_slot * CD_AB_STRIDE, 0); + return uint32_t(d_addr & 0x0000FFFFu) | uint32_t(d_lbo & 0xFFFF0000u); +} + +// One output tile: `full_groups` complete 768-element K-groups. Coordinates +// past logical K are zero-filled by TMA from the unchanged canonical inputs. +static __device__ __forceinline__ MainloopState mainloop_tile( + uint32_t taddr, uint32_t smbase, MainloopState in, uint16_t commit_mask, + uint32_t full_groups + ) { + constexpr uint32_t AB_RING_SLOTS = uint32_t(CD_AB_SLOTS); + constexpr uint32_t SF_RING_SLOTS = uint32_t(CD_SF_SLOTS); + constexpr uint32_t AB_FREE_DELTA = uint32_t(MB_AB_FREE - MB_AB_READY); + constexpr uint32_t SF_READY_OFF = uint32_t(MB_SF_READY); + constexpr uint32_t SF_FREE_DELTA = uint32_t(MB_SF_FREE - MB_SF_READY); + constexpr uint32_t ACC_FREE_OFF = uint32_t(MB_ACC_FREE); + constexpr uint32_t A_OFF = uint32_t(CD_A_OFF); + constexpr uint32_t B_OFF = uint32_t(CD_B_OFF); + + uint32_t d_parity = in.d_parity; + uint32_t ab_phase = in.ab_phase, ab_slot = in.ab_slot; + uint32_t sf_phase = in.sf_phase, sf_slot = in.sf_slot; + + // The caller has already selected exactly one lane for the whole MMA role. + // Keep the helper predicate constant throughout the tile. + const uint32_t el = 1u; + + // Fixed TMEM targets for the single reused scale-factor region: 9 cp + // tiles of 4 columns each, starting past the accumulator region. + constexpr uint32_t SF = uint32_t(SF_TMEM_BASE); + const uint32_t sfa_t0 = taddr + SF + 0; // SFA word 0 (MMA reads here) + const uint32_t sfa_t1 = taddr + SF + 4; + const uint32_t sfa_t2 = taddr + SF + 8; + const uint32_t sfb_t0 = taddr + SF + 12; // SFB word 0 + const uint32_t sfb_t1 = taddr + SF + 16; + const uint32_t sfb_t2 = taddr + SF + 20; + const uint32_t sfb_t3 = taddr + SF + 24; + const uint32_t sfb_t4 = taddr + SF + 28; + const uint32_t sfb_t5 = taddr + SF + 32; + + // Scale-factor cp source descriptors (low word carries the per-slot + // offset; high word is constant). + const uint64_t sfb_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFB_OFF)); + const uint32_t sfb_src_lo = uint32_t(sfb_src & 0xFFFFFFFFu); + const uint32_t sfb_src_hi = uint32_t(sfb_src >> 32); + const uint64_t sfa_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFA_OFF)); + const uint32_t sfa_src_lo = uint32_t(sfa_src & 0xFFFFFFFFu); + const uint32_t sfa_src_hi = uint32_t(sfa_src >> 32); + + // Operand descriptor HIGH words (low words come from ab_desc_lo per step). + const uint32_t a_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + A_OFF) >> 32); + const uint32_t b_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + B_OFF) >> 32); + + // Each K=96 step needs 6 scale bytes per row, split across two staged + // 32-bit words. Odd-issue steps read word 1: bit 31 of the SF address + // selects the with-id encoding, and the id bits fold into the idesc. + constexpr uint32_t SF_ADDR_ID = 0x80000000u; + const uint32_t sfa_w1 = (taddr + SF + 4) | SF_ADDR_ID; + const uint32_t sfb_w1 = (taddr + SF + 20) | SF_ADDR_ID; + const uint32_t w1_id_bits = + ((sfa_w1 >> 1) & 0x60000000u) | ((sfb_w1 >> 26) & 48u); + + // The tile's first MMA overwrites the accumulator buffer this tile owns; + // the @216 "buffer free" wait runs once, before the first group. + const uint32_t acc_wait_parity = d_parity; + d_parity ^= 1u; + const uint32_t d_tmem = d_parity * uint32_t(D_STRIDE) + taddr; + + // Instruction descriptor: kind::mxf4nvf4, M=256, N=256, E2M1 x E2M1, + // UE4M3 block-16, K-major, k_dim=1 (K=96). Base value 0x90400480. + constexpr uint32_t IDESC = ptx::mma_inst_desc_mxf4nvf4_block16( + 256u, 256u, 0u, 0u, ptx::Major::K, ptx::Major::K, + false, false, /*k_dim=*/true, ptx::ScaleFormat::E4M3); + static_assert(IDESC == 0x90400480u, "K=96 NVFP4 idesc drifted"); + const uint32_t w0_id_bits = + ((sfa_t0 >> 1) & 0x60000000u) | ((sfb_t0 >> 26) & 48u); + const uint32_t idesc_w0 = w0_id_bits | IDESC; + const uint32_t idesc_w1 = w1_id_bits | IDESC; + + + uint32_t acc_pred = 0; // first MMA of the tile overwrites + bool need_acc_wait = true; // @216 wait fires once per tile + + if (full_groups >= 1) { + uint32_t g = 0; + do { + // ---- SF batch 0 -> logical steps 0,1 (atoms 0,3) ---- + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // ---- SF batch 1 -> logical steps 2,3 (atoms 6,1) ---- + mbar_wait_relaxed(sf_bar1, sf_phase1); + uint32_t sf_slot2, sf_phase2; + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + // AB0's last reader is MMA2. + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // ---- SF batch 2 -> logical steps 4,5 (atoms 4,7) ---- + mbar_wait_relaxed(sf_bar2, sf_phase2); + uint32_t sf_slot3, sf_phase3; + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; // ring cursor writeback + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // ---- SF batch 3 -> logical steps 6,7 (atoms 2,5) ---- + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; // ring cursor writeback + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + // AB1's last reader is MMA5; this measured placement leaves the + // commit behind the following scale copy. + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + acc_pred = 0xFFFFFFFFu; // accumulate from the second group on + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + need_acc_wait = false; + // AB2's last reader is MMA7. + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, + commit_mask); + ++g; + } while (g != full_groups); + } + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; +} + +} // namespace nvfp4 + +// ============================================================================ +// The kernel. 224 threads = 7 warps, four jobs (blog: the warp table): +// warps 0-3 epilogue: drain TMEM, convert to FP16, store straight to global +// warp 4 the MMA warp (even CTA of each pair only): SF copies + MMAs +// warp 5 TMA producer: A and B tiles into the 6-slot ring +// warp 6 TMA producer: scale factors into the 7-slot ring +// ============================================================================ +namespace nvfp4 { + +// Promise ptxas the exact block shape and express the launch grid in cluster +// units. The launch helper below owns the corresponding grid convention. +#define NVFP4_BLOCK_ATTR __block_size__((TB_SIZE, 1, 1)) +constexpr bool GRID_IN_CLUSTERS = true; + +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void nvfp4_gemm_kernel(const __grid_constant__ CUtensorMap A_tmap, + const __grid_constant__ CUtensorMap B_tmap, + const __grid_constant__ CUtensorMap SFA_tmap, + const __grid_constant__ CUtensorMap SFB_tmap, + __half* C_ptr, int M, int N, int K + ) { + extern __shared__ __align__(1024) char smem_buf[]; + SmemCD& smem = *reinterpret_cast(smem_buf); + + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const uint32_t crank = ptx::cluster_rank(); + const uint32_t m_pair = crank & 1u; // 0 = the pair's even CTA + + const int cluster_grid_m = (M + CLUSTER_M - 1) / CLUSTER_M; + const int grid_n = (N + BLOCK_N - 1) / BLOCK_N; + const int total_tiles = cluster_grid_m * grid_n; + + // Baseline rungs execute a whole K=768 group at the edge. Tensor-map OOB + // fill supplies zeros past logical K without changing the input buffers. + const int full_groups = (K + 767) / 768; + const int num_groups = full_groups; + + // Persistent cluster index. The device always sees the grid in CTA units + // (2, 1, clusters) regardless of the host-side launch-unit regime. + const bool grid_z = (gridDim.z > 1); + const int num_clusters = grid_z ? int(gridDim.z) : int(gridDim.x) / 2; + const int cluster_id = grid_z ? int(blockIdx.z) : int(blockIdx.x) / 2; + + + // Which output tile does persistent work item `t` process? + auto tile_of = [&](int t) -> int { + return t; // natural raster: M fast, N slow + }; + + auto producer_tile_of = [&](int t) -> int { + return tile_of(t); + }; + + // ---- mbarrier init (warp 0, lane 0) ---- + if (warp_id == 0 && ptx::lane_id() == 0) { + // Arrival counts are derived from the topology, never tuned: + // "ready" = one producer expect_tx per slot (bit-24 routes both + // peers' bytes to the even CTA's barrier); "free" = one consumer + // commit per slot reuse. + #pragma unroll + for (int s = 0; s < CD_AB_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.ab_ready[s], 1); + ptx::mbar_init(&smem.hdr.ab_free[s], 1); + } + #pragma unroll + for (int s = 0; s < CD_SF_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.sf_ready[s], 1); + ptx::mbar_init(&smem.hdr.sf_free[s], 1); + } + ptx::mbar_init(&smem.hdr.acc_ready, 1); + ptx::mbar_init(&smem.hdr.acc_free, N_EPI_WARPS * CTA_GROUP); // 4 x 2 = 8 + ptx::mbar_init(&smem.hdr.dealloc, 32); // one warp's lanes, cross-peer + ptx::fence_mbarrier_init_release_cluster(); + } + ptx::cluster_sync_rel_acq(); + + // TMEM allocation: epilogue warp 0 allocates for the pair; warps 0-4 + // rendezvous before reading the granted base. The TMA producer warps + // never touch TMEM and proceed straight to their loops. + uint32_t taddr = 0; + if (warp_id <= 4) { + if (warp_id == 0) + ptx::tcgen05_alloc_2sm(ptx::to_shared(&smem.hdr.tmem_addr), + N_ALLOC_COLS); + ptx::bar_sync(2, 160); + taddr = smem.hdr.tmem_addr; + } + + const int wg = warp_id >> 2; // 0 = epilogue warps, 1 = producers + MMA + if (wg == 1) { + const int pw = warp_id - 4; + if (pw == 1 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 5: lane 0 owns the complete A/B TMA producer loop. + // ================================================================ + // free_phase starts at 1: a fresh mbarrier is phase 0, so the + // first ring's worth of "free" waits falls through. + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + // The pair splits the 256x256 tile: even CTA holds A rows + // [0,128) and B cols [0,128); odd CTA the other halves. + const int m_blk = m_row * CTA_GROUP + int(m_pair); + const int n_blk = n_grp * CTA_GROUP + int(m_pair); + #pragma unroll 1 + for (int g = 0; g < num_groups; ++g) { + // Three full 128-byte windows cover the packed K=768 group. + #pragma unroll + for (int sub = 0; sub < AB_TMA_SUBS; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.ab_free[slot], free_phase); + const uint32_t a_s = + base + CD_A_OFF + uint32_t(slot) * CD_AB_STRIDE; + const uint32_t b_s = + base + CD_B_OFF + uint32_t(slot) * CD_AB_STRIDE; + uint64_t* bar = &smem.hdr.ab_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, AB_TX); + // A rides the multicast form with a self mask so the + // completion routes to the even peer's barrier. + const uint16_t mask_a = uint16_t(1u << m_pair); + // Read the same row-major FP4 allocation as cuBLAS. + // K192 uses four 96-byte boxes; K256 uses three + // 128-byte boxes. The tensor map zero-fills past K. + const CUtensorMap* a_map = &A_tmap; + const CUtensorMap* b_map = &B_tmap; + const int x = g * 384 + sub * AB_LOAD_BYTES; + const int y_a = m_blk * BLOCK_M; + const int y_b = n_blk * BLOCK_N_PEER; + const int z_a = 0; + const int z_b = 0; + ptx::cp_async_bulk_tensor_3d_load_multicast( + a_s, a_map, x, y_a, z_a, + bar, mask_a + ); + ptx::cp_async_bulk_tensor_3d_load_2sm_bit24( + b_s, b_map, x, y_b, z_b, + bar + ); + if (++slot == CD_AB_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + // Drain: wait for the last-filled slot to be consumed, then hand + // it one clean-shutdown expect_tx. + { + const int t_slot = (slot == 0) ? CD_AB_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.ab_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.ab_ready[t_slot], AB_TX); + } + } else if (pw == 2 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 6: scale-factor TMA producer. SFA is per-peer (each CTA's + // own M half); SFB is the one genuinely multicast operand — the + // cta_group::2 copy reads BOTH peers' smem as one logical + // operand, so both copies must be byte-identical (blog: the scale-factor section). + // ================================================================ + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const int m_blk = m_row * CTA_GROUP + int(m_pair); + for (int g = 0; g < num_groups; ++g) { + // 4 SF ring slots per 768-K group (2 K=96 steps each). + #pragma unroll + for (int sub = 0; sub < 4; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.sf_free[slot], free_phase); + const uint32_t sfa_s = + base + CD_SFA_OFF + uint32_t(slot) * CD_SFA_STRIDE; + const uint32_t sfb_s = + base + CD_SFB_OFF + uint32_t(slot) * CD_SFB_STRIDE; + uint64_t* bar = &smem.hdr.sf_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, SF_TX); + // VEC16 is [outer-block][K-group-of-4][128 rows x 4]. + // Three adjacent groups are the 128 x 12 scale slot. + const int sf_group = (g * 4 + sub) * 3; + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfa_s, &SFA_tmap, 0, 0, sf_group, m_blk, bar, + uint16_t(1u << m_pair) + ); + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfb_s + m_pair * uint32_t(CD_SFB_STRIDE / 2), + &SFB_tmap, 0, 0, sf_group, + n_grp * CTA_GROUP + int(m_pair), bar, + uint16_t(0x3) + ); + if (++slot == CD_SF_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + { + const int t_slot = (slot == 0) ? CD_SF_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.sf_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.sf_ready[t_slot], SF_TX); + } + } else if (pw == 0 && m_pair == 0 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 4: the MMA warp — issued once per cluster, on the even + // CTA. The odd CTA's warp 4 does nothing all kernel (blog: the warp table). + // The mainloop needs no tile coordinates: operands arrive through + // the rings; this warp only needs its share of the trip count. + // ================================================================ + const uint32_t smbase = ptx::to_shared(&smem); + const uint16_t commit_mask = 0x3; // both CTAs of this pair + MainloopState st{1u, 0u, 0u, 0u, 0u}; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + st = mainloop_tile(taddr, smbase, st, commit_mask, + uint32_t(full_groups) + ); + // The outer role branch already leaves exactly one issuer. + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;" + :: "r"(smbase + uint32_t(MB_ACC_READY)), + "h"(uint16_t(0x3)) : "memory"); + } + // Drain the final tile's buffer-free arrive before exiting. + mbar_wait_relaxed(smbase + uint32_t(MB_ACC_FREE), st.d_parity); + } + } else { + // ==================================================================== + // Warps 0-3: the epilogue (blog: the epilogue). Each warp owns a 32-row band + // and drains a tile as eight 32-column tcgen05.ld loads, converts + // FP32->FP16 in packed pairs, and stores straight to global — no + // shared-memory bounce. + // ==================================================================== + const int epi_warp = warp_id; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + const uint32_t acc_free_rank = crank & ~1u; // the pair's even CTA + int d_db = 0; + int acc_ready_phase = 0; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const uint32_t tmem_base = taddr + uint32_t(d_db) * D_STRIDE; + mbar_wait_relaxed(ptx::to_shared(&smem.hdr.acc_ready), + uint32_t(acc_ready_phase)); + acc_ready_phase ^= 1; + + const int off_m = (m_row * CTA_GROUP + int(m_pair)) * BLOCK_M; + const int off_n = n_grp * BLOCK_N; + const int row = off_m + epi_warp * 32 + lane_id; + const bool row_in = (row < M); + constexpr int EPI_BANDS = BLOCK_N / EPI_CHUNK; // 8 + #pragma unroll 1 + for (int k = 0; k < EPI_BANDS; ++k) { + // Buffer 0 drains high-to-low and buffer 1 low-to-high, so + // their shared columns [220,256) drain first for both. + const int parity = d_db & 1; + const int band = (parity == 0) ? (EPI_BANDS - 1 - k) : k; + const uint32_t taddr_n = + tmem_base + uint32_t(band) * 32 + taddr_lane; + uint32_t w[32]; + ptx::tcgen05_ld_32x32b_x32(taddr_n, w); + // Return the old buffer as soon as its shared edge is out. + if (k == 1) { + ptx::tcgen05_wait_ld_then_arrive_release( + &smem.hdr.acc_free, acc_free_rank, + /*do_arrive=*/ptx::lane_id() == 0); + } + if (!row_in) continue; + const int col_base = off_n + band * 32; + uint32_t h[16]; + #pragma unroll + for (int j = 0; j < 16; ++j) { + h[j] = ptx::cvt_pack_f16x2(__int_as_float(w[2 * j + 1]), + __int_as_float(w[2 * j])); + } + const size_t row_base = size_t(row) * size_t(N); + const bool aligned32 = ((row_base & 15u) == 0u) && + ((uintptr_t(C_ptr) & 31u) == 0u); + if (aligned32 && (col_base & 15) == 0) { + // Precondition here: row_base is 16-element (32-byte) + // aligned — the aligned32 guard owns it. Restating it to + // ptxas via __builtin_assume compiled this file + // SASS-identical (and measured perf-null on the + // production kernel), so the guard stands alone. + #pragma unroll + for (int q = 0; q < 2; ++q) { + const int c0 = col_base + q * 16; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 16 <= N) { + uint4 lo{h[q*8+0], h[q*8+1], h[q*8+2], h[q*8+3]}; + uint4 hi{h[q*8+4], h[q*8+5], h[q*8+6], h[q*8+7]}; + ptx::st_global_na_ef_b256(cdst, lo, hi); + } else { + // N tail: 128-bit then scalar. + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int cq = c0 + p * 8; + if (cq + 8 <= N) { + uint4 v{h[q*8+p*4+0], h[q*8+p*4+1], + h[q*8+p*4+2], h[q*8+p*4+3]}; + ptx::st_global_na_b128( + reinterpret_cast(cdst + p * 8), v); + } else { + const uint16_t* hb = + reinterpret_cast( + &h[q * 8 + p * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (cq + c < N) cdst[p * 8 + c] = hb[c]; + } + } + } + } + } else { + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int c0 = col_base + q * 8; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 8 <= N) { + uint4 v{h[q*4+0], h[q*4+1], h[q*4+2], h[q*4+3]}; + ptx::st_global_na_b128( + reinterpret_cast(cdst), v); + } else { + const uint16_t* hb = + reinterpret_cast(&h[q * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (c0 + c < N) cdst[c] = hb[c]; + } + } + } + } + if (++d_db == D_DB) d_db = 0; + } + } + + // Teardown: only the epilogue warps rendezvous per CTA; warp 0 then runs + // the cross-peer dealloc handshake. Producer and MMA warps exit freely — + // their drains above already order their work. + if (wg == 0) { + ptx::bar_sync(3, 128); + if (warp_id == 0) { + ptx::tcgen05_relinquish_2sm(); + const uint32_t peer_rank = crank ^ 1u; + // All 32 lanes arrive the PEER's count-32 barrier = one flip. + ptx::mbar_arrive_cluster_release(&smem.hdr.dealloc, peer_rank); + ptx::mbar_wait_parity(&smem.hdr.dealloc, 0); + ptx::tcgen05_dealloc_2sm(taddr, N_ALLOC_COLS); + } + } +} + + +// Keep the host launch convention paired with the kernel's block attribute. +static inline dim3 launch_grid(int clusters) { + return GRID_IN_CLUSTERS ? dim3(1, 1, unsigned(clusters)) + : dim3(2, 1, unsigned(clusters)); +} + +} // namespace nvfp4 + +// ============================================================================ +// Host-side TMA descriptor support. The shared input and benchmark code lives +// in main.cu. +// ============================================================================ +namespace tmap { + +// Thin wrappers over cuTensorMapEncodeTiled. Dim 0 is the innermost +// (stride-1) axis; strides for dim 1+ are in bytes. +inline CUtensorMap encode_tiled_3d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_depth, uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t depth_stride_bytes, + uint32_t box_depth, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[3] = { global_cols, global_rows, global_depth }; + cuuint64_t global_strides[2] = { row_stride_bytes, depth_stride_bytes }; + cuuint32_t box_dim[3] = { box_cols, box_rows, box_depth }; + cuuint32_t element_strides[3] = { 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 3, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +inline CUtensorMap encode_tiled_4d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_slots, uint64_t global_tiles, + uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t tile_stride_bytes, + uint64_t slot_stride_bytes, + uint32_t box_tiles, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[4] = { global_cols, global_rows, global_tiles, + global_slots }; + cuuint64_t global_strides[3] = { row_stride_bytes, tile_stride_bytes, + slot_stride_bytes }; + cuuint32_t box_dim[4] = { box_cols, box_rows, box_tiles, 1u }; + cuuint32_t element_strides[4] = { 1, 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 4, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +} // namespace tmap + +namespace host { + +using namespace nvfp4; + +// ---- implementation-specific tensor maps --------------------------------- + +inline CUtensorMap make_ab_tmap(void* ptr, int rows, int K, + int row_stride) { + return tmap::encode_tiled_3d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*depth=*/1, uint64_t(rows), /*cols=*/uint64_t(K / 2), + /*row_stride=*/uint64_t(row_stride), + /*depth_stride=*/uint64_t(rows) * uint64_t(row_stride), + /*box=*/1, BLOCK_M, AB_LOAD_BYTES, CU_TENSOR_MAP_SWIZZLE_128B); +} + +inline CUtensorMap make_sf_tmap(uint8_t* ptr, int outer, int K) { + const uint64_t sf_inner = (uint64_t(K) + 63) / 64 * 4; + const uint64_t outer_blocks = (uint64_t(outer) + 127) / 128; + return tmap::encode_tiled_4d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*slots=*/outer_blocks, /*tiles=*/sf_inner / 4, + /*rows=*/4, /*cols=*/128, + /*row_stride=*/128, /*tile_stride=*/512, + /*slot_stride=*/sf_inner * 128, + /*box_tiles=*/3, /*box_rows=*/4, /*box_cols=*/128, + CU_TENSOR_MAP_SWIZZLE_NONE); +} +} // namespace host diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/gemm7.cuh b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm7.cuh new file mode 100644 index 000000000..fff64507e --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm7.cuh @@ -0,0 +1,1421 @@ +#pragma once +// r7: feed the exact K tail. +// This is a complete, directly editable ladder snapshot. +// It intentionally duplicates the implementation in the other gemm*.cuh files. + +// Conditional branches are already materialized; edit this implementation directly. +#define NVFP4_BUILD_NAME "r7" +// Read-only metadata for main.cu; these do not select code in this header. +#define NVFP4_HAS_K64_TAIL 0 +#define NVFP4_HAS_SCHEDULE 0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(expr) do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \ + cudaGetErrorName(_e), __FILE__, __LINE__, \ + cudaGetErrorString(_e)); \ + std::abort(); \ + } \ +} while (0) + +#define CU_CHECK(expr) do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + std::fprintf(stderr, "CU error at %s:%d: %s\n", \ + __FILE__, __LINE__, _s ? _s : "?"); \ + std::abort(); \ + } \ +} while (0) + + +// ============================================================================ +// PTX helpers. Thin zero-cost wrappers over the exact instruction forms the +// kernel needs — nothing generic, nothing speculative. +// ============================================================================ +namespace ptx { + +// Generic pointer -> 32-bit `.shared` window byte address. +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Local `.shared` address -> the same offset in CTA `cta_rank`'s smem +// (`.shared::cluster` instructions take this mapped form). +static __device__ __forceinline__ uint32_t mapa_shared_cluster( + uint32_t local_addr, uint32_t cta_rank) { + uint32_t mapped; + asm("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(mapped) : "r"(local_addr), "r"(cta_rank)); + return mapped; +} + +static __device__ __forceinline__ void bar_sync(uint32_t id, uint32_t count) { + asm volatile("barrier.cta.sync %0, %1;" :: "r"(id), "r"(count)); +} + +// Cluster barrier with release/acquire — the publish boundary between the +// mbarrier inits and any cross-CTA use of those mbarriers. +static __device__ __forceinline__ void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// %cluster_ctarank: 0 or 1 in this kernel's (2,1,1) clusters. +static __device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} + +static __device__ __forceinline__ uint32_t lane_id() { + uint32_t r; + asm("mov.u32 %0, %%laneid;" : "=r"(r)); + return r; +} + +// ---- mbarriers ------------------------------------------------------------- + +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + :: "r"(to_shared(bar)), "r"(count)); +} + +// Publish the inits to the async proxy + the peer CTA before any TMA +// completes on them. +static __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// Cross-CTA arrive with release ordering (teardown handshake). +static __device__ __forceinline__ void mbar_arrive_cluster_release( + uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" + :: "r"(mapped)); +} + +// Fused `tcgen05.wait::ld` + "buffer free" arrive, one asm block. The arrive +// must not become visible while TMEM loads are still in flight (it frees the +// accumulator buffer to the next tile's MMA — the bring-up chapter's +// drain overlap); the fused block +// keeps the arrive pinned behind the wait. (An earlier form also manufactured +// a register dependency on a drained value, `dep * 0` folded into the +// address; verified dead on nvcc 13.1 — ptxas const-folds it, SASS +// byte-identical without it — and retired.) +static __device__ __forceinline__ void tcgen05_wait_ld_then_arrive_release( + uint64_t* bar, uint32_t cta_rank, bool do_arrive) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile( + "{\n\t" + ".reg .pred %%pdo;\n\t" + "setp.ne.s32 %%pdo, %1, 0;\n\t" + "tcgen05.wait::ld.sync.aligned;\n\t" + "@%%pdo mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;\n\t" + "}\n\t" + :: "r"(mapped), "r"((int)do_arrive) : "memory"); +} + + +// Blocking parity wait at the consuming site (producer ring/tail waits + +// teardown). The 10 ms suspend hint is the polite form: a full-rate spin +// competes for issue slots and measurably loses; a nanosleep-backoff spin +// measured the same as this hint across five windows, so the simpler form +// ships. The carried one-stage-ahead try_wait peeks that used to gate these +// waits measured as pure overhead and were removed. +static __device__ __forceinline__ void mbar_wait_parity( + uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(to_shared(bar)), "r"(parity)); +} + +// ---- TMA ------------------------------------------------------------------- + +// Producer expect_tx arrive on the cluster-shared ring barrier: both +// peers' transfer bytes account against the even CTA's barrier, so one +// barrier per ring slot covers the whole cluster-wide gather. (Mechanism: +// `& 0xFEFFFFFF` clears bit 24 of the shared address — the hardware's +// peer-CTA bit — re-pointing the arrive at the even CTA's copy.) +static __device__ __forceinline__ void mbar_arrive_expect_tx_cluster( + uint64_t* bar, uint32_t bytes) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes)); +} + +// 3D tiled load, cta_group::2 + multicast (the A feed; the mask is the +// issuing CTA itself in a (2,1) cluster). Completion routes to the even +// CTA's mbarrier via the bit-24 clear. +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z) + : "memory"); +} + +// 3D tiled load, cta_group::2, no multicast (the B feed — each peer loads +// its own 128-wide N-block). +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_2sm_bit24( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst_smem), "l"(tmap), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) + : "memory"); +} + +// 4D tiled load, cta_group::2 + multicast (the scale-factor feed). The +// 512-byte SF tiles of one ring slot form a 4D box; `w` selects the slot. +static __device__ __forceinline__ void cp_async_bulk_tensor_4d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z2, int32_t w, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.4d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z2), "r"(w) + : "memory"); +} + + +// ---- tcgen05 (tensor memory + MMA plumbing) --------------------------------- + +static __device__ __forceinline__ void tcgen05_alloc_2sm( + uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_dealloc_2sm( + uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" + :: "r"(taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_relinquish_2sm() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); +} + +// The widest single-instruction TMEM read: 32 b32 registers per lane +// (blog: the epilogue — a quarter the load/wait count of the narrow forms). +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t taddr, uint32_t* r) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), + "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), + "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), + "=r"(r[16]), "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), + "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), + "=r"(r[24]), "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), + "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr)); +} + +// ---- stores + conversion ----------------------------------------------------- + +static __device__ __forceinline__ uint32_t cvt_pack_f16x2(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// L2::evict_first requires a 256-bit store, but the 128-bit tail can still +// adopt the no-L1-allocation part of the policy at this rung. +static __device__ __forceinline__ void st_global_na_b128( + uint4* ptr, uint4 v) { + asm volatile("st.global.L1::no_allocate.v4.b32 " + "[%0], {%1, %2, %3, %4};" + :: "l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)); +} + +// Add the output cache policy to the 256-bit store: bypass L1 and evict the +// write-once line first from L2 so it does not displace reusable operands. +static __device__ __forceinline__ void st_global_na_ef_b256( + void* ptr, uint4 lo, uint4 hi) { + asm volatile("st.global.L1::no_allocate.L2::evict_first.v8.b32 [%0], " + "{%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"(ptr), "r"(lo.x), "r"(lo.y), "r"(lo.z), "r"(lo.w), + "r"(hi.x), "r"(hi.y), "r"(hi.z), "r"(hi.w)); +} + + +// ---- MMA descriptors --------------------------------------------------------- + +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class ScaleFormat : uint8_t { E4M3 = 0, E8M0 = 1 }; + +// Source descriptor for the tcgen05.cp that stages scale factors from smem +// into TMEM: unswizzled, SBO=128, LBO=16. The low word later carries the +// per-slot byte offset (in 16-byte units). +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_sf_cp_src_desc_cd(uint32_t sf_smem_addr) { + uint64_t addr_field = uint64_t((sf_smem_addr >> 4) & 0x3FC0u); + return addr_field | 0x400800010000ull; +} + +// Operand descriptor for a K=96 (48-byte) slice inside a 128-byte-swizzled +// smem K-window. OR-constant 0x4010404000000000: version=1, absolute-address +// LBO mode (bit 52), 128B swizzle, SBO=1024. Both the address field and the +// absolute-LBO field encode the window's own 1024-aligned base; the +// intra-window offset advances the address field only. +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k96_sw128(uint32_t window_base_1024aligned, int intra_window_off = 0) { + uint64_t addr_field = uint64_t(((window_base_1024aligned >> 4) & 0x7FC0u) + + uint32_t(intra_window_off >> 4)); + uint64_t lbo_field = uint64_t((window_base_1024aligned << 12) & 0x7FC00000u); + return addr_field | lbo_field | 0x4010404000000000ull; +} + +// Instruction descriptor for tcgen05.mma.kind::mxf4nvf4.block_scale. +// The K extent is selected by BIT 31, not by the asm mnemonic: k_dim=0 runs +// K=64 (0x10400480), k_dim=1 runs K=96 (0x90400480, sm_103a-only) — the +// cheapest quarter of the gap if your kernel is built on the K=64 form +// (blog: the MMA warp). Bit 13 flips which TMEM scale-factor column the MMA reads. +__host__ __device__ static __forceinline__ constexpr uint32_t +mma_inst_desc_mxf4nvf4_block16( + uint32_t M, uint32_t N, + uint32_t tmem_sfa_addr, uint32_t tmem_sfb_addr, + Major a_major, Major b_major, + bool negate_a, bool negate_b, + bool k_dim, ScaleFormat sf) { + constexpr uint32_t MXF4_E2M1 = 1u; + uint32_t d = 0; + d |= ((tmem_sfb_addr & 0xC0000000u) >> 30) << 4; // b_sf_id + d |= MXF4_E2M1 << 7; // a_format = E2M1 + d |= MXF4_E2M1 << 10; // b_format = E2M1 + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; // N in units of 8 + d |= (static_cast(sf) & 0x1u) << 23; // 0 = UE4M3 + d |= ((M >> 4) & 0x1Fu) << 24; // m_dim at bits 27-28 + d |= ((tmem_sfa_addr & 0xC0000000u) >> 30) << 29; // a_sf_id + if (k_dim) d |= 1u << 31; // 1 = K=96 + return d; +} + +} // namespace ptx + +// ============================================================================ +// Kernel geometry from the blog's opening kernel section. One tile shape, +// one cluster shape: +// 256x256 output per 2-CTA cluster (128 rows per CTA), K advanced in +// 96-element MMA steps staged as 768-element K-groups of 8 steps, persistent +// grid of 76 clusters = 152 CTAs = one per SM. +// ============================================================================ +namespace nvfp4 { + +constexpr int CTA_GROUP = 2; // cta_group::2 — the MMA spans both CTAs +constexpr int BLOCK_M = 128; // A rows per CTA +constexpr int BLOCK_N = 256; // output tile N extent (per cluster) +constexpr int BLOCK_N_PEER = BLOCK_N / CTA_GROUP; // 128 — B cols per CTA +constexpr int CLUSTER_M = CTA_GROUP * BLOCK_M; // 256 — tile M extent +[[maybe_unused]] constexpr int MMA_K = 96; // K per tcgen05.mma (idesc bit 31) +constexpr int WIN_BYTES = 128; // one 128B-swizzled smem K-window +constexpr int AB_LOAD_BYTES = WIN_BYTES; +constexpr int AB_TMA_SUBS = 3; +constexpr int TB_SIZE = 224; // 7 warps (blog: the warp table) +constexpr int N_EPI_WARPS = 4; // warps 0-3 + +constexpr int SF_TMEM_COLS = 9 * 4; // 3 SFA + 6 SFB cp tiles, 4 cols each +// Two 256-wide accumulators use a 220-column stride and overlap by the 36 +// columns needed for scale factors. +constexpr int D_STRIDE = 220; +constexpr int D_DB = 2; +constexpr int N_D_COLS = D_STRIDE + BLOCK_N; +constexpr int SF_TMEM_BASE = 476; +constexpr int N_ALLOC_COLS = 512; +static_assert(N_D_COLS <= SF_TMEM_BASE && + SF_TMEM_BASE + SF_TMEM_COLS <= N_ALLOC_COLS, + "accumulators + scale factors exceed the 512 TMEM columns"); +constexpr int EPI_CHUNK = 32; // cols per regular drain + +// ---- the shared-memory map (blog: the feed) ----------------------------------- +// Every offset is a compile-time expression. +constexpr int CD_HDR = 1024; // mbarrier header pad +constexpr int CD_AB_SLOTS = 6; +constexpr int CD_SF_SLOTS = 7; +constexpr int CD_AB_STRIDE = BLOCK_M * WIN_BYTES; // 16384 per A or B slot +constexpr int CD_SF_TILE = 512; // one (32 lanes x 16B) x4 tile +constexpr int CD_SFA_TILES = 3; // per SFA slot (M = 1 block) +constexpr int CD_SFB_TILES = 6; // per SFB slot (N = 2 blocks) +constexpr int CD_SFA_STRIDE = CD_SFA_TILES * CD_SF_TILE; // 1536 +constexpr int CD_SFB_STRIDE = CD_SFB_TILES * CD_SF_TILE; // 3072 + +// Expected transfer bytes are derived, never tuned: a wrong-by-one +// expectation doesn't crash, it deadlocks (blog: the feed). +constexpr uint32_t AB_TX = uint32_t(CTA_GROUP) * + uint32_t(BLOCK_M + BLOCK_N_PEER) * uint32_t(AB_LOAD_BYTES); +constexpr uint32_t SF_TX = uint32_t(CD_SFA_STRIDE + CD_SFB_STRIDE) * + uint32_t(CTA_GROUP); // 9216 +static_assert(AB_TX == 65536u && SF_TX == 9216u, + "expect_tx derivation drifted"); + +// The map IS the struct: member sizes and alignment produce every byte +// offset; the raw-address constants the mainloop uses read off it with +// offsetof, and the static_asserts pin the published values. +struct SmemCD { + // mbarrier header: four ring regions in handshake order, then the + // singleton accumulator + teardown barriers (alignas pads it to 1024). + struct alignas(1024) Header { + uint64_t ab_ready[CD_AB_SLOTS]; // TMA -> MMA + uint64_t ab_free [CD_AB_SLOTS]; // MMA -> TMA + uint64_t sf_ready[CD_SF_SLOTS]; // TMA -> MMA + uint64_t sf_free [CD_SF_SLOTS]; // MMA -> TMA + uint64_t acc_ready; // MMA -> epi + uint64_t acc_free; // epi -> MMA + uint64_t dealloc; // 2-CTA teardown + uint32_t tmem_addr; // TMEM base addr storage + } hdr; + uint8_t a [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t b [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t sfa[CD_SF_SLOTS][CD_SFA_STRIDE]; + // SFB starts 1024-aligned (the cp source descriptor drops address bits 4-9). + alignas(1024) uint8_t sfb[CD_SF_SLOTS][CD_SFB_STRIDE]; +}; +static_assert(sizeof(SmemCD::Header) == CD_HDR, "mbarrier header overflows the pad"); + +constexpr int CD_A_OFF = offsetof(SmemCD, a); +constexpr int CD_B_OFF = offsetof(SmemCD, b); +constexpr int CD_SFA_OFF = offsetof(SmemCD, sfa); +constexpr int CD_SFB_OFF = offsetof(SmemCD, sfb); +constexpr int CD_TOTAL = sizeof(SmemCD); +static_assert(CD_TOTAL <= 232448, "smem map exceeds the opt-in budget"); +static_assert(CD_AB_SLOTS != 6 || + (CD_B_OFF == 99328 && CD_SFA_OFF == 197632 && + CD_SFB_OFF == 208896 && CD_TOTAL == 230400), + "the 6-slot byte map drifted from its published values"); + +// The mainloop computes raw `.shared` byte addresses from these. +constexpr int MB_AB_READY = offsetof(SmemCD, hdr.ab_ready); +constexpr int MB_AB_FREE = offsetof(SmemCD, hdr.ab_free); +constexpr int MB_SF_READY = offsetof(SmemCD, hdr.sf_ready); +constexpr int MB_SF_FREE = offsetof(SmemCD, hdr.sf_free); +constexpr int MB_ACC_READY = offsetof(SmemCD, hdr.acc_ready); +constexpr int MB_ACC_FREE = offsetof(SmemCD, hdr.acc_free); + + + +} // namespace nvfp4 + +// ============================================================================ +// The consumer mainloop from the blog's pipeline section. One warp — warp 4 +// of the even CTA — +// issues everything: the tcgen05.cp copies that stage scale factors from the +// smem ring into TMEM, immediately followed by the tcgen05.mma stream, in the +// SAME warp. A same-warp tcgen05 stream is one ordered async pipeline, so the +// write-after-read hazard on the single reused scale-factor TMEM region needs +// no fence and no cross-warp barrier. +// +// Per 768-element K-group, logical K=96 steps 0..7 issue in K order. Their +// starting offsets, expressed as 16-byte descriptor atoms modulo a 128-byte +// swizzle window, are 0,3,6,1,4,7,2,5. Atoms 6 and 7 straddle windows. +// ============================================================================ +namespace nvfp4 { + + +// Ring/parity state carried across tiles by the MMA warp. +struct MainloopState { + uint32_t d_parity; // which TMEM accumulator buffer the next tile writes + uint32_t ab_phase, ab_slot; // A/B ring cursor + uint32_t sf_phase, sf_slot; // scale-factor ring cursor +}; + + +// Relaxed blocking wait (no .acquire): the tcgen05 pipeline's own ordering +// covers these ring waits, and the .acquire form perturbs the scoreboard +// cadence on the issuing warp. Waits sit AT the consuming site — the carried +// one-stage-ahead try_wait peeks that used to gate them measured as pure +// overhead and were removed. +static __device__ __forceinline__ void mbar_wait_relaxed( + uint32_t bar_smem, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(bar_smem), "r"(parity) : "memory"); +} + +// Gated form for the @216 buffer-free wait, which fires once per tile: the +// gate branches over the spin for every group past the first. +static __device__ __forceinline__ void mbar_wait_gated_relaxed( + uint32_t bar_smem, uint32_t parity, bool peek) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p bra DONE_%=;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t" + "DONE_%=:\n\t}\n" + :: "r"(bar_smem), "r"(parity), "r"((uint32_t)peek) : "memory"); +} + + +// One scale-factor staging copy, smem -> TMEM. Single-lane semantics behind +// the converged elect predicate: tcgen05.cp is a single-thread instruction — +// issued bare from a converged warp you get 32 initiations and garbage. +static __device__ __forceinline__ void cp_sf( + uint32_t el, uint32_t dst, uint32_t lo, uint32_t hi) { + const uint64_t desc = (uint64_t(hi) << 32) | uint64_t(lo); + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;\n\t}\n" + :: "r"(dst), "l"(desc), "r"(el) : "memory"); +} + +// Ring-slot "consumed" arrive, multicast to both CTAs of the pair. +// arrive::one must fire exactly once — hence the elect predicate. +static __device__ __forceinline__ void commit_consumed( + uint32_t el, uint32_t smbase, uint32_t mb_off, uint16_t mask) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t}\n" + :: "r"(smbase + mb_off), "h"(mask), "r"(el) : "memory"); +} + +// One K=96 NVFP4 MMA. The asm token says scale_vec::4X regardless — K=96 is +// selected by idesc bit 31 (blog: the MMA warp). acc_pred = 0 overwrites the +// accumulator (the tile's first step), nonzero accumulates. +static __device__ __forceinline__ void mma_cell( + uint32_t el, uint32_t d, + uint32_t desc_a_lo, uint32_t desc_a_hi, + uint32_t desc_b_lo, uint32_t desc_b_hi, + uint32_t idesc, uint32_t sf_a, uint32_t sf_b, uint32_t acc_pred) { + const uint64_t desc_a = (uint64_t(desc_a_hi) << 32) | uint64_t(desc_a_lo); + const uint64_t desc_b = (uint64_t(desc_b_hi) << 32) | uint64_t(desc_b_lo); + asm volatile( + "{\n\t.reg .pred p, acc;\n\t" + "setp.ne.u32 p, %7, 0;\n\t" + "setp.ne.b32 acc, %4, 0;\n\t" + "@p tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X" + " [%0], %1, %2, %3, [%5], [%6], acc;\n\t}\n" + :: "r"(d), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(acc_pred), "r"(sf_a), "r"(sf_b), "r"(el) : "memory"); +} + +// One SF ring slot -> TMEM: 9 tcgen05.cp tiles (3 SFA + 6 SFB). Source +// address advances are in 16-byte descriptor units. The SFB tile order +// 0,3,1,4,2,5 interleaves the two 128-wide N-blocks. +static __device__ __forceinline__ void cp_sf_buffer( + uint32_t el, uint32_t slot, + uint32_t sfa0, uint32_t sfa1, uint32_t sfa2, + uint32_t sfb0, uint32_t sfb1, uint32_t sfb2, + uint32_t sfb3, uint32_t sfb4, uint32_t sfb5, + uint32_t sfa_hi, uint32_t sfa_lo, uint32_t sfb_hi, uint32_t sfb_lo) { + constexpr uint32_t SFA_SLOT_U = CD_SFA_STRIDE / 16u; // 96 + constexpr uint32_t SFB_SLOT_U = CD_SFB_STRIDE / 16u; // 192 + constexpr uint32_t TILE_U = CD_SF_TILE / 16u; // 32 + const uint32_t a = slot * SFA_SLOT_U + sfa_lo; + cp_sf(el, sfa0, a, sfa_hi); + cp_sf(el, sfa1, a + 1u*TILE_U, sfa_hi); + cp_sf(el, sfa2, a + 2u*TILE_U, sfa_hi); + const uint32_t b = slot * SFB_SLOT_U + sfb_lo; + cp_sf(el, sfb0, b + 0u*TILE_U, sfb_hi); + cp_sf(el, sfb1, b + 3u*TILE_U, sfb_hi); + cp_sf(el, sfb2, b + 1u*TILE_U, sfb_hi); + cp_sf(el, sfb3, b + 4u*TILE_U, sfb_hi); + cp_sf(el, sfb4, b + 2u*TILE_U, sfb_hi); + cp_sf(el, sfb5, b + 5u*TILE_U, sfb_hi); +} + +// Operand descriptor LOW word for A/B ring slot `slot`, K-step `k_atom` +// (0..7 inside the 768-element group; each step is 48 bytes = +3 in the +// 16-byte address field). +static __device__ __forceinline__ uint32_t ab_desc_lo( + uint32_t smbase, uint32_t ab_base_off, uint32_t slot, uint32_t k_atom) { + const uint64_t d = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + slot * CD_AB_STRIDE, int(k_atom) * 16); + return uint32_t(d & 0xFFFFFFFFu); +} + +// K-steps 6 and 7 cross a 128-byte window boundary. Their descriptor address +// comes from `addr_slot`, while the absolute-LBO field comes from the next +// slot whose swizzle window contains the second half. +static __device__ __forceinline__ uint32_t ab_desc_lo_straddle( + uint32_t smbase, uint32_t ab_base_off, + uint32_t addr_slot, uint32_t lbo_slot, uint32_t k_atom) { + const uint64_t d_addr = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + addr_slot * CD_AB_STRIDE, int(k_atom) * 16); + const uint64_t d_lbo = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + lbo_slot * CD_AB_STRIDE, 0); + return uint32_t(d_addr & 0x0000FFFFu) | uint32_t(d_lbo & 0xFFFF0000u); +} + +// One output tile: complete 768-element K-groups followed by one partial group +// containing `tail_cells` live K=96 steps. Dead cells fold their liveness into +// the issue predicate, so the warp stays converged. +static __device__ __forceinline__ MainloopState mainloop_tile( + uint32_t taddr, uint32_t smbase, MainloopState in, uint16_t commit_mask, + uint32_t full_groups, + uint32_t tail_cells + ) { + constexpr uint32_t AB_RING_SLOTS = uint32_t(CD_AB_SLOTS); + constexpr uint32_t SF_RING_SLOTS = uint32_t(CD_SF_SLOTS); + constexpr uint32_t AB_FREE_DELTA = uint32_t(MB_AB_FREE - MB_AB_READY); + constexpr uint32_t SF_READY_OFF = uint32_t(MB_SF_READY); + constexpr uint32_t SF_FREE_DELTA = uint32_t(MB_SF_FREE - MB_SF_READY); + constexpr uint32_t ACC_FREE_OFF = uint32_t(MB_ACC_FREE); + constexpr uint32_t A_OFF = uint32_t(CD_A_OFF); + constexpr uint32_t B_OFF = uint32_t(CD_B_OFF); + + uint32_t d_parity = in.d_parity; + uint32_t ab_phase = in.ab_phase, ab_slot = in.ab_slot; + uint32_t sf_phase = in.sf_phase, sf_slot = in.sf_slot; + + // The caller has already selected exactly one lane for the whole MMA role. + // Keep the helper predicate constant throughout the tile. + const uint32_t el = 1u; + + // Fixed TMEM targets for the single reused scale-factor region: 9 cp + // tiles of 4 columns each, starting past the accumulator region. + constexpr uint32_t SF = uint32_t(SF_TMEM_BASE); + const uint32_t sfa_t0 = taddr + SF + 0; // SFA word 0 (MMA reads here) + const uint32_t sfa_t1 = taddr + SF + 4; + const uint32_t sfa_t2 = taddr + SF + 8; + const uint32_t sfb_t0 = taddr + SF + 12; // SFB word 0 + const uint32_t sfb_t1 = taddr + SF + 16; + const uint32_t sfb_t2 = taddr + SF + 20; + const uint32_t sfb_t3 = taddr + SF + 24; + const uint32_t sfb_t4 = taddr + SF + 28; + const uint32_t sfb_t5 = taddr + SF + 32; + + // Scale-factor cp source descriptors (low word carries the per-slot + // offset; high word is constant). + const uint64_t sfb_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFB_OFF)); + const uint32_t sfb_src_lo = uint32_t(sfb_src & 0xFFFFFFFFu); + const uint32_t sfb_src_hi = uint32_t(sfb_src >> 32); + const uint64_t sfa_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFA_OFF)); + const uint32_t sfa_src_lo = uint32_t(sfa_src & 0xFFFFFFFFu); + const uint32_t sfa_src_hi = uint32_t(sfa_src >> 32); + + // Operand descriptor HIGH words (low words come from ab_desc_lo per step). + const uint32_t a_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + A_OFF) >> 32); + const uint32_t b_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + B_OFF) >> 32); + + // Each K=96 step needs 6 scale bytes per row, split across two staged + // 32-bit words. Odd-issue steps read word 1: bit 31 of the SF address + // selects the with-id encoding, and the id bits fold into the idesc. + constexpr uint32_t SF_ADDR_ID = 0x80000000u; + const uint32_t sfa_w1 = (taddr + SF + 4) | SF_ADDR_ID; + const uint32_t sfb_w1 = (taddr + SF + 20) | SF_ADDR_ID; + const uint32_t w1_id_bits = + ((sfa_w1 >> 1) & 0x60000000u) | ((sfb_w1 >> 26) & 48u); + + // The tile's first MMA overwrites the accumulator buffer this tile owns; + // the @216 "buffer free" wait runs once, before the first group. + const uint32_t acc_wait_parity = d_parity; + d_parity ^= 1u; + const uint32_t d_tmem = d_parity * uint32_t(D_STRIDE) + taddr; + + // Instruction descriptor: kind::mxf4nvf4, M=256, N=256, E2M1 x E2M1, + // UE4M3 block-16, K-major, k_dim=1 (K=96). Base value 0x90400480. + constexpr uint32_t IDESC = ptx::mma_inst_desc_mxf4nvf4_block16( + 256u, 256u, 0u, 0u, ptx::Major::K, ptx::Major::K, + false, false, /*k_dim=*/true, ptx::ScaleFormat::E4M3); + static_assert(IDESC == 0x90400480u, "K=96 NVFP4 idesc drifted"); + const uint32_t w0_id_bits = + ((sfa_t0 >> 1) & 0x60000000u) | ((sfb_t0 >> 26) & 48u); + const uint32_t idesc_w0 = w0_id_bits | IDESC; + const uint32_t idesc_w1 = w1_id_bits | IDESC; + + + uint32_t acc_pred = 0; // first MMA of the tile overwrites + bool need_acc_wait = true; // @216 wait fires once per tile + + if (full_groups >= 1) { + uint32_t g = 0; + do { + // ---- SF batch 0 -> logical steps 0,1 (atoms 0,3) ---- + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // ---- SF batch 1 -> logical steps 2,3 (atoms 6,1) ---- + mbar_wait_relaxed(sf_bar1, sf_phase1); + uint32_t sf_slot2, sf_phase2; + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + // AB0's last reader is MMA2. + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // ---- SF batch 2 -> logical steps 4,5 (atoms 4,7) ---- + mbar_wait_relaxed(sf_bar2, sf_phase2); + uint32_t sf_slot3, sf_phase3; + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; // ring cursor writeback + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // ---- SF batch 3 -> logical steps 6,7 (atoms 2,5) ---- + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; // ring cursor writeback + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + // AB1's last reader is MMA5; this measured placement leaves the + // commit behind the following scale copy. + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + acc_pred = 0xFFFFFFFFu; // accumulate from the second group on + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + need_acc_wait = false; + // AB2's last reader is MMA7. + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, + commit_mask); + ++g; + } while (g != full_groups); + } + + if (tail_cells == 0u) + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; + + + // ---- the exact-K tail group: tail_cells live K=96 steps out of 8 ---- + // Cell liveness rides the elect predicate (a 0-predicate op never enters + // the tcgen05 pipe); SF batch b exists iff tail_cells > 2b, and skipped + // batches advance NO ring cursor — the SF producer skips symmetrically. + const uint32_t el_c1 = (tail_cells >= 2u) ? el : 0u; + const uint32_t el_c2 = (tail_cells >= 3u) ? el : 0u; + const uint32_t el_c3 = (tail_cells >= 4u) ? el : 0u; + const uint32_t el_c4 = (tail_cells >= 5u) ? el : 0u; + const uint32_t el_c5 = (tail_cells >= 6u) ? el : 0u; + const uint32_t el_c6 = (tail_cells >= 7u) ? el : 0u; + const uint32_t el_c7 = (tail_cells >= 8u) ? el : 0u; + const bool b1_live = (tail_cells > 2u); + const bool b2_live = (tail_cells > 4u); + const bool b3_live = (tail_cells > 6u); + + // SF batch 0 (cells 0,1) — a live tail always has cell 0. + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el_c1, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // SF batch 1 (cells 2,3). + uint32_t sf_slot2 = sf_slot1, sf_phase2 = sf_phase1; + if (b1_live) { + mbar_wait_relaxed(sf_bar1, sf_phase1); + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + } + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el_c2, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, commit_mask); + mma_cell(el_c3, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // SF batch 2 (cells 4,5). + uint32_t sf_slot3 = sf_slot2, sf_phase3 = sf_phase2; + if (b2_live) { + mbar_wait_relaxed(sf_bar2, sf_phase2); + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + } + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el_c4, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el_c5, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // SF batch 3 (cells 6,7). + if (b3_live) { + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + } else { + sf_slot = sf_slot3; + sf_phase = sf_phase3; + } + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, commit_mask); + mma_cell(el_c6, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + mma_cell(el_c7, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, commit_mask); + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; +} + +} // namespace nvfp4 + +// ============================================================================ +// The kernel. 224 threads = 7 warps, four jobs (blog: the warp table): +// warps 0-3 epilogue: drain TMEM, convert to FP16, store straight to global +// warp 4 the MMA warp (even CTA of each pair only): SF copies + MMAs +// warp 5 TMA producer: A and B tiles into the 6-slot ring +// warp 6 TMA producer: scale factors into the 7-slot ring +// ============================================================================ +namespace nvfp4 { + +// Promise ptxas the exact block shape and express the launch grid in cluster +// units. The launch helper below owns the corresponding grid convention. +#define NVFP4_BLOCK_ATTR __block_size__((TB_SIZE, 1, 1)) +constexpr bool GRID_IN_CLUSTERS = true; + +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void nvfp4_gemm_kernel(const __grid_constant__ CUtensorMap A_tmap, + const __grid_constant__ CUtensorMap B_tmap, + const __grid_constant__ CUtensorMap SFA_tmap, + const __grid_constant__ CUtensorMap SFB_tmap, + __half* C_ptr, int M, int N, int K + ) { + extern __shared__ __align__(1024) char smem_buf[]; + SmemCD& smem = *reinterpret_cast(smem_buf); + + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const uint32_t crank = ptx::cluster_rank(); + const uint32_t m_pair = crank & 1u; // 0 = the pair's even CTA + + const int cluster_grid_m = (M + CLUSTER_M - 1) / CLUSTER_M; + const int grid_n = (N + BLOCK_N - 1) / BLOCK_N; + const int total_tiles = cluster_grid_m * grid_n; + + // Exact-K tail geometry (the exact-K feed, Other shapes; the host gates K % 16 == 0). + const int full_groups = K / 768; + const int tail_cells = (K - full_groups * 768 + (MMA_K - 1)) / MMA_K; + const int t_sf = (tail_cells + 1) / 2; // live SF slots in the tail + const int num_groups = full_groups + (tail_cells ? 1 : 0); + + // Persistent cluster index. The device always sees the grid in CTA units + // (2, 1, clusters) regardless of the host-side launch-unit regime. + const bool grid_z = (gridDim.z > 1); + const int num_clusters = grid_z ? int(gridDim.z) : int(gridDim.x) / 2; + const int cluster_id = grid_z ? int(blockIdx.z) : int(blockIdx.x) / 2; + + + // Which output tile does persistent work item `t` process? + auto tile_of = [&](int t) -> int { + return t; // natural raster: M fast, N slow + }; + + auto producer_tile_of = [&](int t) -> int { + return tile_of(t); + }; + + // ---- mbarrier init (warp 0, lane 0) ---- + if (warp_id == 0 && ptx::lane_id() == 0) { + // Arrival counts are derived from the topology, never tuned: + // "ready" = one producer expect_tx per slot (bit-24 routes both + // peers' bytes to the even CTA's barrier); "free" = one consumer + // commit per slot reuse. + #pragma unroll + for (int s = 0; s < CD_AB_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.ab_ready[s], 1); + ptx::mbar_init(&smem.hdr.ab_free[s], 1); + } + #pragma unroll + for (int s = 0; s < CD_SF_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.sf_ready[s], 1); + ptx::mbar_init(&smem.hdr.sf_free[s], 1); + } + ptx::mbar_init(&smem.hdr.acc_ready, 1); + ptx::mbar_init(&smem.hdr.acc_free, N_EPI_WARPS * CTA_GROUP); // 4 x 2 = 8 + ptx::mbar_init(&smem.hdr.dealloc, 32); // one warp's lanes, cross-peer + ptx::fence_mbarrier_init_release_cluster(); + } + ptx::cluster_sync_rel_acq(); + + // TMEM allocation: epilogue warp 0 allocates for the pair; warps 0-4 + // rendezvous before reading the granted base. The TMA producer warps + // never touch TMEM and proceed straight to their loops. + uint32_t taddr = 0; + if (warp_id <= 4) { + if (warp_id == 0) + ptx::tcgen05_alloc_2sm(ptx::to_shared(&smem.hdr.tmem_addr), + N_ALLOC_COLS); + ptx::bar_sync(2, 160); + taddr = smem.hdr.tmem_addr; + } + + const int wg = warp_id >> 2; // 0 = epilogue warps, 1 = producers + MMA + if (wg == 1) { + const int pw = warp_id - 4; + if (pw == 1 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 5: lane 0 owns the complete A/B TMA producer loop. + // ================================================================ + // free_phase starts at 1: a fresh mbarrier is phase 0, so the + // first ring's worth of "free" waits falls through. + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + // The pair splits the 256x256 tile: even CTA holds A rows + // [0,128) and B cols [0,128); odd CTA the other halves. + const int m_blk = m_row * CTA_GROUP + int(m_pair); + const int n_blk = n_grp * CTA_GROUP + int(m_pair); + #pragma unroll 1 + for (int g = 0; g < num_groups; ++g) { + // Three full 128-byte windows cover the packed K=768 group. + #pragma unroll + for (int sub = 0; sub < AB_TMA_SUBS; ++sub) { + ptx::mbar_wait_parity(&smem.hdr.ab_free[slot], free_phase); + const uint32_t a_s = + base + CD_A_OFF + uint32_t(slot) * CD_AB_STRIDE; + const uint32_t b_s = + base + CD_B_OFF + uint32_t(slot) * CD_AB_STRIDE; + uint64_t* bar = &smem.hdr.ab_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, AB_TX); + // A rides the multicast form with a self mask so the + // completion routes to the even peer's barrier. + const uint16_t mask_a = uint16_t(1u << m_pair); + // Read the same row-major FP4 allocation as cuBLAS. + // K192 uses four 96-byte boxes; K256 uses three + // 128-byte boxes. The tensor map zero-fills past K. + const CUtensorMap* a_map = &A_tmap; + const CUtensorMap* b_map = &B_tmap; + const int x = g * 384 + sub * AB_LOAD_BYTES; + const int y_a = m_blk * BLOCK_M; + const int y_b = n_blk * BLOCK_N_PEER; + const int z_a = 0; + const int z_b = 0; + ptx::cp_async_bulk_tensor_3d_load_multicast( + a_s, a_map, x, y_a, z_a, + bar, mask_a + ); + ptx::cp_async_bulk_tensor_3d_load_2sm_bit24( + b_s, b_map, x, y_b, z_b, + bar + ); + if (++slot == CD_AB_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + // Drain: wait for the last-filled slot to be consumed, then hand + // it one clean-shutdown expect_tx. + { + const int t_slot = (slot == 0) ? CD_AB_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.ab_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.ab_ready[t_slot], AB_TX); + } + } else if (pw == 2 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 6: scale-factor TMA producer. SFA is per-peer (each CTA's + // own M half); SFB is the one genuinely multicast operand — the + // cta_group::2 copy reads BOTH peers' smem as one logical + // operand, so both copies must be byte-identical (blog: the scale-factor section). + // ================================================================ + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const int m_blk = m_row * CTA_GROUP + int(m_pair); + for (int g = 0; g < num_groups; ++g) { + // 4 SF ring slots per 768-K group (2 K=96 steps each). + #pragma unroll + for (int sub = 0; sub < 4; ++sub) { + // The tail group stores only its live SF slots; dead + // slots are skipped entirely (no wait, no arrive, no + // cursor advance — the consumer skips symmetrically). + if (g == full_groups && sub >= t_sf) break; + ptx::mbar_wait_parity(&smem.hdr.sf_free[slot], free_phase); + const uint32_t sfa_s = + base + CD_SFA_OFF + uint32_t(slot) * CD_SFA_STRIDE; + const uint32_t sfb_s = + base + CD_SFB_OFF + uint32_t(slot) * CD_SFB_STRIDE; + uint64_t* bar = &smem.hdr.sf_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, SF_TX); + // VEC16 is [outer-block][K-group-of-4][128 rows x 4]. + // Three adjacent groups are the 128 x 12 scale slot. + const int sf_group = (g * 4 + sub) * 3; + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfa_s, &SFA_tmap, 0, 0, sf_group, m_blk, bar, + uint16_t(1u << m_pair) + ); + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfb_s + m_pair * uint32_t(CD_SFB_STRIDE / 2), + &SFB_tmap, 0, 0, sf_group, + n_grp * CTA_GROUP + int(m_pair), bar, + uint16_t(0x3) + ); + if (++slot == CD_SF_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + { + const int t_slot = (slot == 0) ? CD_SF_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.sf_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.sf_ready[t_slot], SF_TX); + } + } else if (pw == 0 && m_pair == 0 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 4: the MMA warp — issued once per cluster, on the even + // CTA. The odd CTA's warp 4 does nothing all kernel (blog: the warp table). + // The mainloop needs no tile coordinates: operands arrive through + // the rings; this warp only needs its share of the trip count. + // ================================================================ + const uint32_t smbase = ptx::to_shared(&smem); + const uint16_t commit_mask = 0x3; // both CTAs of this pair + MainloopState st{1u, 0u, 0u, 0u, 0u}; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + st = mainloop_tile(taddr, smbase, st, commit_mask, + uint32_t(full_groups), + uint32_t(tail_cells) + ); + // The outer role branch already leaves exactly one issuer. + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;" + :: "r"(smbase + uint32_t(MB_ACC_READY)), + "h"(uint16_t(0x3)) : "memory"); + } + // Drain the final tile's buffer-free arrive before exiting. + mbar_wait_relaxed(smbase + uint32_t(MB_ACC_FREE), st.d_parity); + } + } else { + // ==================================================================== + // Warps 0-3: the epilogue (blog: the epilogue). Each warp owns a 32-row band + // and drains a tile as eight 32-column tcgen05.ld loads, converts + // FP32->FP16 in packed pairs, and stores straight to global — no + // shared-memory bounce. + // ==================================================================== + const int epi_warp = warp_id; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + const uint32_t acc_free_rank = crank & ~1u; // the pair's even CTA + int d_db = 0; + int acc_ready_phase = 0; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const uint32_t tmem_base = taddr + uint32_t(d_db) * D_STRIDE; + mbar_wait_relaxed(ptx::to_shared(&smem.hdr.acc_ready), + uint32_t(acc_ready_phase)); + acc_ready_phase ^= 1; + + const int off_m = (m_row * CTA_GROUP + int(m_pair)) * BLOCK_M; + const int off_n = n_grp * BLOCK_N; + const int row = off_m + epi_warp * 32 + lane_id; + const bool row_in = (row < M); + constexpr int EPI_BANDS = BLOCK_N / EPI_CHUNK; // 8 + #pragma unroll 1 + for (int k = 0; k < EPI_BANDS; ++k) { + // Buffer 0 drains high-to-low and buffer 1 low-to-high, so + // their shared columns [220,256) drain first for both. + const int parity = d_db & 1; + const int band = (parity == 0) ? (EPI_BANDS - 1 - k) : k; + const uint32_t taddr_n = + tmem_base + uint32_t(band) * 32 + taddr_lane; + uint32_t w[32]; + ptx::tcgen05_ld_32x32b_x32(taddr_n, w); + // Return the old buffer as soon as its shared edge is out. + if (k == 1) { + ptx::tcgen05_wait_ld_then_arrive_release( + &smem.hdr.acc_free, acc_free_rank, + /*do_arrive=*/ptx::lane_id() == 0); + } + if (!row_in) continue; + const int col_base = off_n + band * 32; + uint32_t h[16]; + #pragma unroll + for (int j = 0; j < 16; ++j) { + h[j] = ptx::cvt_pack_f16x2(__int_as_float(w[2 * j + 1]), + __int_as_float(w[2 * j])); + } + const size_t row_base = size_t(row) * size_t(N); + const bool aligned32 = ((row_base & 15u) == 0u) && + ((uintptr_t(C_ptr) & 31u) == 0u); + if (aligned32 && (col_base & 15) == 0) { + // Precondition here: row_base is 16-element (32-byte) + // aligned — the aligned32 guard owns it. Restating it to + // ptxas via __builtin_assume compiled this file + // SASS-identical (and measured perf-null on the + // production kernel), so the guard stands alone. + #pragma unroll + for (int q = 0; q < 2; ++q) { + const int c0 = col_base + q * 16; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 16 <= N) { + uint4 lo{h[q*8+0], h[q*8+1], h[q*8+2], h[q*8+3]}; + uint4 hi{h[q*8+4], h[q*8+5], h[q*8+6], h[q*8+7]}; + ptx::st_global_na_ef_b256(cdst, lo, hi); + } else { + // N tail: 128-bit then scalar. + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int cq = c0 + p * 8; + if (cq + 8 <= N) { + uint4 v{h[q*8+p*4+0], h[q*8+p*4+1], + h[q*8+p*4+2], h[q*8+p*4+3]}; + ptx::st_global_na_b128( + reinterpret_cast(cdst + p * 8), v); + } else { + const uint16_t* hb = + reinterpret_cast( + &h[q * 8 + p * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (cq + c < N) cdst[p * 8 + c] = hb[c]; + } + } + } + } + } else { + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int c0 = col_base + q * 8; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 8 <= N) { + uint4 v{h[q*4+0], h[q*4+1], h[q*4+2], h[q*4+3]}; + ptx::st_global_na_b128( + reinterpret_cast(cdst), v); + } else { + const uint16_t* hb = + reinterpret_cast(&h[q * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (c0 + c < N) cdst[c] = hb[c]; + } + } + } + } + if (++d_db == D_DB) d_db = 0; + } + } + + // Teardown: only the epilogue warps rendezvous per CTA; warp 0 then runs + // the cross-peer dealloc handshake. Producer and MMA warps exit freely — + // their drains above already order their work. + if (wg == 0) { + ptx::bar_sync(3, 128); + if (warp_id == 0) { + ptx::tcgen05_relinquish_2sm(); + const uint32_t peer_rank = crank ^ 1u; + // All 32 lanes arrive the PEER's count-32 barrier = one flip. + ptx::mbar_arrive_cluster_release(&smem.hdr.dealloc, peer_rank); + ptx::mbar_wait_parity(&smem.hdr.dealloc, 0); + ptx::tcgen05_dealloc_2sm(taddr, N_ALLOC_COLS); + } + } +} + + +// Keep the host launch convention paired with the kernel's block attribute. +static inline dim3 launch_grid(int clusters) { + return GRID_IN_CLUSTERS ? dim3(1, 1, unsigned(clusters)) + : dim3(2, 1, unsigned(clusters)); +} + +} // namespace nvfp4 + +// ============================================================================ +// Host-side TMA descriptor support. The shared input and benchmark code lives +// in main.cu. +// ============================================================================ +namespace tmap { + +// Thin wrappers over cuTensorMapEncodeTiled. Dim 0 is the innermost +// (stride-1) axis; strides for dim 1+ are in bytes. +inline CUtensorMap encode_tiled_3d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_depth, uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t depth_stride_bytes, + uint32_t box_depth, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[3] = { global_cols, global_rows, global_depth }; + cuuint64_t global_strides[2] = { row_stride_bytes, depth_stride_bytes }; + cuuint32_t box_dim[3] = { box_cols, box_rows, box_depth }; + cuuint32_t element_strides[3] = { 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 3, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +inline CUtensorMap encode_tiled_4d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_slots, uint64_t global_tiles, + uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t tile_stride_bytes, + uint64_t slot_stride_bytes, + uint32_t box_tiles, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[4] = { global_cols, global_rows, global_tiles, + global_slots }; + cuuint64_t global_strides[3] = { row_stride_bytes, tile_stride_bytes, + slot_stride_bytes }; + cuuint32_t box_dim[4] = { box_cols, box_rows, box_tiles, 1u }; + cuuint32_t element_strides[4] = { 1, 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 4, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +} // namespace tmap + +namespace host { + +using namespace nvfp4; + +// ---- implementation-specific tensor maps --------------------------------- + +inline CUtensorMap make_ab_tmap(void* ptr, int rows, int K, + int row_stride) { + return tmap::encode_tiled_3d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*depth=*/1, uint64_t(rows), /*cols=*/uint64_t(K / 2), + /*row_stride=*/uint64_t(row_stride), + /*depth_stride=*/uint64_t(rows) * uint64_t(row_stride), + /*box=*/1, BLOCK_M, AB_LOAD_BYTES, CU_TENSOR_MAP_SWIZZLE_128B); +} + +inline CUtensorMap make_sf_tmap(uint8_t* ptr, int outer, int K) { + const uint64_t sf_inner = (uint64_t(K) + 63) / 64 * 4; + const uint64_t outer_blocks = (uint64_t(outer) + 127) / 128; + return tmap::encode_tiled_4d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*slots=*/outer_blocks, /*tiles=*/sf_inner / 4, + /*rows=*/4, /*cols=*/128, + /*row_stride=*/128, /*tile_stride=*/512, + /*slot_stride=*/sf_inner * 128, + /*box_tiles=*/3, /*box_rows=*/4, /*box_cols=*/128, + CU_TENSOR_MAP_SWIZZLE_NONE); +} +} // namespace host diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/gemm8.cuh b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm8.cuh new file mode 100644 index 000000000..144255976 --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm8.cuh @@ -0,0 +1,1688 @@ +#pragma once +// r8: dispatch compatible tails to K64. +// This is a complete, directly editable ladder snapshot. +// It intentionally duplicates the implementation in the other gemm*.cuh files. + +// Conditional branches are already materialized; edit this implementation directly. +#define NVFP4_BUILD_NAME "r8" +// Read-only metadata for main.cu; these do not select code in this header. +#define NVFP4_HAS_K64_TAIL 1 +#define NVFP4_HAS_SCHEDULE 0 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(expr) do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \ + cudaGetErrorName(_e), __FILE__, __LINE__, \ + cudaGetErrorString(_e)); \ + std::abort(); \ + } \ +} while (0) + +#define CU_CHECK(expr) do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + std::fprintf(stderr, "CU error at %s:%d: %s\n", \ + __FILE__, __LINE__, _s ? _s : "?"); \ + std::abort(); \ + } \ +} while (0) + + +// ============================================================================ +// PTX helpers. Thin zero-cost wrappers over the exact instruction forms the +// kernel needs — nothing generic, nothing speculative. +// ============================================================================ +namespace ptx { + +// Generic pointer -> 32-bit `.shared` window byte address. +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Local `.shared` address -> the same offset in CTA `cta_rank`'s smem +// (`.shared::cluster` instructions take this mapped form). +static __device__ __forceinline__ uint32_t mapa_shared_cluster( + uint32_t local_addr, uint32_t cta_rank) { + uint32_t mapped; + asm("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(mapped) : "r"(local_addr), "r"(cta_rank)); + return mapped; +} + +static __device__ __forceinline__ void bar_sync(uint32_t id, uint32_t count) { + asm volatile("barrier.cta.sync %0, %1;" :: "r"(id), "r"(count)); +} + +// Cluster barrier with release/acquire — the publish boundary between the +// mbarrier inits and any cross-CTA use of those mbarriers. +static __device__ __forceinline__ void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// %cluster_ctarank: 0 or 1 in this kernel's (2,1,1) clusters. +static __device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} + +static __device__ __forceinline__ uint32_t lane_id() { + uint32_t r; + asm("mov.u32 %0, %%laneid;" : "=r"(r)); + return r; +} + +// ---- mbarriers ------------------------------------------------------------- + +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + :: "r"(to_shared(bar)), "r"(count)); +} + +// Publish the inits to the async proxy + the peer CTA before any TMA +// completes on them. +static __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// Cross-CTA arrive with release ordering (teardown handshake). +static __device__ __forceinline__ void mbar_arrive_cluster_release( + uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" + :: "r"(mapped)); +} + +// Fused `tcgen05.wait::ld` + "buffer free" arrive, one asm block. The arrive +// must not become visible while TMEM loads are still in flight (it frees the +// accumulator buffer to the next tile's MMA — the bring-up chapter's +// drain overlap); the fused block +// keeps the arrive pinned behind the wait. (An earlier form also manufactured +// a register dependency on a drained value, `dep * 0` folded into the +// address; verified dead on nvcc 13.1 — ptxas const-folds it, SASS +// byte-identical without it — and retired.) +static __device__ __forceinline__ void tcgen05_wait_ld_then_arrive_release( + uint64_t* bar, uint32_t cta_rank, bool do_arrive) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile( + "{\n\t" + ".reg .pred %%pdo;\n\t" + "setp.ne.s32 %%pdo, %1, 0;\n\t" + "tcgen05.wait::ld.sync.aligned;\n\t" + "@%%pdo mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;\n\t" + "}\n\t" + :: "r"(mapped), "r"((int)do_arrive) : "memory"); +} + + +// Blocking parity wait at the consuming site (producer ring/tail waits + +// teardown). The 10 ms suspend hint is the polite form: a full-rate spin +// competes for issue slots and measurably loses; a nanosleep-backoff spin +// measured the same as this hint across five windows, so the simpler form +// ships. The carried one-stage-ahead try_wait peeks that used to gate these +// waits measured as pure overhead and were removed. +static __device__ __forceinline__ void mbar_wait_parity( + uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(to_shared(bar)), "r"(parity)); +} + +// ---- TMA ------------------------------------------------------------------- + +// Producer expect_tx arrive on the cluster-shared ring barrier: both +// peers' transfer bytes account against the even CTA's barrier, so one +// barrier per ring slot covers the whole cluster-wide gather. (Mechanism: +// `& 0xFEFFFFFF` clears bit 24 of the shared address — the hardware's +// peer-CTA bit — re-pointing the arrive at the even CTA's copy.) +static __device__ __forceinline__ void mbar_arrive_expect_tx_cluster( + uint64_t* bar, uint32_t bytes) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes)); +} + +// 3D tiled load, cta_group::2 + multicast (the A feed; the mask is the +// issuing CTA itself in a (2,1) cluster). Completion routes to the even +// CTA's mbarrier via the bit-24 clear. +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z) + : "memory"); +} + +// 3D tiled load, cta_group::2, no multicast (the B feed — each peer loads +// its own 128-wide N-block). +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_2sm_bit24( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst_smem), "l"(tmap), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) + : "memory"); +} + +// 4D tiled load, cta_group::2 + multicast (the scale-factor feed). The +// 512-byte SF tiles of one ring slot form a 4D box; `w` selects the slot. +static __device__ __forceinline__ void cp_async_bulk_tensor_4d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z2, int32_t w, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.4d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z2), "r"(w) + : "memory"); +} + + +// ---- tcgen05 (tensor memory + MMA plumbing) --------------------------------- + +static __device__ __forceinline__ void tcgen05_alloc_2sm( + uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_dealloc_2sm( + uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" + :: "r"(taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_relinquish_2sm() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); +} + +// The widest single-instruction TMEM read: 32 b32 registers per lane +// (blog: the epilogue — a quarter the load/wait count of the narrow forms). +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t taddr, uint32_t* r) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), + "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), + "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), + "=r"(r[16]), "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), + "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), + "=r"(r[24]), "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), + "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr)); +} + +// ---- stores + conversion ----------------------------------------------------- + +static __device__ __forceinline__ uint32_t cvt_pack_f16x2(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// L2::evict_first requires a 256-bit store, but the 128-bit tail can still +// adopt the no-L1-allocation part of the policy at this rung. +static __device__ __forceinline__ void st_global_na_b128( + uint4* ptr, uint4 v) { + asm volatile("st.global.L1::no_allocate.v4.b32 " + "[%0], {%1, %2, %3, %4};" + :: "l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)); +} + +// Add the output cache policy to the 256-bit store: bypass L1 and evict the +// write-once line first from L2 so it does not displace reusable operands. +static __device__ __forceinline__ void st_global_na_ef_b256( + void* ptr, uint4 lo, uint4 hi) { + asm volatile("st.global.L1::no_allocate.L2::evict_first.v8.b32 [%0], " + "{%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"(ptr), "r"(lo.x), "r"(lo.y), "r"(lo.z), "r"(lo.w), + "r"(hi.x), "r"(hi.y), "r"(hi.z), "r"(hi.w)); +} + + +// ---- MMA descriptors --------------------------------------------------------- + +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class ScaleFormat : uint8_t { E4M3 = 0, E8M0 = 1 }; + +// Source descriptor for the tcgen05.cp that stages scale factors from smem +// into TMEM: unswizzled, SBO=128, LBO=16. The low word later carries the +// per-slot byte offset (in 16-byte units). +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_sf_cp_src_desc_cd(uint32_t sf_smem_addr) { + uint64_t addr_field = uint64_t((sf_smem_addr >> 4) & 0x3FC0u); + return addr_field | 0x400800010000ull; +} + +// Operand descriptor for a K=96 (48-byte) slice inside a 128-byte-swizzled +// smem K-window. OR-constant 0x4010404000000000: version=1, absolute-address +// LBO mode (bit 52), 128B swizzle, SBO=1024. Both the address field and the +// absolute-LBO field encode the window's own 1024-aligned base; the +// intra-window offset advances the address field only. +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k96_sw128(uint32_t window_base_1024aligned, int intra_window_off = 0) { + uint64_t addr_field = uint64_t(((window_base_1024aligned >> 4) & 0x7FC0u) + + uint32_t(intra_window_off >> 4)); + uint64_t lbo_field = uint64_t((window_base_1024aligned << 12) & 0x7FC00000u); + return addr_field | lbo_field | 0x4010404000000000ull; +} + +// Instruction descriptor for tcgen05.mma.kind::mxf4nvf4.block_scale. +// The K extent is selected by BIT 31, not by the asm mnemonic: k_dim=0 runs +// K=64 (0x10400480), k_dim=1 runs K=96 (0x90400480, sm_103a-only) — the +// cheapest quarter of the gap if your kernel is built on the K=64 form +// (blog: the MMA warp). Bit 13 flips which TMEM scale-factor column the MMA reads. +__host__ __device__ static __forceinline__ constexpr uint32_t +mma_inst_desc_mxf4nvf4_block16( + uint32_t M, uint32_t N, + uint32_t tmem_sfa_addr, uint32_t tmem_sfb_addr, + Major a_major, Major b_major, + bool negate_a, bool negate_b, + bool k_dim, ScaleFormat sf) { + constexpr uint32_t MXF4_E2M1 = 1u; + uint32_t d = 0; + d |= ((tmem_sfb_addr & 0xC0000000u) >> 30) << 4; // b_sf_id + d |= MXF4_E2M1 << 7; // a_format = E2M1 + d |= MXF4_E2M1 << 10; // b_format = E2M1 + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; // N in units of 8 + d |= (static_cast(sf) & 0x1u) << 23; // 0 = UE4M3 + d |= ((M >> 4) & 0x1Fu) << 24; // m_dim at bits 27-28 + d |= ((tmem_sfa_addr & 0xC0000000u) >> 30) << 29; // a_sf_id + if (k_dim) d |= 1u << 31; // 1 = K=96 + return d; +} + +} // namespace ptx + +// ============================================================================ +// Kernel geometry from the blog's opening kernel section. One tile shape, +// one cluster shape: +// 256x256 output per 2-CTA cluster (128 rows per CTA), K advanced in +// 96-element MMA steps staged as 768-element K-groups of 8 steps, persistent +// grid of 76 clusters = 152 CTAs = one per SM. +// ============================================================================ +namespace nvfp4 { + +constexpr int CTA_GROUP = 2; // cta_group::2 — the MMA spans both CTAs +constexpr int BLOCK_M = 128; // A rows per CTA +constexpr int BLOCK_N = 256; // output tile N extent (per cluster) +constexpr int BLOCK_N_PEER = BLOCK_N / CTA_GROUP; // 128 — B cols per CTA +constexpr int CLUSTER_M = CTA_GROUP * BLOCK_M; // 256 — tile M extent +[[maybe_unused]] constexpr int MMA_K = 96; // K per tcgen05.mma (idesc bit 31) +constexpr int WIN_BYTES = 128; // one 128B-swizzled smem K-window +constexpr int AB_LOAD_BYTES = WIN_BYTES; +constexpr int AB_TMA_SUBS = 3; +constexpr int TB_SIZE = 224; // 7 warps (blog: the warp table) +constexpr int N_EPI_WARPS = 4; // warps 0-3 + +constexpr int SF_TMEM_COLS = 9 * 4; // 3 SFA + 6 SFB cp tiles, 4 cols each +// Two 256-wide accumulators use a 220-column stride and overlap by the 36 +// columns needed for scale factors. +constexpr int D_STRIDE = 220; +constexpr int D_DB = 2; +constexpr int N_D_COLS = D_STRIDE + BLOCK_N; +constexpr int SF_TMEM_BASE = 476; +constexpr int N_ALLOC_COLS = 512; +static_assert(N_D_COLS <= SF_TMEM_BASE && + SF_TMEM_BASE + SF_TMEM_COLS <= N_ALLOC_COLS, + "accumulators + scale factors exceed the 512 TMEM columns"); +constexpr int EPI_CHUNK = 32; // cols per regular drain + +// ---- the shared-memory map (blog: the feed) ----------------------------------- +// Every offset is a compile-time expression. +constexpr int CD_HDR = 1024; // mbarrier header pad +constexpr int CD_AB_SLOTS = 6; +constexpr int CD_SF_SLOTS = 7; +constexpr int CD_AB_STRIDE = BLOCK_M * WIN_BYTES; // 16384 per A or B slot +constexpr int CD_SF_TILE = 512; // one (32 lanes x 16B) x4 tile +constexpr int CD_SFA_TILES = 3; // per SFA slot (M = 1 block) +constexpr int CD_SFB_TILES = 6; // per SFB slot (N = 2 blocks) +constexpr int CD_SFA_STRIDE = CD_SFA_TILES * CD_SF_TILE; // 1536 +constexpr int CD_SFB_STRIDE = CD_SFB_TILES * CD_SF_TILE; // 3072 + +// Expected transfer bytes are derived, never tuned: a wrong-by-one +// expectation doesn't crash, it deadlocks (blog: the feed). +constexpr uint32_t AB_TX = uint32_t(CTA_GROUP) * + uint32_t(BLOCK_M + BLOCK_N_PEER) * uint32_t(AB_LOAD_BYTES); +constexpr uint32_t SF_TX = uint32_t(CD_SFA_STRIDE + CD_SFB_STRIDE) * + uint32_t(CTA_GROUP); // 9216 +static_assert(AB_TX == 65536u && SF_TX == 9216u, + "expect_tx derivation drifted"); + +// The map IS the struct: member sizes and alignment produce every byte +// offset; the raw-address constants the mainloop uses read off it with +// offsetof, and the static_asserts pin the published values. +struct SmemCD { + // mbarrier header: four ring regions in handshake order, then the + // singleton accumulator + teardown barriers (alignas pads it to 1024). + struct alignas(1024) Header { + uint64_t ab_ready[CD_AB_SLOTS]; // TMA -> MMA + uint64_t ab_free [CD_AB_SLOTS]; // MMA -> TMA + uint64_t sf_ready[CD_SF_SLOTS]; // TMA -> MMA + uint64_t sf_free [CD_SF_SLOTS]; // MMA -> TMA + uint64_t acc_ready; // MMA -> epi + uint64_t acc_free; // epi -> MMA + uint64_t dealloc; // 2-CTA teardown + uint32_t tmem_addr; // TMEM base addr storage + } hdr; + uint8_t a [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t b [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t sfa[CD_SF_SLOTS][CD_SFA_STRIDE]; + // SFB starts 1024-aligned (the cp source descriptor drops address bits 4-9). + alignas(1024) uint8_t sfb[CD_SF_SLOTS][CD_SFB_STRIDE]; +}; +static_assert(sizeof(SmemCD::Header) == CD_HDR, "mbarrier header overflows the pad"); + +constexpr int CD_A_OFF = offsetof(SmemCD, a); +constexpr int CD_B_OFF = offsetof(SmemCD, b); +constexpr int CD_SFA_OFF = offsetof(SmemCD, sfa); +constexpr int CD_SFB_OFF = offsetof(SmemCD, sfb); +constexpr int CD_TOTAL = sizeof(SmemCD); +static_assert(CD_TOTAL <= 232448, "smem map exceeds the opt-in budget"); +static_assert(CD_AB_SLOTS != 6 || + (CD_B_OFF == 99328 && CD_SFA_OFF == 197632 && + CD_SFB_OFF == 208896 && CD_TOTAL == 230400), + "the 6-slot byte map drifted from its published values"); + +// The mainloop computes raw `.shared` byte addresses from these. +constexpr int MB_AB_READY = offsetof(SmemCD, hdr.ab_ready); +constexpr int MB_AB_FREE = offsetof(SmemCD, hdr.ab_free); +constexpr int MB_SF_READY = offsetof(SmemCD, hdr.sf_ready); +constexpr int MB_SF_FREE = offsetof(SmemCD, hdr.sf_free); +constexpr int MB_ACC_READY = offsetof(SmemCD, hdr.acc_ready); +constexpr int MB_ACC_FREE = offsetof(SmemCD, hdr.acc_free); + + + +} // namespace nvfp4 + +// ============================================================================ +// The consumer mainloop from the blog's pipeline section. One warp — warp 4 +// of the even CTA — +// issues everything: the tcgen05.cp copies that stage scale factors from the +// smem ring into TMEM, immediately followed by the tcgen05.mma stream, in the +// SAME warp. A same-warp tcgen05 stream is one ordered async pipeline, so the +// write-after-read hazard on the single reused scale-factor TMEM region needs +// no fence and no cross-warp barrier. +// +// Per 768-element K-group, logical K=96 steps 0..7 issue in K order. Their +// starting offsets, expressed as 16-byte descriptor atoms modulo a 128-byte +// swizzle window, are 0,3,6,1,4,7,2,5. Atoms 6 and 7 straddle windows. +// ============================================================================ +namespace nvfp4 { + +// Integral tag for the K64-tail subclass dispatch (per-(a,b) folded variants). +template struct kdi { static constexpr uint32_t value = V; }; + +// Ring/parity state carried across tiles by the MMA warp. +struct MainloopState { + uint32_t d_parity; // which TMEM accumulator buffer the next tile writes + uint32_t ab_phase, ab_slot; // A/B ring cursor + uint32_t sf_phase, sf_slot; // scale-factor ring cursor +}; + + +// Relaxed blocking wait (no .acquire): the tcgen05 pipeline's own ordering +// covers these ring waits, and the .acquire form perturbs the scoreboard +// cadence on the issuing warp. Waits sit AT the consuming site — the carried +// one-stage-ahead try_wait peeks that used to gate them measured as pure +// overhead and were removed. +static __device__ __forceinline__ void mbar_wait_relaxed( + uint32_t bar_smem, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(bar_smem), "r"(parity) : "memory"); +} + +// Gated form for the @216 buffer-free wait, which fires once per tile: the +// gate branches over the spin for every group past the first. +static __device__ __forceinline__ void mbar_wait_gated_relaxed( + uint32_t bar_smem, uint32_t parity, bool peek) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p bra DONE_%=;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t" + "DONE_%=:\n\t}\n" + :: "r"(bar_smem), "r"(parity), "r"((uint32_t)peek) : "memory"); +} + + +// One scale-factor staging copy, smem -> TMEM. Single-lane semantics behind +// the converged elect predicate: tcgen05.cp is a single-thread instruction — +// issued bare from a converged warp you get 32 initiations and garbage. +static __device__ __forceinline__ void cp_sf( + uint32_t el, uint32_t dst, uint32_t lo, uint32_t hi) { + const uint64_t desc = (uint64_t(hi) << 32) | uint64_t(lo); + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;\n\t}\n" + :: "r"(dst), "l"(desc), "r"(el) : "memory"); +} + +// Ring-slot "consumed" arrive, multicast to both CTAs of the pair. +// arrive::one must fire exactly once — hence the elect predicate. +static __device__ __forceinline__ void commit_consumed( + uint32_t el, uint32_t smbase, uint32_t mb_off, uint16_t mask) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t}\n" + :: "r"(smbase + mb_off), "h"(mask), "r"(el) : "memory"); +} + +// One K=96 NVFP4 MMA. The asm token says scale_vec::4X regardless — K=96 is +// selected by idesc bit 31 (blog: the MMA warp). acc_pred = 0 overwrites the +// accumulator (the tile's first step), nonzero accumulates. +static __device__ __forceinline__ void mma_cell( + uint32_t el, uint32_t d, + uint32_t desc_a_lo, uint32_t desc_a_hi, + uint32_t desc_b_lo, uint32_t desc_b_hi, + uint32_t idesc, uint32_t sf_a, uint32_t sf_b, uint32_t acc_pred) { + const uint64_t desc_a = (uint64_t(desc_a_hi) << 32) | uint64_t(desc_a_lo); + const uint64_t desc_b = (uint64_t(desc_b_hi) << 32) | uint64_t(desc_b_lo); + asm volatile( + "{\n\t.reg .pred p, acc;\n\t" + "setp.ne.u32 p, %7, 0;\n\t" + "setp.ne.b32 acc, %4, 0;\n\t" + "@p tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X" + " [%0], %1, %2, %3, [%5], [%6], acc;\n\t}\n" + :: "r"(d), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(acc_pred), "r"(sf_a), "r"(sf_b), "r"(el) : "memory"); +} + +// One SF ring slot -> TMEM: 9 tcgen05.cp tiles (3 SFA + 6 SFB). Source +// address advances are in 16-byte descriptor units. The SFB tile order +// 0,3,1,4,2,5 interleaves the two 128-wide N-blocks. +static __device__ __forceinline__ void cp_sf_buffer( + uint32_t el, uint32_t slot, + uint32_t sfa0, uint32_t sfa1, uint32_t sfa2, + uint32_t sfb0, uint32_t sfb1, uint32_t sfb2, + uint32_t sfb3, uint32_t sfb4, uint32_t sfb5, + uint32_t sfa_hi, uint32_t sfa_lo, uint32_t sfb_hi, uint32_t sfb_lo) { + constexpr uint32_t SFA_SLOT_U = CD_SFA_STRIDE / 16u; // 96 + constexpr uint32_t SFB_SLOT_U = CD_SFB_STRIDE / 16u; // 192 + constexpr uint32_t TILE_U = CD_SF_TILE / 16u; // 32 + const uint32_t a = slot * SFA_SLOT_U + sfa_lo; + cp_sf(el, sfa0, a, sfa_hi); + cp_sf(el, sfa1, a + 1u*TILE_U, sfa_hi); + cp_sf(el, sfa2, a + 2u*TILE_U, sfa_hi); + const uint32_t b = slot * SFB_SLOT_U + sfb_lo; + cp_sf(el, sfb0, b + 0u*TILE_U, sfb_hi); + cp_sf(el, sfb1, b + 3u*TILE_U, sfb_hi); + cp_sf(el, sfb2, b + 1u*TILE_U, sfb_hi); + cp_sf(el, sfb3, b + 4u*TILE_U, sfb_hi); + cp_sf(el, sfb4, b + 2u*TILE_U, sfb_hi); + cp_sf(el, sfb5, b + 5u*TILE_U, sfb_hi); +} + +// Operand descriptor LOW word for A/B ring slot `slot`, K-step `k_atom` +// (0..7 inside the 768-element group; each step is 48 bytes = +3 in the +// 16-byte address field). +static __device__ __forceinline__ uint32_t ab_desc_lo( + uint32_t smbase, uint32_t ab_base_off, uint32_t slot, uint32_t k_atom) { + const uint64_t d = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + slot * CD_AB_STRIDE, int(k_atom) * 16); + return uint32_t(d & 0xFFFFFFFFu); +} + +// K-steps 6 and 7 cross a 128-byte window boundary. Their descriptor address +// comes from `addr_slot`, while the absolute-LBO field comes from the next +// slot whose swizzle window contains the second half. +static __device__ __forceinline__ uint32_t ab_desc_lo_straddle( + uint32_t smbase, uint32_t ab_base_off, + uint32_t addr_slot, uint32_t lbo_slot, uint32_t k_atom) { + const uint64_t d_addr = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + addr_slot * CD_AB_STRIDE, int(k_atom) * 16); + const uint64_t d_lbo = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + lbo_slot * CD_AB_STRIDE, 0); + return uint32_t(d_addr & 0x0000FFFFu) | uint32_t(d_lbo & 0xFFFF0000u); +} + +// One output tile: complete 768-element K-groups followed by one partial group +// containing `tail_cells` live K=96 steps. Dead cells fold their liveness into +// the issue predicate, so the warp stays converged. +static __device__ __forceinline__ MainloopState mainloop_tile( + uint32_t taddr, uint32_t smbase, MainloopState in, uint16_t commit_mask, + uint32_t full_groups, + uint32_t tail_cells, + uint32_t b64, uint32_t t_win + ) { + constexpr uint32_t AB_RING_SLOTS = uint32_t(CD_AB_SLOTS); + constexpr uint32_t SF_RING_SLOTS = uint32_t(CD_SF_SLOTS); + constexpr uint32_t AB_FREE_DELTA = uint32_t(MB_AB_FREE - MB_AB_READY); + constexpr uint32_t SF_READY_OFF = uint32_t(MB_SF_READY); + constexpr uint32_t SF_FREE_DELTA = uint32_t(MB_SF_FREE - MB_SF_READY); + constexpr uint32_t ACC_FREE_OFF = uint32_t(MB_ACC_FREE); + constexpr uint32_t A_OFF = uint32_t(CD_A_OFF); + constexpr uint32_t B_OFF = uint32_t(CD_B_OFF); + + uint32_t d_parity = in.d_parity; + uint32_t ab_phase = in.ab_phase, ab_slot = in.ab_slot; + uint32_t sf_phase = in.sf_phase, sf_slot = in.sf_slot; + + // The caller has already selected exactly one lane for the whole MMA role. + // Keep the helper predicate constant throughout the tile. + const uint32_t el = 1u; + + // Fixed TMEM targets for the single reused scale-factor region: 9 cp + // tiles of 4 columns each, starting past the accumulator region. + constexpr uint32_t SF = uint32_t(SF_TMEM_BASE); + const uint32_t sfa_t0 = taddr + SF + 0; // SFA word 0 (MMA reads here) + const uint32_t sfa_t1 = taddr + SF + 4; + const uint32_t sfa_t2 = taddr + SF + 8; + const uint32_t sfb_t0 = taddr + SF + 12; // SFB word 0 + const uint32_t sfb_t1 = taddr + SF + 16; + const uint32_t sfb_t2 = taddr + SF + 20; + const uint32_t sfb_t3 = taddr + SF + 24; + const uint32_t sfb_t4 = taddr + SF + 28; + const uint32_t sfb_t5 = taddr + SF + 32; + + // Scale-factor cp source descriptors (low word carries the per-slot + // offset; high word is constant). + const uint64_t sfb_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFB_OFF)); + const uint32_t sfb_src_lo = uint32_t(sfb_src & 0xFFFFFFFFu); + const uint32_t sfb_src_hi = uint32_t(sfb_src >> 32); + const uint64_t sfa_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFA_OFF)); + const uint32_t sfa_src_lo = uint32_t(sfa_src & 0xFFFFFFFFu); + const uint32_t sfa_src_hi = uint32_t(sfa_src >> 32); + + // Operand descriptor HIGH words (low words come from ab_desc_lo per step). + const uint32_t a_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + A_OFF) >> 32); + const uint32_t b_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + B_OFF) >> 32); + + // Each K=96 step needs 6 scale bytes per row, split across two staged + // 32-bit words. Odd-issue steps read word 1: bit 31 of the SF address + // selects the with-id encoding, and the id bits fold into the idesc. + constexpr uint32_t SF_ADDR_ID = 0x80000000u; + const uint32_t sfa_w1 = (taddr + SF + 4) | SF_ADDR_ID; + const uint32_t sfb_w1 = (taddr + SF + 20) | SF_ADDR_ID; + const uint32_t w1_id_bits = + ((sfa_w1 >> 1) & 0x60000000u) | ((sfb_w1 >> 26) & 48u); + + // The tile's first MMA overwrites the accumulator buffer this tile owns; + // the @216 "buffer free" wait runs once, before the first group. + const uint32_t acc_wait_parity = d_parity; + d_parity ^= 1u; + const uint32_t d_tmem = d_parity * uint32_t(D_STRIDE) + taddr; + + // Instruction descriptor: kind::mxf4nvf4, M=256, N=256, E2M1 x E2M1, + // UE4M3 block-16, K-major, k_dim=1 (K=96). Base value 0x90400480. + constexpr uint32_t IDESC = ptx::mma_inst_desc_mxf4nvf4_block16( + 256u, 256u, 0u, 0u, ptx::Major::K, ptx::Major::K, + false, false, /*k_dim=*/true, ptx::ScaleFormat::E4M3); + static_assert(IDESC == 0x90400480u, "K=96 NVFP4 idesc drifted"); + const uint32_t w0_id_bits = + ((sfa_t0 >> 1) & 0x60000000u) | ((sfb_t0 >> 26) & 48u); + const uint32_t idesc_w0 = w0_id_bits | IDESC; + const uint32_t idesc_w1 = w1_id_bits | IDESC; + + + uint32_t acc_pred = 0; // first MMA of the tile overwrites + bool need_acc_wait = true; // @216 wait fires once per tile + + if (full_groups >= 1) { + uint32_t g = 0; + do { + // ---- SF batch 0 -> logical steps 0,1 (atoms 0,3) ---- + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // ---- SF batch 1 -> logical steps 2,3 (atoms 6,1) ---- + mbar_wait_relaxed(sf_bar1, sf_phase1); + uint32_t sf_slot2, sf_phase2; + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + // AB0's last reader is MMA2. + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // ---- SF batch 2 -> logical steps 4,5 (atoms 4,7) ---- + mbar_wait_relaxed(sf_bar2, sf_phase2); + uint32_t sf_slot3, sf_phase3; + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; // ring cursor writeback + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // ---- SF batch 3 -> logical steps 6,7 (atoms 2,5) ---- + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; // ring cursor writeback + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + // AB1's last reader is MMA5; this measured placement leaves the + // commit behind the following scale copy. + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + acc_pred = 0xFFFFFFFFu; // accumulate from the second group on + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + need_acc_wait = false; + // AB2's last reader is MMA7. + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, + commit_mask); + ++g; + } while (g != full_groups); + } + + if (tail_cells == 0u) + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; + + // ---- the K64-dispatch tail: r = a x K96 + b x K64, a even, b in {1,2}. + // ONE fully-folded tail variant per active (a, b) subclass, selected by + // a once-per-tile branch ladder: every per-position encoding below is a + // compile-time constant inside its variant (the fixed-shape production + // receipt: full folding costs zero select/branch artifact, while the + // blended runtime-select form measured -0.75pp and a live uniform + // preamble evicts the hot mainloop's descriptor math off the uniform + // register file — +147 R2UR in the full-group loop, measured in SASS). + // Encodings are the metal-proven set: k_dim=0 idesc, LEGACY operand + // descriptors (bit 52 clear, even 16-byte units — abs-LBO pairs with + // k_dim=1 only), plain id-0 SF words ("block16 + K=64 mandates sf_id + // 0"). Cell count a + b and t_sf equal the K96 tail's, so the SF ring + // walk is IDENTICAL; only T_WIN = ceil((r/2)/128) A/B windows are + // consumed (the producer skips the same subs). + const auto k64_tail = [&](auto a_tag, auto b_tag) -> MainloopState { + constexpr uint32_t A_C = decltype(a_tag)::value; // even + constexpr uint32_t B64V = decltype(b_tag)::value; // 1 or 2 + constexpr uint32_t N_C = A_C + B64V; // == tail_cells + constexpr uint32_t T_WIN = + ((A_C * 96u + B64V * 64u) / 2u + 127u) / 128u; + static_assert((IDESC & 0x7FFFFFFFu) == 0x10400480u, + "K=64 NVFP4 idesc drifted"); + // The two K64 idesc forms share the runtime id bits already staged + // for the K96 skeleton; the K64#2 word-1-plain form derives its own. + const uint32_t idesc_k64a = idesc_w0 & 0x7FFFFFFFu; + const uint32_t idesc_k64b = ((sfa_t1 >> 1) & 0x60000000u) + | ((sfb_t2 >> 26) & 48u) + | (IDESC & 0x7FFFFFFFu); + const uint32_t a_hi64 = a_desc_hi & ~(1u << 20); // desc bit 52 + const uint32_t b_hi64 = b_desc_hi & ~(1u << 20); + // Ring slots/phases for the tail windows: pure cursor math. + const bool w0_wrap = (ab_slot + 1u == AB_RING_SLOTS); + const uint32_t s1 = w0_wrap ? 0u : ab_slot + 1u; + const uint32_t ph1 = ab_phase ^ (w0_wrap ? 1u : 0u); + const bool w1_wrap = (s1 + 1u == AB_RING_SLOTS); + const uint32_t s2 = w1_wrap ? 0u : s1 + 1u; + const uint32_t ph2 = ph1 ^ (w1_wrap ? 1u : 0u); + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + const uint32_t ab_bar1 = smbase + (s1 << 3); + const uint32_t ab_bar2 = smbase + (s2 << 3); + + // SF batch 0 (cells 0,1) — always live. + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar0, ab_phase); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + // Cell 0: K96 (A_C > 0) or the lone/leading K64 (A_C == 0). + if constexpr (A_C == 0u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 0) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, ab_slot, 0) & 0xFFFFu, b_hi64, + idesc_k64a, sfa_t0, sfb_t0, acc_pred); + } else { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); + } + // Cell 1: K96 / K64#2 (A_C == 0, unit 2) / dead. + if constexpr (N_C > 1u) { + if constexpr (A_C == 0u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 2) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, ab_slot, 2) & 0xFFFFu, b_hi64, + idesc_k64b, sfa_t1, sfb_t2, 0xFFFFFFFFu); + } else { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); + } + } + + // SF batch 1 (cells 2,3). + uint32_t sf_slot2 = sf_slot1, sf_phase2 = sf_phase1; + if constexpr (N_C > 2u) { + mbar_wait_relaxed(sf_bar1, sf_phase1); + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + } + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + if constexpr (T_WIN >= 2u) mbar_wait_relaxed(ab_bar1, ph1); + // Cell 2: K96 straddle (A_C > 2) / K64#1 at unit 6 of window 0 + // (A_C == 2; no crossing, legacy) / dead. + if constexpr (A_C == 2u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 6) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, ab_slot, 6) & 0xFFFFu, b_hi64, + idesc_k64a, sfa_t0, sfb_t0, 0xFFFFFFFFu); + } else if constexpr (A_C > 2u) { + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, s1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, s1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); + } + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, + commit_mask); + // Cell 3: K96 / K64#2 (A_C == 2, unit 0 of window 1) / dead. + if constexpr (A_C == 2u && B64V == 2u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s1, 0) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, s1, 0) & 0xFFFFu, b_hi64, + idesc_k64b, sfa_t1, sfb_t2, 0xFFFFFFFFu); + } else if constexpr (A_C > 3u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, s1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); + } + + // SF batch 2 (cells 4,5). + uint32_t sf_slot3 = sf_slot2, sf_phase3 = sf_phase2; + if constexpr (N_C > 4u) { + mbar_wait_relaxed(sf_bar2, sf_phase2); + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + } + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + // Cell 4: K96 / K64#1 at unit 4 of window 1 (A_C == 4) / dead. + if constexpr (A_C == 4u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s1, 4) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, s1, 4) & 0xFFFFu, b_hi64, + idesc_k64a, sfa_t0, sfb_t0, 0xFFFFFFFFu); + } else if constexpr (A_C > 4u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, s1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); + } + // A/B cursor writeback: entry + T_WIN (the producer advances the same). + if constexpr (T_WIN >= 3u) { + mbar_wait_relaxed(ab_bar2, ph2); + const bool w2_wrap = (s2 + 1u == AB_RING_SLOTS); + ab_slot = w2_wrap ? 0u : s2 + 1u; + ab_phase = ph2 ^ (w2_wrap ? 1u : 0u); + } else if constexpr (T_WIN == 2u) { + ab_slot = s2; + ab_phase = ph2; + } else { + ab_slot = s1; + ab_phase = ph1; + } + // Cell 5: K96 straddle (A_C == 6) / K64#2 at unit 6 of window 1 + // (A_C == 4, no crossing) / dead. + if constexpr (A_C == 4u && B64V == 2u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s1, 6) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, s1, 6) & 0xFFFFu, b_hi64, + idesc_k64b, sfa_t1, sfb_t2, 0xFFFFFFFFu); + } else if constexpr (A_C > 5u) { + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, s1, s2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, s1, s2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); + } + + // SF batch 3 (cells 6,7). + if constexpr (N_C > 6u) { + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + } else { + sf_slot = sf_slot3; + sf_phase = sf_phase3; + } + commit_consumed(T_WIN >= 2u ? el : 0u, smbase, + (ab_bar1 + AB_FREE_DELTA) - smbase, commit_mask); + // Cell 6: K96 / K64#1 at unit 2 of window 2 (A_C == 6) / dead. + if constexpr (A_C == 6u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s2, 2) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, s2, 2) & 0xFFFFu, b_hi64, + idesc_k64a, sfa_t0, sfb_t0, 0xFFFFFFFFu); + } + // Cell 7: K64#2 at unit 4 of window 2 (A_C == 6) / dead (no K96 + // cell 7 exists in the active class: A_C <= 6). + if constexpr (A_C == 6u && B64V == 2u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s2, 4) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, s2, 4) & 0xFFFFu, b_hi64, + idesc_k64b, sfa_t1, sfb_t2, 0xFFFFFFFFu); + } + commit_consumed(T_WIN >= 3u ? el : 0u, smbase, + (ab_bar2 + AB_FREE_DELTA) - smbase, commit_mask); + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; + }; + if (b64 != 0u) { + // Once-per-tile subclass dispatch (a = tail_cells - b, even). + const uint32_t a_c = tail_cells - b64; + if (b64 == 1u) { + if (a_c == 0u) return k64_tail(kdi<0>{}, kdi<1>{}); + if (a_c == 2u) return k64_tail(kdi<2>{}, kdi<1>{}); + if (a_c == 4u) return k64_tail(kdi<4>{}, kdi<1>{}); + return k64_tail(kdi<6>{}, kdi<1>{}); + } + if (a_c == 0u) return k64_tail(kdi<0>{}, kdi<2>{}); + if (a_c == 2u) return k64_tail(kdi<2>{}, kdi<2>{}); + if (a_c == 4u) return k64_tail(kdi<4>{}, kdi<2>{}); + return k64_tail(kdi<6>{}, kdi<2>{}); + } + + // ---- the exact-K tail group: tail_cells live K=96 steps out of 8 ---- + // Cell liveness rides the elect predicate (a 0-predicate op never enters + // the tcgen05 pipe); SF batch b exists iff tail_cells > 2b, and skipped + // batches advance NO ring cursor — the SF producer skips symmetrically. + const uint32_t el_c1 = (tail_cells >= 2u) ? el : 0u; + const uint32_t el_c2 = (tail_cells >= 3u) ? el : 0u; + const uint32_t el_c3 = (tail_cells >= 4u) ? el : 0u; + const uint32_t el_c4 = (tail_cells >= 5u) ? el : 0u; + const uint32_t el_c5 = (tail_cells >= 6u) ? el : 0u; + const uint32_t el_c6 = (tail_cells >= 7u) ? el : 0u; + const uint32_t el_c7 = (tail_cells >= 8u) ? el : 0u; + const bool b1_live = (tail_cells > 2u); + const bool b2_live = (tail_cells > 4u); + const bool b3_live = (tail_cells > 6u); + + // SF batch 0 (cells 0,1) — a live tail always has cell 0. + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el_c1, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // SF batch 1 (cells 2,3). + uint32_t sf_slot2 = sf_slot1, sf_phase2 = sf_phase1; + if (b1_live) { + mbar_wait_relaxed(sf_bar1, sf_phase1); + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + } + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el_c2, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, commit_mask); + mma_cell(el_c3, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // SF batch 2 (cells 4,5). + uint32_t sf_slot3 = sf_slot2, sf_phase3 = sf_phase2; + if (b2_live) { + mbar_wait_relaxed(sf_bar2, sf_phase2); + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + } + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el_c4, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el_c5, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // SF batch 3 (cells 6,7). + if (b3_live) { + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + } else { + sf_slot = sf_slot3; + sf_phase = sf_phase3; + } + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, commit_mask); + mma_cell(el_c6, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + mma_cell(el_c7, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, commit_mask); + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; +} + +} // namespace nvfp4 + +// ============================================================================ +// The kernel. 224 threads = 7 warps, four jobs (blog: the warp table): +// warps 0-3 epilogue: drain TMEM, convert to FP16, store straight to global +// warp 4 the MMA warp (even CTA of each pair only): SF copies + MMAs +// warp 5 TMA producer: A and B tiles into the 6-slot ring +// warp 6 TMA producer: scale factors into the 7-slot ring +// ============================================================================ +namespace nvfp4 { + +// Promise ptxas the exact block shape and express the launch grid in cluster +// units. The launch helper below owns the corresponding grid convention. +#define NVFP4_BLOCK_ATTR __block_size__((TB_SIZE, 1, 1)) +constexpr bool GRID_IN_CLUSTERS = true; + +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void nvfp4_gemm_kernel(const __grid_constant__ CUtensorMap A_tmap, + const __grid_constant__ CUtensorMap B_tmap, + const __grid_constant__ CUtensorMap SFA_tmap, + const __grid_constant__ CUtensorMap SFB_tmap, + __half* C_ptr, int M, int N, int K + ) { + extern __shared__ __align__(1024) char smem_buf[]; + SmemCD& smem = *reinterpret_cast(smem_buf); + + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const uint32_t crank = ptx::cluster_rank(); + const uint32_t m_pair = crank & 1u; // 0 = the pair's even CTA + + const int cluster_grid_m = (M + CLUSTER_M - 1) / CLUSTER_M; + const int grid_n = (N + BLOCK_N - 1) / BLOCK_N; + const int total_tiles = cluster_grid_m * grid_n; + + // Exact-K tail geometry (the exact-K feed, Other shapes; the host gates K % 16 == 0). + const int full_groups = K / 768; + const int tail_cells = (K - full_groups * 768 + (MMA_K - 1)) / MMA_K; + const int t_sf = (tail_cells + 1) / 2; // live SF slots in the tail + const int num_groups = full_groups + (tail_cells ? 1 : 0); + // K64-dispatch class (banner): r = a x 96 + b x 64, minimal b in {0,1,2}; + // active only for EVEN a (r mod 192 in {64,128}: the with-id SF word is + // architecturally sf_id=0-only at k_dim=0, and even a also excludes the + // window-crossing decomposition). Cell count a + b equals tail_cells and + // t_sf is unchanged, so only the A/B window count differs on the K64 path. + const int k_rem = K - full_groups * 768; + const int b64_try = (k_rem % 96 == 64) ? 1 : (k_rem % 96 == 32) ? 2 : 0; + const int a64 = tail_cells - b64_try; + const int b64 = (b64_try != 0 && k_rem % 32 == 0 && a64 >= 0 && + !(a64 & 1)) ? b64_try : 0; + // Live A/B tail windows on the K64 path (b == 0 keeps all 3). + const int t_win = b64 ? (k_rem / 2 + 127) / 128 : 3; + + // Persistent cluster index. The device always sees the grid in CTA units + // (2, 1, clusters) regardless of the host-side launch-unit regime. + const bool grid_z = (gridDim.z > 1); + const int num_clusters = grid_z ? int(gridDim.z) : int(gridDim.x) / 2; + const int cluster_id = grid_z ? int(blockIdx.z) : int(blockIdx.x) / 2; + + + // Which output tile does persistent work item `t` process? + auto tile_of = [&](int t) -> int { + return t; // natural raster: M fast, N slow + }; + + auto producer_tile_of = [&](int t) -> int { + return tile_of(t); + }; + + // ---- mbarrier init (warp 0, lane 0) ---- + if (warp_id == 0 && ptx::lane_id() == 0) { + // Arrival counts are derived from the topology, never tuned: + // "ready" = one producer expect_tx per slot (bit-24 routes both + // peers' bytes to the even CTA's barrier); "free" = one consumer + // commit per slot reuse. + #pragma unroll + for (int s = 0; s < CD_AB_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.ab_ready[s], 1); + ptx::mbar_init(&smem.hdr.ab_free[s], 1); + } + #pragma unroll + for (int s = 0; s < CD_SF_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.sf_ready[s], 1); + ptx::mbar_init(&smem.hdr.sf_free[s], 1); + } + ptx::mbar_init(&smem.hdr.acc_ready, 1); + ptx::mbar_init(&smem.hdr.acc_free, N_EPI_WARPS * CTA_GROUP); // 4 x 2 = 8 + ptx::mbar_init(&smem.hdr.dealloc, 32); // one warp's lanes, cross-peer + ptx::fence_mbarrier_init_release_cluster(); + } + ptx::cluster_sync_rel_acq(); + + // TMEM allocation: epilogue warp 0 allocates for the pair; warps 0-4 + // rendezvous before reading the granted base. The TMA producer warps + // never touch TMEM and proceed straight to their loops. + uint32_t taddr = 0; + if (warp_id <= 4) { + if (warp_id == 0) + ptx::tcgen05_alloc_2sm(ptx::to_shared(&smem.hdr.tmem_addr), + N_ALLOC_COLS); + ptx::bar_sync(2, 160); + taddr = smem.hdr.tmem_addr; + } + + const int wg = warp_id >> 2; // 0 = epilogue warps, 1 = producers + MMA + if (wg == 1) { + const int pw = warp_id - 4; + if (pw == 1 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 5: lane 0 owns the complete A/B TMA producer loop. + // ================================================================ + // free_phase starts at 1: a fresh mbarrier is phase 0, so the + // first ring's worth of "free" waits falls through. + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + // The pair splits the 256x256 tile: even CTA holds A rows + // [0,128) and B cols [0,128); odd CTA the other halves. + const int m_blk = m_row * CTA_GROUP + int(m_pair); + const int n_blk = n_grp * CTA_GROUP + int(m_pair); + #pragma unroll 1 + for (int g = 0; g < num_groups; ++g) { + // Three full 128-byte windows cover the packed K=768 group. + #pragma unroll + for (int sub = 0; sub < AB_TMA_SUBS; ++sub) { + // The K64-dispatch tail ends at logical K: windows + // past t_win are never read — skip them entirely (no + // wait, no expect_tx, no cursor advance; the consumer + // skips symmetrically). Inactive remainders keep + // t_win == 3, so this never fires for them. + if (g == full_groups && sub >= t_win) break; + ptx::mbar_wait_parity(&smem.hdr.ab_free[slot], free_phase); + const uint32_t a_s = + base + CD_A_OFF + uint32_t(slot) * CD_AB_STRIDE; + const uint32_t b_s = + base + CD_B_OFF + uint32_t(slot) * CD_AB_STRIDE; + uint64_t* bar = &smem.hdr.ab_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, AB_TX); + // A rides the multicast form with a self mask so the + // completion routes to the even peer's barrier. + const uint16_t mask_a = uint16_t(1u << m_pair); + // Read the same row-major FP4 allocation as cuBLAS. + // K192 uses four 96-byte boxes; K256 uses three + // 128-byte boxes. The tensor map zero-fills past K. + const CUtensorMap* a_map = &A_tmap; + const CUtensorMap* b_map = &B_tmap; + const int x = g * 384 + sub * AB_LOAD_BYTES; + const int y_a = m_blk * BLOCK_M; + const int y_b = n_blk * BLOCK_N_PEER; + const int z_a = 0; + const int z_b = 0; + ptx::cp_async_bulk_tensor_3d_load_multicast( + a_s, a_map, x, y_a, z_a, + bar, mask_a + ); + ptx::cp_async_bulk_tensor_3d_load_2sm_bit24( + b_s, b_map, x, y_b, z_b, + bar + ); + if (++slot == CD_AB_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + // Drain: wait for the last-filled slot to be consumed, then hand + // it one clean-shutdown expect_tx. + { + const int t_slot = (slot == 0) ? CD_AB_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.ab_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.ab_ready[t_slot], AB_TX); + } + } else if (pw == 2 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 6: scale-factor TMA producer. SFA is per-peer (each CTA's + // own M half); SFB is the one genuinely multicast operand — the + // cta_group::2 copy reads BOTH peers' smem as one logical + // operand, so both copies must be byte-identical (blog: the scale-factor section). + // ================================================================ + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const int m_blk = m_row * CTA_GROUP + int(m_pair); + for (int g = 0; g < num_groups; ++g) { + // 4 SF ring slots per 768-K group (2 K=96 steps each). + #pragma unroll + for (int sub = 0; sub < 4; ++sub) { + // The tail group stores only its live SF slots; dead + // slots are skipped entirely (no wait, no arrive, no + // cursor advance — the consumer skips symmetrically). + if (g == full_groups && sub >= t_sf) break; + ptx::mbar_wait_parity(&smem.hdr.sf_free[slot], free_phase); + const uint32_t sfa_s = + base + CD_SFA_OFF + uint32_t(slot) * CD_SFA_STRIDE; + const uint32_t sfb_s = + base + CD_SFB_OFF + uint32_t(slot) * CD_SFB_STRIDE; + uint64_t* bar = &smem.hdr.sf_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, SF_TX); + // VEC16 is [outer-block][K-group-of-4][128 rows x 4]. + // Three adjacent groups are the 128 x 12 scale slot. + const int sf_group = (g * 4 + sub) * 3; + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfa_s, &SFA_tmap, 0, 0, sf_group, m_blk, bar, + uint16_t(1u << m_pair) + ); + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfb_s + m_pair * uint32_t(CD_SFB_STRIDE / 2), + &SFB_tmap, 0, 0, sf_group, + n_grp * CTA_GROUP + int(m_pair), bar, + uint16_t(0x3) + ); + if (++slot == CD_SF_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + { + const int t_slot = (slot == 0) ? CD_SF_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.sf_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.sf_ready[t_slot], SF_TX); + } + } else if (pw == 0 && m_pair == 0 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 4: the MMA warp — issued once per cluster, on the even + // CTA. The odd CTA's warp 4 does nothing all kernel (blog: the warp table). + // The mainloop needs no tile coordinates: operands arrive through + // the rings; this warp only needs its share of the trip count. + // ================================================================ + const uint32_t smbase = ptx::to_shared(&smem); + const uint16_t commit_mask = 0x3; // both CTAs of this pair + MainloopState st{1u, 0u, 0u, 0u, 0u}; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + st = mainloop_tile(taddr, smbase, st, commit_mask, + uint32_t(full_groups), + uint32_t(tail_cells), + uint32_t(b64), uint32_t(t_win) + ); + // The outer role branch already leaves exactly one issuer. + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;" + :: "r"(smbase + uint32_t(MB_ACC_READY)), + "h"(uint16_t(0x3)) : "memory"); + } + // Drain the final tile's buffer-free arrive before exiting. + mbar_wait_relaxed(smbase + uint32_t(MB_ACC_FREE), st.d_parity); + } + } else { + // ==================================================================== + // Warps 0-3: the epilogue (blog: the epilogue). Each warp owns a 32-row band + // and drains a tile as eight 32-column tcgen05.ld loads, converts + // FP32->FP16 in packed pairs, and stores straight to global — no + // shared-memory bounce. + // ==================================================================== + const int epi_warp = warp_id; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + const uint32_t acc_free_rank = crank & ~1u; // the pair's even CTA + int d_db = 0; + int acc_ready_phase = 0; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const uint32_t tmem_base = taddr + uint32_t(d_db) * D_STRIDE; + mbar_wait_relaxed(ptx::to_shared(&smem.hdr.acc_ready), + uint32_t(acc_ready_phase)); + acc_ready_phase ^= 1; + + const int off_m = (m_row * CTA_GROUP + int(m_pair)) * BLOCK_M; + const int off_n = n_grp * BLOCK_N; + const int row = off_m + epi_warp * 32 + lane_id; + const bool row_in = (row < M); + constexpr int EPI_BANDS = BLOCK_N / EPI_CHUNK; // 8 + #pragma unroll 1 + for (int k = 0; k < EPI_BANDS; ++k) { + // Buffer 0 drains high-to-low and buffer 1 low-to-high, so + // their shared columns [220,256) drain first for both. + const int parity = d_db & 1; + const int band = (parity == 0) ? (EPI_BANDS - 1 - k) : k; + const uint32_t taddr_n = + tmem_base + uint32_t(band) * 32 + taddr_lane; + uint32_t w[32]; + ptx::tcgen05_ld_32x32b_x32(taddr_n, w); + // Return the old buffer as soon as its shared edge is out. + if (k == 1) { + ptx::tcgen05_wait_ld_then_arrive_release( + &smem.hdr.acc_free, acc_free_rank, + /*do_arrive=*/ptx::lane_id() == 0); + } + if (!row_in) continue; + const int col_base = off_n + band * 32; + uint32_t h[16]; + #pragma unroll + for (int j = 0; j < 16; ++j) { + h[j] = ptx::cvt_pack_f16x2(__int_as_float(w[2 * j + 1]), + __int_as_float(w[2 * j])); + } + const size_t row_base = size_t(row) * size_t(N); + const bool aligned32 = ((row_base & 15u) == 0u) && + ((uintptr_t(C_ptr) & 31u) == 0u); + if (aligned32 && (col_base & 15) == 0) { + // Precondition here: row_base is 16-element (32-byte) + // aligned — the aligned32 guard owns it. Restating it to + // ptxas via __builtin_assume compiled this file + // SASS-identical (and measured perf-null on the + // production kernel), so the guard stands alone. + #pragma unroll + for (int q = 0; q < 2; ++q) { + const int c0 = col_base + q * 16; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 16 <= N) { + uint4 lo{h[q*8+0], h[q*8+1], h[q*8+2], h[q*8+3]}; + uint4 hi{h[q*8+4], h[q*8+5], h[q*8+6], h[q*8+7]}; + ptx::st_global_na_ef_b256(cdst, lo, hi); + } else { + // N tail: 128-bit then scalar. + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int cq = c0 + p * 8; + if (cq + 8 <= N) { + uint4 v{h[q*8+p*4+0], h[q*8+p*4+1], + h[q*8+p*4+2], h[q*8+p*4+3]}; + ptx::st_global_na_b128( + reinterpret_cast(cdst + p * 8), v); + } else { + const uint16_t* hb = + reinterpret_cast( + &h[q * 8 + p * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (cq + c < N) cdst[p * 8 + c] = hb[c]; + } + } + } + } + } else { + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int c0 = col_base + q * 8; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 8 <= N) { + uint4 v{h[q*4+0], h[q*4+1], h[q*4+2], h[q*4+3]}; + ptx::st_global_na_b128( + reinterpret_cast(cdst), v); + } else { + const uint16_t* hb = + reinterpret_cast(&h[q * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (c0 + c < N) cdst[c] = hb[c]; + } + } + } + } + if (++d_db == D_DB) d_db = 0; + } + } + + // Teardown: only the epilogue warps rendezvous per CTA; warp 0 then runs + // the cross-peer dealloc handshake. Producer and MMA warps exit freely — + // their drains above already order their work. + if (wg == 0) { + ptx::bar_sync(3, 128); + if (warp_id == 0) { + ptx::tcgen05_relinquish_2sm(); + const uint32_t peer_rank = crank ^ 1u; + // All 32 lanes arrive the PEER's count-32 barrier = one flip. + ptx::mbar_arrive_cluster_release(&smem.hdr.dealloc, peer_rank); + ptx::mbar_wait_parity(&smem.hdr.dealloc, 0); + ptx::tcgen05_dealloc_2sm(taddr, N_ALLOC_COLS); + } + } +} + + +// Keep the host launch convention paired with the kernel's block attribute. +static inline dim3 launch_grid(int clusters) { + return GRID_IN_CLUSTERS ? dim3(1, 1, unsigned(clusters)) + : dim3(2, 1, unsigned(clusters)); +} + +} // namespace nvfp4 + +// ============================================================================ +// Host-side TMA descriptor support. The shared input and benchmark code lives +// in main.cu. +// ============================================================================ +namespace tmap { + +// Thin wrappers over cuTensorMapEncodeTiled. Dim 0 is the innermost +// (stride-1) axis; strides for dim 1+ are in bytes. +inline CUtensorMap encode_tiled_3d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_depth, uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t depth_stride_bytes, + uint32_t box_depth, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[3] = { global_cols, global_rows, global_depth }; + cuuint64_t global_strides[2] = { row_stride_bytes, depth_stride_bytes }; + cuuint32_t box_dim[3] = { box_cols, box_rows, box_depth }; + cuuint32_t element_strides[3] = { 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 3, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +inline CUtensorMap encode_tiled_4d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_slots, uint64_t global_tiles, + uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t tile_stride_bytes, + uint64_t slot_stride_bytes, + uint32_t box_tiles, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[4] = { global_cols, global_rows, global_tiles, + global_slots }; + cuuint64_t global_strides[3] = { row_stride_bytes, tile_stride_bytes, + slot_stride_bytes }; + cuuint32_t box_dim[4] = { box_cols, box_rows, box_tiles, 1u }; + cuuint32_t element_strides[4] = { 1, 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 4, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +} // namespace tmap + +namespace host { + +using namespace nvfp4; + +// ---- implementation-specific tensor maps --------------------------------- + +inline CUtensorMap make_ab_tmap(void* ptr, int rows, int K, + int row_stride) { + return tmap::encode_tiled_3d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*depth=*/1, uint64_t(rows), /*cols=*/uint64_t(K / 2), + /*row_stride=*/uint64_t(row_stride), + /*depth_stride=*/uint64_t(rows) * uint64_t(row_stride), + /*box=*/1, BLOCK_M, AB_LOAD_BYTES, CU_TENSOR_MAP_SWIZZLE_128B); +} + +inline CUtensorMap make_sf_tmap(uint8_t* ptr, int outer, int K) { + const uint64_t sf_inner = (uint64_t(K) + 63) / 64 * 4; + const uint64_t outer_blocks = (uint64_t(outer) + 127) / 128; + return tmap::encode_tiled_4d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*slots=*/outer_blocks, /*tiles=*/sf_inner / 4, + /*rows=*/4, /*cols=*/128, + /*row_stride=*/128, /*tile_stride=*/512, + /*slot_stride=*/sf_inner * 128, + /*box_tiles=*/3, /*box_rows=*/4, /*box_cols=*/128, + CU_TENSOR_MAP_SWIZZLE_NONE); +} +} // namespace host diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/gemm9.cuh b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm9.cuh new file mode 100644 index 000000000..87d5bf134 --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/gemm9.cuh @@ -0,0 +1,2770 @@ +#pragma once +// r9: assign A reuse by L2 side and select the visit order per shape. +// This is a complete, directly editable ladder snapshot. +// It intentionally duplicates the implementation in the other gemm*.cuh files. + +// Conditional branches are already materialized; edit this implementation directly. +#define NVFP4_BUILD_NAME "r9" +// Read-only metadata for main.cu; these do not select code in this header. +#define NVFP4_HAS_K64_TAIL 1 +#define NVFP4_HAS_SCHEDULE 1 + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUDA_CHECK(expr) do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + std::fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \ + cudaGetErrorName(_e), __FILE__, __LINE__, \ + cudaGetErrorString(_e)); \ + std::abort(); \ + } \ +} while (0) + +#define CU_CHECK(expr) do { \ + CUresult _e = (expr); \ + if (_e != CUDA_SUCCESS) { \ + const char* _s = nullptr; \ + cuGetErrorString(_e, &_s); \ + std::fprintf(stderr, "CU error at %s:%d: %s\n", \ + __FILE__, __LINE__, _s ? _s : "?"); \ + std::abort(); \ + } \ +} while (0) + + +// ============================================================================ +// PTX helpers. Thin zero-cost wrappers over the exact instruction forms the +// kernel needs — nothing generic, nothing speculative. +// ============================================================================ +namespace ptx { + +// Generic pointer -> 32-bit `.shared` window byte address. +template +static __device__ __forceinline__ uint32_t to_shared(T* ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +// Local `.shared` address -> the same offset in CTA `cta_rank`'s smem +// (`.shared::cluster` instructions take this mapped form). +static __device__ __forceinline__ uint32_t mapa_shared_cluster( + uint32_t local_addr, uint32_t cta_rank) { + uint32_t mapped; + asm("mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(mapped) : "r"(local_addr), "r"(cta_rank)); + return mapped; +} + +static __device__ __forceinline__ void bar_sync(uint32_t id, uint32_t count) { + asm volatile("barrier.cta.sync %0, %1;" :: "r"(id), "r"(count)); +} + +// Cluster barrier with release/acquire — the publish boundary between the +// mbarrier inits and any cross-CTA use of those mbarriers. +static __device__ __forceinline__ void cluster_sync_rel_acq() { + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); +} + +// %cluster_ctarank: 0 or 1 in this kernel's (2,1,1) clusters. +static __device__ __forceinline__ uint32_t cluster_rank() { + uint32_t r; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(r)); + return r; +} + +static __device__ __forceinline__ uint32_t lane_id() { + uint32_t r; + asm("mov.u32 %0, %%laneid;" : "=r"(r)); + return r; +} + +// ---- mbarriers ------------------------------------------------------------- + +static __device__ __forceinline__ void mbar_init(uint64_t* bar, uint32_t count) { + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + :: "r"(to_shared(bar)), "r"(count)); +} + +// Publish the inits to the async proxy + the peer CTA before any TMA +// completes on them. +static __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { + asm volatile("fence.mbarrier_init.release.cluster;"); +} + +// Cross-CTA arrive with release ordering (teardown handshake). +static __device__ __forceinline__ void mbar_arrive_cluster_release( + uint64_t* bar, uint32_t cta_rank) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile("mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;" + :: "r"(mapped)); +} + +// Fused `tcgen05.wait::ld` + "buffer free" arrive, one asm block. The arrive +// must not become visible while TMEM loads are still in flight (it frees the +// accumulator buffer to the next tile's MMA — the bring-up chapter's +// drain overlap); the fused block +// keeps the arrive pinned behind the wait. (An earlier form also manufactured +// a register dependency on a drained value, `dep * 0` folded into the +// address; verified dead on nvcc 13.1 — ptxas const-folds it, SASS +// byte-identical without it — and retired.) +static __device__ __forceinline__ void tcgen05_wait_ld_then_arrive_release( + uint64_t* bar, uint32_t cta_rank, bool do_arrive) { + const uint32_t mapped = mapa_shared_cluster(to_shared(bar), cta_rank); + asm volatile( + "{\n\t" + ".reg .pred %%pdo;\n\t" + "setp.ne.s32 %%pdo, %1, 0;\n\t" + "tcgen05.wait::ld.sync.aligned;\n\t" + "@%%pdo mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0], 1;\n\t" + "}\n\t" + :: "r"(mapped), "r"((int)do_arrive) : "memory"); +} + + +// Blocking parity wait at the consuming site (producer ring/tail waits + +// teardown). The 10 ms suspend hint is the polite form: a full-rate spin +// competes for issue slots and measurably loses; a nanosleep-backoff spin +// measured the same as this hint across five windows, so the simpler form +// ships. The carried one-stage-ahead try_wait peeks that used to gate these +// waits measured as pure overhead and were removed. +static __device__ __forceinline__ void mbar_wait_parity( + uint64_t* bar, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(to_shared(bar)), "r"(parity)); +} + +// ---- TMA ------------------------------------------------------------------- + +// Producer expect_tx arrive on the cluster-shared ring barrier: both +// peers' transfer bytes account against the even CTA's barrier, so one +// barrier per ring slot covers the whole cluster-wide gather. (Mechanism: +// `& 0xFEFFFFFF` clears bit 24 of the shared address — the hardware's +// peer-CTA bit — re-pointing the arrive at the even CTA's copy.) +static __device__ __forceinline__ void mbar_arrive_expect_tx_cluster( + uint64_t* bar, uint32_t bytes) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes)); +} + +// 3D tiled load, cta_group::2 + multicast (the A feed; the mask is the +// issuing CTA itself in a (2,1) cluster). Completion routes to the even +// CTA's mbarrier via the bit-24 clear. +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z) + : "memory"); +} + +// 3D tiled load, cta_group::2, no multicast (the B feed — each peer loads +// its own 128-wide N-block). +static __device__ __forceinline__ void cp_async_bulk_tensor_3d_load_2sm_bit24( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z, uint64_t* bar) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.3d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst_smem), "l"(tmap), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) + : "memory"); +} + +// 4D tiled load, cta_group::2 + multicast (the scale-factor feed). The +// 512-byte SF tiles of one ring slot form a 4D box; `w` selects the slot. +static __device__ __forceinline__ void cp_async_bulk_tensor_4d_load_multicast( + uint32_t dst_smem, const CUtensorMap* tmap, + int32_t x, int32_t y, int32_t z2, int32_t w, uint64_t* bar, + uint16_t multicast_mask) { + const uint32_t mbar_addr = to_shared(bar) & 0xFEFFFFFFu; + asm volatile( + "cp.async.bulk.tensor.4d.cta_group::2.shared::cluster.global" + ".mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + :: "r"(dst_smem), "l"(tmap), "r"(mbar_addr), "h"(multicast_mask), + "r"(x), "r"(y), "r"(z2), "r"(w) + : "memory"); +} + + +// ---- tcgen05 (tensor memory + MMA plumbing) --------------------------------- + +static __device__ __forceinline__ void tcgen05_alloc_2sm( + uint32_t smem_addr_for_taddr, uint32_t n_cols) { + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" + :: "r"(smem_addr_for_taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_dealloc_2sm( + uint32_t taddr, uint32_t n_cols) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" + :: "r"(taddr), "r"(n_cols)); +} + +static __device__ __forceinline__ void tcgen05_relinquish_2sm() { + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); +} + +// The widest single-instruction TMEM read: 32 b32 registers per lane +// (blog: the epilogue — a quarter the load/wait count of the narrow forms). +static __device__ __forceinline__ void tcgen05_ld_32x32b_x32( + uint32_t taddr, uint32_t* r) { + asm volatile("tcgen05.ld.sync.aligned.32x32b.x32.b32 " + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]), + "=r"(r[4]), "=r"(r[5]), "=r"(r[6]), "=r"(r[7]), + "=r"(r[8]), "=r"(r[9]), "=r"(r[10]), "=r"(r[11]), + "=r"(r[12]), "=r"(r[13]), "=r"(r[14]), "=r"(r[15]), + "=r"(r[16]), "=r"(r[17]), "=r"(r[18]), "=r"(r[19]), + "=r"(r[20]), "=r"(r[21]), "=r"(r[22]), "=r"(r[23]), + "=r"(r[24]), "=r"(r[25]), "=r"(r[26]), "=r"(r[27]), + "=r"(r[28]), "=r"(r[29]), "=r"(r[30]), "=r"(r[31]) + : "r"(taddr)); +} + +// ---- stores + conversion ----------------------------------------------------- + +static __device__ __forceinline__ uint32_t cvt_pack_f16x2(float a, float b) { + uint32_t d; + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;" : "=r"(d) : "f"(a), "f"(b)); + return d; +} + +// L2::evict_first requires a 256-bit store, but the 128-bit tail can still +// adopt the no-L1-allocation part of the policy at this rung. +static __device__ __forceinline__ void st_global_na_b128( + uint4* ptr, uint4 v) { + asm volatile("st.global.L1::no_allocate.v4.b32 " + "[%0], {%1, %2, %3, %4};" + :: "l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)); +} + +// Add the output cache policy to the 256-bit store: bypass L1 and evict the +// write-once line first from L2 so it does not displace reusable operands. +static __device__ __forceinline__ void st_global_na_ef_b256( + void* ptr, uint4 lo, uint4 hi) { + asm volatile("st.global.L1::no_allocate.L2::evict_first.v8.b32 [%0], " + "{%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"(ptr), "r"(lo.x), "r"(lo.y), "r"(lo.z), "r"(lo.w), + "r"(hi.x), "r"(hi.y), "r"(hi.z), "r"(hi.w)); +} + + +// ---- MMA descriptors --------------------------------------------------------- + +enum class Major : uint8_t { K = 0, MN = 1 }; +enum class ScaleFormat : uint8_t { E4M3 = 0, E8M0 = 1 }; + +// Source descriptor for the tcgen05.cp that stages scale factors from smem +// into TMEM: unswizzled, SBO=128, LBO=16. The low word later carries the +// per-slot byte offset (in 16-byte units). +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_sf_cp_src_desc_cd(uint32_t sf_smem_addr) { + uint64_t addr_field = uint64_t((sf_smem_addr >> 4) & 0x3FC0u); + return addr_field | 0x400800010000ull; +} + +// Operand descriptor for a K=96 (48-byte) slice inside a 128-byte-swizzled +// smem K-window. OR-constant 0x4010404000000000: version=1, absolute-address +// LBO mode (bit 52), 128B swizzle, SBO=1024. Both the address field and the +// absolute-LBO field encode the window's own 1024-aligned base; the +// intra-window offset advances the address field only. +__host__ __device__ static __forceinline__ constexpr uint64_t +mma_smem_desc_k96_sw128(uint32_t window_base_1024aligned, int intra_window_off = 0) { + uint64_t addr_field = uint64_t(((window_base_1024aligned >> 4) & 0x7FC0u) + + uint32_t(intra_window_off >> 4)); + uint64_t lbo_field = uint64_t((window_base_1024aligned << 12) & 0x7FC00000u); + return addr_field | lbo_field | 0x4010404000000000ull; +} + +// Instruction descriptor for tcgen05.mma.kind::mxf4nvf4.block_scale. +// The K extent is selected by BIT 31, not by the asm mnemonic: k_dim=0 runs +// K=64 (0x10400480), k_dim=1 runs K=96 (0x90400480, sm_103a-only) — the +// cheapest quarter of the gap if your kernel is built on the K=64 form +// (blog: the MMA warp). Bit 13 flips which TMEM scale-factor column the MMA reads. +__host__ __device__ static __forceinline__ constexpr uint32_t +mma_inst_desc_mxf4nvf4_block16( + uint32_t M, uint32_t N, + uint32_t tmem_sfa_addr, uint32_t tmem_sfb_addr, + Major a_major, Major b_major, + bool negate_a, bool negate_b, + bool k_dim, ScaleFormat sf) { + constexpr uint32_t MXF4_E2M1 = 1u; + uint32_t d = 0; + d |= ((tmem_sfb_addr & 0xC0000000u) >> 30) << 4; // b_sf_id + d |= MXF4_E2M1 << 7; // a_format = E2M1 + d |= MXF4_E2M1 << 10; // b_format = E2M1 + if (negate_a) d |= 1u << 13; + if (negate_b) d |= 1u << 14; + d |= (static_cast(a_major) & 0x1u) << 15; + d |= (static_cast(b_major) & 0x1u) << 16; + d |= ((N >> 3) & 0x3Fu) << 17; // N in units of 8 + d |= (static_cast(sf) & 0x1u) << 23; // 0 = UE4M3 + d |= ((M >> 4) & 0x1Fu) << 24; // m_dim at bits 27-28 + d |= ((tmem_sfa_addr & 0xC0000000u) >> 30) << 29; // a_sf_id + if (k_dim) d |= 1u << 31; // 1 = K=96 + return d; +} + +} // namespace ptx + +// ============================================================================ +// Kernel geometry from the blog's opening kernel section. One tile shape, +// one cluster shape: +// 256x256 output per 2-CTA cluster (128 rows per CTA), K advanced in +// 96-element MMA steps staged as 768-element K-groups of 8 steps, persistent +// grid of 76 clusters = 152 CTAs = one per SM. +// ============================================================================ +namespace nvfp4 { + +constexpr int CTA_GROUP = 2; // cta_group::2 — the MMA spans both CTAs +constexpr int BLOCK_M = 128; // A rows per CTA +constexpr int BLOCK_N = 256; // output tile N extent (per cluster) +constexpr int BLOCK_N_PEER = BLOCK_N / CTA_GROUP; // 128 — B cols per CTA +constexpr int CLUSTER_M = CTA_GROUP * BLOCK_M; // 256 — tile M extent +[[maybe_unused]] constexpr int MMA_K = 96; // K per tcgen05.mma (idesc bit 31) +constexpr int WIN_BYTES = 128; // one 128B-swizzled smem K-window +constexpr int AB_LOAD_BYTES = WIN_BYTES; +constexpr int AB_TMA_SUBS = 3; +constexpr int TB_SIZE = 224; // 7 warps (blog: the warp table) +constexpr int N_EPI_WARPS = 4; // warps 0-3 + +constexpr int SF_TMEM_COLS = 9 * 4; // 3 SFA + 6 SFB cp tiles, 4 cols each +// Two 256-wide accumulators use a 220-column stride and overlap by the 36 +// columns needed for scale factors. +constexpr int D_STRIDE = 220; +constexpr int D_DB = 2; +constexpr int N_D_COLS = D_STRIDE + BLOCK_N; +constexpr int SF_TMEM_BASE = 476; +constexpr int N_ALLOC_COLS = 512; +static_assert(N_D_COLS <= SF_TMEM_BASE && + SF_TMEM_BASE + SF_TMEM_COLS <= N_ALLOC_COLS, + "accumulators + scale factors exceed the 512 TMEM columns"); +constexpr int EPI_CHUNK = 32; // cols per regular drain + +// ---- the shared-memory map (blog: the feed) ----------------------------------- +// Every offset is a compile-time expression. +constexpr int CD_HDR = 1024; // mbarrier header pad +constexpr int CD_AB_SLOTS = 6; +constexpr int CD_SF_SLOTS = 7; +constexpr int CD_AB_STRIDE = BLOCK_M * WIN_BYTES; // 16384 per A or B slot +constexpr int CD_SF_TILE = 512; // one (32 lanes x 16B) x4 tile +constexpr int CD_SFA_TILES = 3; // per SFA slot (M = 1 block) +constexpr int CD_SFB_TILES = 6; // per SFB slot (N = 2 blocks) +constexpr int CD_SFA_STRIDE = CD_SFA_TILES * CD_SF_TILE; // 1536 +constexpr int CD_SFB_STRIDE = CD_SFB_TILES * CD_SF_TILE; // 3072 + +// Expected transfer bytes are derived, never tuned: a wrong-by-one +// expectation doesn't crash, it deadlocks (blog: the feed). +constexpr uint32_t AB_TX = uint32_t(CTA_GROUP) * + uint32_t(BLOCK_M + BLOCK_N_PEER) * uint32_t(AB_LOAD_BYTES); +constexpr uint32_t SF_TX = uint32_t(CD_SFA_STRIDE + CD_SFB_STRIDE) * + uint32_t(CTA_GROUP); // 9216 +static_assert(AB_TX == 65536u && SF_TX == 9216u, + "expect_tx derivation drifted"); + +// The map IS the struct: member sizes and alignment produce every byte +// offset; the raw-address constants the mainloop uses read off it with +// offsetof, and the static_asserts pin the published values. +struct SmemCD { + // mbarrier header: four ring regions in handshake order, then the + // singleton accumulator + teardown barriers (alignas pads it to 1024). + struct alignas(1024) Header { + uint64_t ab_ready[CD_AB_SLOTS]; // TMA -> MMA + uint64_t ab_free [CD_AB_SLOTS]; // MMA -> TMA + uint64_t sf_ready[CD_SF_SLOTS]; // TMA -> MMA + uint64_t sf_free [CD_SF_SLOTS]; // MMA -> TMA + uint64_t acc_ready; // MMA -> epi + uint64_t acc_free; // epi -> MMA + uint64_t dealloc; // 2-CTA teardown + uint32_t tmem_addr; // TMEM base addr storage + } hdr; + uint8_t a [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t b [CD_AB_SLOTS][CD_AB_STRIDE]; + uint8_t sfa[CD_SF_SLOTS][CD_SFA_STRIDE]; + // SFB starts 1024-aligned (the cp source descriptor drops address bits 4-9). + alignas(1024) uint8_t sfb[CD_SF_SLOTS][CD_SFB_STRIDE]; +}; +static_assert(sizeof(SmemCD::Header) == CD_HDR, "mbarrier header overflows the pad"); + +constexpr int CD_A_OFF = offsetof(SmemCD, a); +constexpr int CD_B_OFF = offsetof(SmemCD, b); +constexpr int CD_SFA_OFF = offsetof(SmemCD, sfa); +constexpr int CD_SFB_OFF = offsetof(SmemCD, sfb); +constexpr int CD_TOTAL = sizeof(SmemCD); +static_assert(CD_TOTAL <= 232448, "smem map exceeds the opt-in budget"); +static_assert(CD_AB_SLOTS != 6 || + (CD_B_OFF == 99328 && CD_SFA_OFF == 197632 && + CD_SFB_OFF == 208896 && CD_TOTAL == 230400), + "the 6-slot byte map drifted from its published values"); + +// The mainloop computes raw `.shared` byte addresses from these. +constexpr int MB_AB_READY = offsetof(SmemCD, hdr.ab_ready); +constexpr int MB_AB_FREE = offsetof(SmemCD, hdr.ab_free); +constexpr int MB_SF_READY = offsetof(SmemCD, hdr.sf_ready); +constexpr int MB_SF_FREE = offsetof(SmemCD, hdr.sf_free); +constexpr int MB_ACC_READY = offsetof(SmemCD, hdr.acc_ready); +constexpr int MB_ACC_FREE = offsetof(SmemCD, hdr.acc_free); + + +// ---- schedule table + placement audit (blog: the L2 trick) -------------------- +constexpr int L2A_SM_CAP = 256; +constexpr int L2A_CLUSTER_CAP = 128; +constexpr int L2A_ROUTE_MODES = 2; // slot 0: gate table, slot 1: schedule +constexpr int L2A_ROUTE_WORK_CAP = 64 * 64; +__device__ __constant__ int l2a_sm_side[L2A_SM_CAP]; +__device__ __constant__ int l2a_cluster_side[L2A_CLUSTER_CAP]; +__device__ __constant__ int l2a_route_tables[L2A_ROUTE_MODES * L2A_ROUTE_WORK_CAP]; +__device__ unsigned l2a_placement_errors; + +} // namespace nvfp4 + +// ============================================================================ +// The consumer mainloop from the blog's pipeline section. One warp — warp 4 +// of the even CTA — +// issues everything: the tcgen05.cp copies that stage scale factors from the +// smem ring into TMEM, immediately followed by the tcgen05.mma stream, in the +// SAME warp. A same-warp tcgen05 stream is one ordered async pipeline, so the +// write-after-read hazard on the single reused scale-factor TMEM region needs +// no fence and no cross-warp barrier. +// +// Per 768-element K-group, logical K=96 steps 0..7 issue in K order. Their +// starting offsets, expressed as 16-byte descriptor atoms modulo a 128-byte +// swizzle window, are 0,3,6,1,4,7,2,5. Atoms 6 and 7 straddle windows. +// ============================================================================ +namespace nvfp4 { + +// Integral tag for the K64-tail subclass dispatch (per-(a,b) folded variants). +template struct kdi { static constexpr uint32_t value = V; }; + +// Ring/parity state carried across tiles by the MMA warp. +struct MainloopState { + uint32_t d_parity; // which TMEM accumulator buffer the next tile writes + uint32_t ab_phase, ab_slot; // A/B ring cursor + uint32_t sf_phase, sf_slot; // scale-factor ring cursor +}; + + +// Relaxed blocking wait (no .acquire): the tcgen05 pipeline's own ordering +// covers these ring waits, and the .acquire form perturbs the scoreboard +// cadence on the issuing warp. Waits sit AT the consuming site — the carried +// one-stage-ahead try_wait peeks that used to gate them measured as pure +// overhead and were removed. +static __device__ __forceinline__ void mbar_wait_relaxed( + uint32_t bar_smem, uint32_t parity) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t}\n" + :: "r"(bar_smem), "r"(parity) : "memory"); +} + +// Gated form for the @216 buffer-free wait, which fires once per tile: the +// gate branches over the spin for every group past the first. +static __device__ __forceinline__ void mbar_wait_gated_relaxed( + uint32_t bar_smem, uint32_t parity, bool peek) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p bra DONE_%=;\n\t" + "WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1, 10000000;\n\t" + "@!p bra WAIT_%=;\n\t" + "DONE_%=:\n\t}\n" + :: "r"(bar_smem), "r"(parity), "r"((uint32_t)peek) : "memory"); +} + +// redux.sync.min over 32 equal lanes is an identity, but ptxas trusts its +// result as warp-uniform. Schedule bit-math then stays on the uniform register +// datapath instead of paying a broadcast per tcgen05 operation. +static __device__ __forceinline__ uint32_t to_uniform(uint32_t v) { + uint32_t r; + asm volatile("redux.sync.min.u32 %0, %1, 0xffffffff;\n\t" : "=r"(r) : "r"(v)); + return r; +} + +// One scale-factor staging copy, smem -> TMEM. Single-lane semantics behind +// the converged elect predicate: tcgen05.cp is a single-thread instruction — +// issued bare from a converged warp you get 32 initiations and garbage. +static __device__ __forceinline__ void cp_sf( + uint32_t el, uint32_t dst, uint32_t lo, uint32_t hi) { + const uint64_t desc = (uint64_t(hi) << 32) | uint64_t(lo); + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.cp.cta_group::2.32x128b.warpx4 [%0], %1;\n\t}\n" + :: "r"(dst), "l"(desc), "r"(el) : "memory"); +} + +// Ring-slot "consumed" arrive, multicast to both CTAs of the pair. +// arrive::one must fire exactly once — hence the elect predicate. +static __device__ __forceinline__ void commit_consumed( + uint32_t el, uint32_t smbase, uint32_t mb_off, uint16_t mask) { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "setp.ne.u32 p, %2, 0;\n\t" + "@p tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t}\n" + :: "r"(smbase + mb_off), "h"(mask), "r"(el) : "memory"); +} + +// One K=96 NVFP4 MMA. The asm token says scale_vec::4X regardless — K=96 is +// selected by idesc bit 31 (blog: the MMA warp). acc_pred = 0 overwrites the +// accumulator (the tile's first step), nonzero accumulates. +static __device__ __forceinline__ void mma_cell( + uint32_t el, uint32_t d, + uint32_t desc_a_lo, uint32_t desc_a_hi, + uint32_t desc_b_lo, uint32_t desc_b_hi, + uint32_t idesc, uint32_t sf_a, uint32_t sf_b, uint32_t acc_pred) { + const uint64_t desc_a = (uint64_t(desc_a_hi) << 32) | uint64_t(desc_a_lo); + const uint64_t desc_b = (uint64_t(desc_b_hi) << 32) | uint64_t(desc_b_lo); + asm volatile( + "{\n\t.reg .pred p, acc;\n\t" + "setp.ne.u32 p, %7, 0;\n\t" + "setp.ne.b32 acc, %4, 0;\n\t" + "@p tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X" + " [%0], %1, %2, %3, [%5], [%6], acc;\n\t}\n" + :: "r"(d), "l"(desc_a), "l"(desc_b), "r"(idesc), + "r"(acc_pred), "r"(sf_a), "r"(sf_b), "r"(el) : "memory"); +} + +// One SF ring slot -> TMEM: 9 tcgen05.cp tiles (3 SFA + 6 SFB). Source +// address advances are in 16-byte descriptor units. The SFB tile order +// 0,3,1,4,2,5 interleaves the two 128-wide N-blocks. +static __device__ __forceinline__ void cp_sf_buffer( + uint32_t el, uint32_t slot, + uint32_t sfa0, uint32_t sfa1, uint32_t sfa2, + uint32_t sfb0, uint32_t sfb1, uint32_t sfb2, + uint32_t sfb3, uint32_t sfb4, uint32_t sfb5, + uint32_t sfa_hi, uint32_t sfa_lo, uint32_t sfb_hi, uint32_t sfb_lo) { + constexpr uint32_t SFA_SLOT_U = CD_SFA_STRIDE / 16u; // 96 + constexpr uint32_t SFB_SLOT_U = CD_SFB_STRIDE / 16u; // 192 + constexpr uint32_t TILE_U = CD_SF_TILE / 16u; // 32 + const uint32_t a = slot * SFA_SLOT_U + sfa_lo; + cp_sf(el, sfa0, a, sfa_hi); + cp_sf(el, sfa1, a + 1u*TILE_U, sfa_hi); + cp_sf(el, sfa2, a + 2u*TILE_U, sfa_hi); + const uint32_t b = slot * SFB_SLOT_U + sfb_lo; + cp_sf(el, sfb0, b + 0u*TILE_U, sfb_hi); + cp_sf(el, sfb1, b + 3u*TILE_U, sfb_hi); + cp_sf(el, sfb2, b + 1u*TILE_U, sfb_hi); + cp_sf(el, sfb3, b + 4u*TILE_U, sfb_hi); + cp_sf(el, sfb4, b + 2u*TILE_U, sfb_hi); + cp_sf(el, sfb5, b + 5u*TILE_U, sfb_hi); +} + +// Operand descriptor LOW word for A/B ring slot `slot`, K-step `k_atom` +// (0..7 inside the 768-element group; each step is 48 bytes = +3 in the +// 16-byte address field). +static __device__ __forceinline__ uint32_t ab_desc_lo( + uint32_t smbase, uint32_t ab_base_off, uint32_t slot, uint32_t k_atom) { + const uint64_t d = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + slot * CD_AB_STRIDE, int(k_atom) * 16); + return uint32_t(d & 0xFFFFFFFFu); +} + +// K-steps 6 and 7 cross a 128-byte window boundary. Their descriptor address +// comes from `addr_slot`, while the absolute-LBO field comes from the next +// slot whose swizzle window contains the second half. +static __device__ __forceinline__ uint32_t ab_desc_lo_straddle( + uint32_t smbase, uint32_t ab_base_off, + uint32_t addr_slot, uint32_t lbo_slot, uint32_t k_atom) { + const uint64_t d_addr = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + addr_slot * CD_AB_STRIDE, int(k_atom) * 16); + const uint64_t d_lbo = ptx::mma_smem_desc_k96_sw128( + smbase + ab_base_off + lbo_slot * CD_AB_STRIDE, 0); + return uint32_t(d_addr & 0x0000FFFFu) | uint32_t(d_lbo & 0xFFFF0000u); +} + +// One output tile: complete 768-element K-groups followed by one partial group +// containing `tail_cells` live K=96 steps. Dead cells fold their liveness into +// the issue predicate, so the warp stays converged. +static __device__ __forceinline__ MainloopState mainloop_tile( + uint32_t taddr, uint32_t smbase, MainloopState in, uint16_t commit_mask, + uint32_t full_groups, + uint32_t tail_cells, + uint32_t b64, uint32_t t_win + ) { + constexpr uint32_t AB_RING_SLOTS = uint32_t(CD_AB_SLOTS); + constexpr uint32_t SF_RING_SLOTS = uint32_t(CD_SF_SLOTS); + constexpr uint32_t AB_FREE_DELTA = uint32_t(MB_AB_FREE - MB_AB_READY); + constexpr uint32_t SF_READY_OFF = uint32_t(MB_SF_READY); + constexpr uint32_t SF_FREE_DELTA = uint32_t(MB_SF_FREE - MB_SF_READY); + constexpr uint32_t ACC_FREE_OFF = uint32_t(MB_ACC_FREE); + constexpr uint32_t A_OFF = uint32_t(CD_A_OFF); + constexpr uint32_t B_OFF = uint32_t(CD_B_OFF); + + uint32_t d_parity = in.d_parity; + uint32_t ab_phase = in.ab_phase, ab_slot = in.ab_slot; + uint32_t sf_phase = in.sf_phase, sf_slot = in.sf_slot; + + // The caller has already selected exactly one lane for the whole MMA role. + // Keep the helper predicate constant throughout the tile. + const uint32_t el = 1u; + + // Fixed TMEM targets for the single reused scale-factor region: 9 cp + // tiles of 4 columns each, starting past the accumulator region. + constexpr uint32_t SF = uint32_t(SF_TMEM_BASE); + const uint32_t sfa_t0 = taddr + SF + 0; // SFA word 0 (MMA reads here) + const uint32_t sfa_t1 = taddr + SF + 4; + const uint32_t sfa_t2 = taddr + SF + 8; + const uint32_t sfb_t0 = taddr + SF + 12; // SFB word 0 + const uint32_t sfb_t1 = taddr + SF + 16; + const uint32_t sfb_t2 = taddr + SF + 20; + const uint32_t sfb_t3 = taddr + SF + 24; + const uint32_t sfb_t4 = taddr + SF + 28; + const uint32_t sfb_t5 = taddr + SF + 32; + + // Scale-factor cp source descriptors (low word carries the per-slot + // offset; high word is constant). + const uint64_t sfb_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFB_OFF)); + const uint32_t sfb_src_lo = uint32_t(sfb_src & 0xFFFFFFFFu); + const uint32_t sfb_src_hi = uint32_t(sfb_src >> 32); + const uint64_t sfa_src = ptx::mma_sf_cp_src_desc_cd(smbase + uint32_t(CD_SFA_OFF)); + const uint32_t sfa_src_lo = uint32_t(sfa_src & 0xFFFFFFFFu); + const uint32_t sfa_src_hi = uint32_t(sfa_src >> 32); + + // Operand descriptor HIGH words (low words come from ab_desc_lo per step). + const uint32_t a_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + A_OFF) >> 32); + const uint32_t b_desc_hi = + uint32_t(ptx::mma_smem_desc_k96_sw128(smbase + B_OFF) >> 32); + + // Each K=96 step needs 6 scale bytes per row, split across two staged + // 32-bit words. Odd-issue steps read word 1: bit 31 of the SF address + // selects the with-id encoding, and the id bits fold into the idesc. + constexpr uint32_t SF_ADDR_ID = 0x80000000u; + const uint32_t sfa_w1 = (taddr + SF + 4) | SF_ADDR_ID; + const uint32_t sfb_w1 = (taddr + SF + 20) | SF_ADDR_ID; + const uint32_t w1_id_bits = + ((sfa_w1 >> 1) & 0x60000000u) | ((sfb_w1 >> 26) & 48u); + + // The tile's first MMA overwrites the accumulator buffer this tile owns; + // the @216 "buffer free" wait runs once, before the first group. + const uint32_t acc_wait_parity = d_parity; + d_parity ^= 1u; + const uint32_t d_tmem = d_parity * uint32_t(D_STRIDE) + taddr; + + // Instruction descriptor: kind::mxf4nvf4, M=256, N=256, E2M1 x E2M1, + // UE4M3 block-16, K-major, k_dim=1 (K=96). Base value 0x90400480. + constexpr uint32_t IDESC = ptx::mma_inst_desc_mxf4nvf4_block16( + 256u, 256u, 0u, 0u, ptx::Major::K, ptx::Major::K, + false, false, /*k_dim=*/true, ptx::ScaleFormat::E4M3); + static_assert(IDESC == 0x90400480u, "K=96 NVFP4 idesc drifted"); + const uint32_t w0_id_bits = + ((sfa_t0 >> 1) & 0x60000000u) | ((sfb_t0 >> 26) & 48u); + const uint32_t idesc_w0 = w0_id_bits | IDESC; + const uint32_t idesc_w1 = w1_id_bits | IDESC; + + + uint32_t acc_pred = 0; // first MMA of the tile overwrites + bool need_acc_wait = true; // @216 wait fires once per tile + + if (full_groups >= 1) { + uint32_t g = 0; + do { + // ---- SF batch 0 -> logical steps 0,1 (atoms 0,3) ---- + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // ---- SF batch 1 -> logical steps 2,3 (atoms 6,1) ---- + mbar_wait_relaxed(sf_bar1, sf_phase1); + uint32_t sf_slot2, sf_phase2; + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + // AB0's last reader is MMA2. + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // ---- SF batch 2 -> logical steps 4,5 (atoms 4,7) ---- + mbar_wait_relaxed(sf_bar2, sf_phase2); + uint32_t sf_slot3, sf_phase3; + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; // ring cursor writeback + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // ---- SF batch 3 -> logical steps 6,7 (atoms 2,5) ---- + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; // ring cursor writeback + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + // AB1's last reader is MMA5; this measured placement leaves the + // commit behind the following scale copy. + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, + commit_mask); + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + acc_pred = 0xFFFFFFFFu; // accumulate from the second group on + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + need_acc_wait = false; + // AB2's last reader is MMA7. + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, + commit_mask); + ++g; + } while (g != full_groups); + } + + if (tail_cells == 0u) + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; + + // ---- the K64-dispatch tail: r = a x K96 + b x K64, a even, b in {1,2}. + // ONE fully-folded tail variant per active (a, b) subclass, selected by + // a once-per-tile branch ladder: every per-position encoding below is a + // compile-time constant inside its variant (the fixed-shape production + // receipt: full folding costs zero select/branch artifact, while the + // blended runtime-select form measured -0.75pp and a live uniform + // preamble evicts the hot mainloop's descriptor math off the uniform + // register file — +147 R2UR in the full-group loop, measured in SASS). + // Encodings are the metal-proven set: k_dim=0 idesc, LEGACY operand + // descriptors (bit 52 clear, even 16-byte units — abs-LBO pairs with + // k_dim=1 only), plain id-0 SF words ("block16 + K=64 mandates sf_id + // 0"). Cell count a + b and t_sf equal the K96 tail's, so the SF ring + // walk is IDENTICAL; only T_WIN = ceil((r/2)/128) A/B windows are + // consumed (the producer skips the same subs). + const auto k64_tail = [&](auto a_tag, auto b_tag) -> MainloopState { + constexpr uint32_t A_C = decltype(a_tag)::value; // even + constexpr uint32_t B64V = decltype(b_tag)::value; // 1 or 2 + constexpr uint32_t N_C = A_C + B64V; // == tail_cells + constexpr uint32_t T_WIN = + ((A_C * 96u + B64V * 64u) / 2u + 127u) / 128u; + static_assert((IDESC & 0x7FFFFFFFu) == 0x10400480u, + "K=64 NVFP4 idesc drifted"); + // The two K64 idesc forms share the runtime id bits already staged + // for the K96 skeleton; the K64#2 word-1-plain form derives its own. + const uint32_t idesc_k64a = idesc_w0 & 0x7FFFFFFFu; + const uint32_t idesc_k64b = ((sfa_t1 >> 1) & 0x60000000u) + | ((sfb_t2 >> 26) & 48u) + | (IDESC & 0x7FFFFFFFu); + const uint32_t a_hi64 = a_desc_hi & ~(1u << 20); // desc bit 52 + const uint32_t b_hi64 = b_desc_hi & ~(1u << 20); + // Ring slots/phases for the tail windows: pure cursor math. + const bool w0_wrap = (ab_slot + 1u == AB_RING_SLOTS); + const uint32_t s1 = w0_wrap ? 0u : ab_slot + 1u; + const uint32_t ph1 = ab_phase ^ (w0_wrap ? 1u : 0u); + const bool w1_wrap = (s1 + 1u == AB_RING_SLOTS); + const uint32_t s2 = w1_wrap ? 0u : s1 + 1u; + const uint32_t ph2 = ph1 ^ (w1_wrap ? 1u : 0u); + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + const uint32_t ab_bar1 = smbase + (s1 << 3); + const uint32_t ab_bar2 = smbase + (s2 << 3); + + // SF batch 0 (cells 0,1) — always live. + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, + commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar0, ab_phase); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + // Cell 0: K96 (A_C > 0) or the lone/leading K64 (A_C == 0). + if constexpr (A_C == 0u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 0) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, ab_slot, 0) & 0xFFFFu, b_hi64, + idesc_k64a, sfa_t0, sfb_t0, acc_pred); + } else { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); + } + // Cell 1: K96 / K64#2 (A_C == 0, unit 2) / dead. + if constexpr (N_C > 1u) { + if constexpr (A_C == 0u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 2) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, ab_slot, 2) & 0xFFFFu, b_hi64, + idesc_k64b, sfa_t1, sfb_t2, 0xFFFFFFFFu); + } else { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); + } + } + + // SF batch 1 (cells 2,3). + uint32_t sf_slot2 = sf_slot1, sf_phase2 = sf_phase1; + if constexpr (N_C > 2u) { + mbar_wait_relaxed(sf_bar1, sf_phase1); + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + } + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + if constexpr (T_WIN >= 2u) mbar_wait_relaxed(ab_bar1, ph1); + // Cell 2: K96 straddle (A_C > 2) / K64#1 at unit 6 of window 0 + // (A_C == 2; no crossing, legacy) / dead. + if constexpr (A_C == 2u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, ab_slot, 6) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, ab_slot, 6) & 0xFFFFu, b_hi64, + idesc_k64a, sfa_t0, sfb_t0, 0xFFFFFFFFu); + } else if constexpr (A_C > 2u) { + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, s1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, s1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); + } + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, + commit_mask); + // Cell 3: K96 / K64#2 (A_C == 2, unit 0 of window 1) / dead. + if constexpr (A_C == 2u && B64V == 2u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s1, 0) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, s1, 0) & 0xFFFFu, b_hi64, + idesc_k64b, sfa_t1, sfb_t2, 0xFFFFFFFFu); + } else if constexpr (A_C > 3u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, s1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); + } + + // SF batch 2 (cells 4,5). + uint32_t sf_slot3 = sf_slot2, sf_phase3 = sf_phase2; + if constexpr (N_C > 4u) { + mbar_wait_relaxed(sf_bar2, sf_phase2); + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + } + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + // Cell 4: K96 / K64#1 at unit 4 of window 1 (A_C == 4) / dead. + if constexpr (A_C == 4u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s1, 4) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, s1, 4) & 0xFFFFu, b_hi64, + idesc_k64a, sfa_t0, sfb_t0, 0xFFFFFFFFu); + } else if constexpr (A_C > 4u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, s1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); + } + // A/B cursor writeback: entry + T_WIN (the producer advances the same). + if constexpr (T_WIN >= 3u) { + mbar_wait_relaxed(ab_bar2, ph2); + const bool w2_wrap = (s2 + 1u == AB_RING_SLOTS); + ab_slot = w2_wrap ? 0u : s2 + 1u; + ab_phase = ph2 ^ (w2_wrap ? 1u : 0u); + } else if constexpr (T_WIN == 2u) { + ab_slot = s2; + ab_phase = ph2; + } else { + ab_slot = s1; + ab_phase = ph1; + } + // Cell 5: K96 straddle (A_C == 6) / K64#2 at unit 6 of window 1 + // (A_C == 4, no crossing) / dead. + if constexpr (A_C == 4u && B64V == 2u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s1, 6) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, s1, 6) & 0xFFFFu, b_hi64, + idesc_k64b, sfa_t1, sfb_t2, 0xFFFFFFFFu); + } else if constexpr (A_C > 5u) { + mma_cell(el, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, s1, s2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, s1, s2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); + } + + // SF batch 3 (cells 6,7). + if constexpr (N_C > 6u) { + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + } else { + sf_slot = sf_slot3; + sf_phase = sf_phase3; + } + commit_consumed(T_WIN >= 2u ? el : 0u, smbase, + (ab_bar1 + AB_FREE_DELTA) - smbase, commit_mask); + // Cell 6: K96 / K64#1 at unit 2 of window 2 (A_C == 6) / dead. + if constexpr (A_C == 6u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s2, 2) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, s2, 2) & 0xFFFFu, b_hi64, + idesc_k64a, sfa_t0, sfb_t0, 0xFFFFFFFFu); + } + // Cell 7: K64#2 at unit 4 of window 2 (A_C == 6) / dead (no K96 + // cell 7 exists in the active class: A_C <= 6). + if constexpr (A_C == 6u && B64V == 2u) { + mma_cell(el, d_tmem, + ab_desc_lo(smbase, A_OFF, s2, 4) & 0xFFFFu, a_hi64, + ab_desc_lo(smbase, B_OFF, s2, 4) & 0xFFFFu, b_hi64, + idesc_k64b, sfa_t1, sfb_t2, 0xFFFFFFFFu); + } + commit_consumed(T_WIN >= 3u ? el : 0u, smbase, + (ab_bar2 + AB_FREE_DELTA) - smbase, commit_mask); + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; + }; + if (b64 != 0u) { + // Once-per-tile subclass dispatch (a = tail_cells - b, even). + const uint32_t a_c = tail_cells - b64; + if (b64 == 1u) { + if (a_c == 0u) return k64_tail(kdi<0>{}, kdi<1>{}); + if (a_c == 2u) return k64_tail(kdi<2>{}, kdi<1>{}); + if (a_c == 4u) return k64_tail(kdi<4>{}, kdi<1>{}); + return k64_tail(kdi<6>{}, kdi<1>{}); + } + if (a_c == 0u) return k64_tail(kdi<0>{}, kdi<2>{}); + if (a_c == 2u) return k64_tail(kdi<2>{}, kdi<2>{}); + if (a_c == 4u) return k64_tail(kdi<4>{}, kdi<2>{}); + return k64_tail(kdi<6>{}, kdi<2>{}); + } + + // ---- the exact-K tail group: tail_cells live K=96 steps out of 8 ---- + // Cell liveness rides the elect predicate (a 0-predicate op never enters + // the tcgen05 pipe); SF batch b exists iff tail_cells > 2b, and skipped + // batches advance NO ring cursor — the SF producer skips symmetrically. + const uint32_t el_c1 = (tail_cells >= 2u) ? el : 0u; + const uint32_t el_c2 = (tail_cells >= 3u) ? el : 0u; + const uint32_t el_c3 = (tail_cells >= 4u) ? el : 0u; + const uint32_t el_c4 = (tail_cells >= 5u) ? el : 0u; + const uint32_t el_c5 = (tail_cells >= 6u) ? el : 0u; + const uint32_t el_c6 = (tail_cells >= 7u) ? el : 0u; + const uint32_t el_c7 = (tail_cells >= 8u) ? el : 0u; + const bool b1_live = (tail_cells > 2u); + const bool b2_live = (tail_cells > 4u); + const bool b3_live = (tail_cells > 6u); + + // SF batch 0 (cells 0,1) — a live tail always has cell 0. + const uint32_t sf_bar0 = smbase + (sf_slot << 3) + SF_READY_OFF; + mbar_wait_relaxed(sf_bar0, sf_phase); + uint32_t sf_slot1, sf_phase1; + { + const uint32_t nx = sf_slot + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot1 = wrap ? 0u : nx; + sf_phase1 = sf_phase ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar0 + SF_FREE_DELTA) - smbase, commit_mask); + const uint32_t sf_bar1 = smbase + (sf_slot1 << 3) + SF_READY_OFF; + const uint32_t ab_bar0 = smbase + (ab_slot << 3); + mbar_wait_relaxed(ab_bar0, ab_phase); + uint32_t ab_slot1, ab_phase1; + { + const uint32_t nx = ab_slot + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot1 = wrap ? 0u : nx; + ab_phase1 = ab_phase ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar1 = smbase + (ab_slot1 << 3); + { + const uint32_t acc_bar = smbase + ACC_FREE_OFF; + mbar_wait_gated_relaxed(acc_bar, acc_wait_parity, !need_acc_wait); + } + mma_cell(el, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 0), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 0), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, acc_pred); // step 0, atom 0 + mma_cell(el_c1, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot, 3), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot, 3), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 1, atom 3 + + // SF batch 1 (cells 2,3). + uint32_t sf_slot2 = sf_slot1, sf_phase2 = sf_phase1; + if (b1_live) { + mbar_wait_relaxed(sf_bar1, sf_phase1); + { + const uint32_t nx = sf_slot1 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot2 = wrap ? 0u : nx; + sf_phase2 = sf_phase1 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot1, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar1 + SF_FREE_DELTA) - smbase, + commit_mask); + } + const uint32_t sf_bar2 = smbase + (sf_slot2 << 3) + SF_READY_OFF; + mbar_wait_relaxed(ab_bar1, ab_phase1); + uint32_t ab_slot2, ab_phase2; + { + const uint32_t nx = ab_slot1 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot2 = wrap ? 0u : nx; + ab_phase2 = ab_phase1 ^ (wrap ? 1u : 0u); + } + const uint32_t ab_bar2 = smbase + (ab_slot2 << 3); + mma_cell(el_c2, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot, ab_slot1, 6), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot, ab_slot1, 6), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 2, atom 6 + commit_consumed(el, smbase, (ab_bar0 + AB_FREE_DELTA) - smbase, commit_mask); + mma_cell(el_c3, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 1), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 1), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 3, atom 1 + + // SF batch 2 (cells 4,5). + uint32_t sf_slot3 = sf_slot2, sf_phase3 = sf_phase2; + if (b2_live) { + mbar_wait_relaxed(sf_bar2, sf_phase2); + { + const uint32_t nx = sf_slot2 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot3 = wrap ? 0u : nx; + sf_phase3 = sf_phase2 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot2, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar2 + SF_FREE_DELTA) - smbase, + commit_mask); + } + const uint32_t sf_bar3 = smbase + (sf_slot3 << 3) + SF_READY_OFF; + mma_cell(el_c4, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot1, 4), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot1, 4), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 4, atom 4 + mbar_wait_relaxed(ab_bar2, ab_phase2); + { + const uint32_t nx = ab_slot2 + 1; + const bool wrap = (nx == AB_RING_SLOTS); + ab_slot = wrap ? 0u : nx; + ab_phase = ab_phase2 ^ (wrap ? 1u : 0u); + } + mma_cell(el_c5, d_tmem, + ab_desc_lo_straddle(smbase, A_OFF, ab_slot1, ab_slot2, 7), a_desc_hi, + ab_desc_lo_straddle(smbase, B_OFF, ab_slot1, ab_slot2, 7), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 5, atom 7 + + // SF batch 3 (cells 6,7). + if (b3_live) { + mbar_wait_relaxed(sf_bar3, sf_phase3); + { + const uint32_t nx = sf_slot3 + 1; + const bool wrap = (nx == SF_RING_SLOTS); + sf_slot = wrap ? 0u : nx; + sf_phase = sf_phase3 ^ (wrap ? 1u : 0u); + } + cp_sf_buffer(el, sf_slot3, sfa_t0, sfa_t1, sfa_t2, sfb_t0, sfb_t1, + sfb_t2, sfb_t3, sfb_t4, sfb_t5, + sfa_src_hi, sfa_src_lo, sfb_src_hi, sfb_src_lo); + commit_consumed(el, smbase, (sf_bar3 + SF_FREE_DELTA) - smbase, + commit_mask); + } else { + sf_slot = sf_slot3; + sf_phase = sf_phase3; + } + commit_consumed(el, smbase, (ab_bar1 + AB_FREE_DELTA) - smbase, commit_mask); + mma_cell(el_c6, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 2), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 2), b_desc_hi, + idesc_w0, sfa_t0, sfb_t0, 0xFFFFFFFFu); // step 6, atom 2 + mma_cell(el_c7, d_tmem, ab_desc_lo(smbase, A_OFF, ab_slot2, 5), a_desc_hi, + ab_desc_lo(smbase, B_OFF, ab_slot2, 5), b_desc_hi, + idesc_w1, sfa_w1, sfb_w1, 0xFFFFFFFFu); // step 7, atom 5 + commit_consumed(el, smbase, (ab_bar2 + AB_FREE_DELTA) - smbase, commit_mask); + + return MainloopState{d_parity, ab_phase, ab_slot, sf_phase, sf_slot}; +} + +} // namespace nvfp4 + +// ============================================================================ +// The kernel. 224 threads = 7 warps, four jobs (blog: the warp table): +// warps 0-3 epilogue: drain TMEM, convert to FP16, store straight to global +// warp 4 the MMA warp (even CTA of each pair only): SF copies + MMAs +// warp 5 TMA producer: A and B tiles into the 6-slot ring +// warp 6 TMA producer: scale factors into the 7-slot ring +// ============================================================================ +namespace nvfp4 { + +// Promise ptxas the exact block shape and express the launch grid in cluster +// units. The launch helper below owns the corresponding grid convention. +#define NVFP4_BLOCK_ATTR __block_size__((TB_SIZE, 1, 1)) +constexpr bool GRID_IN_CLUSTERS = true; + +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void nvfp4_gemm_kernel(const __grid_constant__ CUtensorMap A_tmap, + const __grid_constant__ CUtensorMap B_tmap, + const __grid_constant__ CUtensorMap SFA_tmap, + const __grid_constant__ CUtensorMap SFB_tmap, + __half* C_ptr, int M, int N, int K, + int table_offset + ) { + extern __shared__ __align__(1024) char smem_buf[]; + SmemCD& smem = *reinterpret_cast(smem_buf); + + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const uint32_t crank = ptx::cluster_rank(); + const uint32_t m_pair = crank & 1u; // 0 = the pair's even CTA + + const int cluster_grid_m = (M + CLUSTER_M - 1) / CLUSTER_M; + const int grid_n = (N + BLOCK_N - 1) / BLOCK_N; + const int total_tiles = cluster_grid_m * grid_n; + + // Exact-K tail geometry (the exact-K feed, Other shapes; the host gates K % 16 == 0). + const int full_groups = K / 768; + const int tail_cells = (K - full_groups * 768 + (MMA_K - 1)) / MMA_K; + const int t_sf = (tail_cells + 1) / 2; // live SF slots in the tail + const int num_groups = full_groups + (tail_cells ? 1 : 0); + // K64-dispatch class (banner): r = a x 96 + b x 64, minimal b in {0,1,2}; + // active only for EVEN a (r mod 192 in {64,128}: the with-id SF word is + // architecturally sf_id=0-only at k_dim=0, and even a also excludes the + // window-crossing decomposition). Cell count a + b equals tail_cells and + // t_sf is unchanged, so only the A/B window count differs on the K64 path. + const int k_rem = K - full_groups * 768; + const int b64_try = (k_rem % 96 == 64) ? 1 : (k_rem % 96 == 32) ? 2 : 0; + const int a64 = tail_cells - b64_try; + const int b64 = (b64_try != 0 && k_rem % 32 == 0 && a64 >= 0 && + !(a64 & 1)) ? b64_try : 0; + // Live A/B tail windows on the K64 path (b == 0 keeps all 3). + const int t_win = b64 ? (k_rem / 2 + 127) / 128 : 3; + + // Persistent cluster index. The device always sees the grid in CTA units + // (2, 1, clusters) regardless of the host-side launch-unit regime. + const bool grid_z = (gridDim.z > 1); + const int num_clusters = grid_z ? int(gridDim.z) : int(gridDim.x) / 2; + const int cluster_id = grid_z ? int(blockIdx.z) : int(blockIdx.x) / 2; + + // Placement audit (blog: the L2 trick): one predicated atom.global.add checks + // this SM's censused side against the plan's side for this cluster slot, + // on every launch, timed runs included. A nonzero counter voids the run. + { + unsigned smid; + asm volatile("mov.u32 %0, %%smid;" : "=r"(smid)); + const unsigned bad = unsigned(tid == 0) & + unsigned(l2a_sm_side[smid] != l2a_cluster_side[cluster_id]); + asm volatile( + "{\n\t.reg .pred p;\n\t.reg .u32 old;\n\t" + "setp.ne.u32 p, %1, 0;\n\t" + "@p atom.global.add.u32 old, [%0], 1;\n\t}" + :: "l"(&l2a_placement_errors), "r"(bad) : "memory"); + } + + // Which output tile does persistent work item `t` process? + auto tile_of = [&](int t) -> int { + // The ownership schedule is data: a small constant table built on the + // host. to_uniform keeps its decode on the uniform register datapath. + return int(to_uniform(uint32_t(l2a_route_tables[table_offset + t]))); + }; + + auto producer_tile_of = [&](int t) -> int { + // This loop runs only on lane 0, so its route-table load is already + // scalar. tile_of() instead performs a warp-wide uniform reduction. + return int(l2a_route_tables[table_offset + t]); + }; + + // ---- mbarrier init (warp 0, lane 0) ---- + if (warp_id == 0 && ptx::lane_id() == 0) { + // Arrival counts are derived from the topology, never tuned: + // "ready" = one producer expect_tx per slot (bit-24 routes both + // peers' bytes to the even CTA's barrier); "free" = one consumer + // commit per slot reuse. + #pragma unroll + for (int s = 0; s < CD_AB_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.ab_ready[s], 1); + ptx::mbar_init(&smem.hdr.ab_free[s], 1); + } + #pragma unroll + for (int s = 0; s < CD_SF_SLOTS; ++s) { + ptx::mbar_init(&smem.hdr.sf_ready[s], 1); + ptx::mbar_init(&smem.hdr.sf_free[s], 1); + } + ptx::mbar_init(&smem.hdr.acc_ready, 1); + ptx::mbar_init(&smem.hdr.acc_free, N_EPI_WARPS * CTA_GROUP); // 4 x 2 = 8 + ptx::mbar_init(&smem.hdr.dealloc, 32); // one warp's lanes, cross-peer + ptx::fence_mbarrier_init_release_cluster(); + } + ptx::cluster_sync_rel_acq(); + + // TMEM allocation: epilogue warp 0 allocates for the pair; warps 0-4 + // rendezvous before reading the granted base. The TMA producer warps + // never touch TMEM and proceed straight to their loops. + uint32_t taddr = 0; + if (warp_id <= 4) { + if (warp_id == 0) + ptx::tcgen05_alloc_2sm(ptx::to_shared(&smem.hdr.tmem_addr), + N_ALLOC_COLS); + ptx::bar_sync(2, 160); + taddr = smem.hdr.tmem_addr; + } + + const int wg = warp_id >> 2; // 0 = epilogue warps, 1 = producers + MMA + if (wg == 1) { + const int pw = warp_id - 4; + if (pw == 1 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 5: lane 0 owns the complete A/B TMA producer loop. + // ================================================================ + // free_phase starts at 1: a fresh mbarrier is phase 0, so the + // first ring's worth of "free" waits falls through. + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + // The pair splits the 256x256 tile: even CTA holds A rows + // [0,128) and B cols [0,128); odd CTA the other halves. + const int m_blk = m_row * CTA_GROUP + int(m_pair); + const int n_blk = n_grp * CTA_GROUP + int(m_pair); + #pragma unroll 1 + for (int g = 0; g < num_groups; ++g) { + // Three full 128-byte windows cover the packed K=768 group. + #pragma unroll + for (int sub = 0; sub < AB_TMA_SUBS; ++sub) { + // The K64-dispatch tail ends at logical K: windows + // past t_win are never read — skip them entirely (no + // wait, no expect_tx, no cursor advance; the consumer + // skips symmetrically). Inactive remainders keep + // t_win == 3, so this never fires for them. + if (g == full_groups && sub >= t_win) break; + ptx::mbar_wait_parity(&smem.hdr.ab_free[slot], free_phase); + const uint32_t a_s = + base + CD_A_OFF + uint32_t(slot) * CD_AB_STRIDE; + const uint32_t b_s = + base + CD_B_OFF + uint32_t(slot) * CD_AB_STRIDE; + uint64_t* bar = &smem.hdr.ab_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, AB_TX); + // A rides the multicast form with a self mask so the + // completion routes to the even peer's barrier. + const uint16_t mask_a = uint16_t(1u << m_pair); + // Read the same row-major FP4 allocation as cuBLAS. + // K192 uses four 96-byte boxes; K256 uses three + // 128-byte boxes. The tensor map zero-fills past K. + const CUtensorMap* a_map = &A_tmap; + const CUtensorMap* b_map = &B_tmap; + const int x = g * 384 + sub * AB_LOAD_BYTES; + const int y_a = m_blk * BLOCK_M; + const int y_b = n_blk * BLOCK_N_PEER; + const int z_a = 0; + const int z_b = 0; + ptx::cp_async_bulk_tensor_3d_load_multicast( + a_s, a_map, x, y_a, z_a, + bar, mask_a + ); + ptx::cp_async_bulk_tensor_3d_load_2sm_bit24( + b_s, b_map, x, y_b, z_b, + bar + ); + if (++slot == CD_AB_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + // Drain: wait for the last-filled slot to be consumed, then hand + // it one clean-shutdown expect_tx. + { + const int t_slot = (slot == 0) ? CD_AB_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.ab_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.ab_ready[t_slot], AB_TX); + } + } else if (pw == 2 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 6: scale-factor TMA producer. SFA is per-peer (each CTA's + // own M half); SFB is the one genuinely multicast operand — the + // cta_group::2 copy reads BOTH peers' smem as one logical + // operand, so both copies must be byte-identical (blog: the scale-factor section). + // ================================================================ + int slot = 0, free_phase = 1; + const uint32_t base = ptx::to_shared(&smem); + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = producer_tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const int m_blk = m_row * CTA_GROUP + int(m_pair); + for (int g = 0; g < num_groups; ++g) { + // 4 SF ring slots per 768-K group (2 K=96 steps each). + #pragma unroll + for (int sub = 0; sub < 4; ++sub) { + // The tail group stores only its live SF slots; dead + // slots are skipped entirely (no wait, no arrive, no + // cursor advance — the consumer skips symmetrically). + if (g == full_groups && sub >= t_sf) break; + ptx::mbar_wait_parity(&smem.hdr.sf_free[slot], free_phase); + const uint32_t sfa_s = + base + CD_SFA_OFF + uint32_t(slot) * CD_SFA_STRIDE; + const uint32_t sfb_s = + base + CD_SFB_OFF + uint32_t(slot) * CD_SFB_STRIDE; + uint64_t* bar = &smem.hdr.sf_ready[slot]; + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster(bar, SF_TX); + // VEC16 is [outer-block][K-group-of-4][128 rows x 4]. + // Three adjacent groups are the 128 x 12 scale slot. + const int sf_group = (g * 4 + sub) * 3; + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfa_s, &SFA_tmap, 0, 0, sf_group, m_blk, bar, + uint16_t(1u << m_pair) + ); + ptx::cp_async_bulk_tensor_4d_load_multicast( + sfb_s + m_pair * uint32_t(CD_SFB_STRIDE / 2), + &SFB_tmap, 0, 0, sf_group, + n_grp * CTA_GROUP + int(m_pair), bar, + uint16_t(0x3) + ); + if (++slot == CD_SF_SLOTS) { slot = 0; free_phase ^= 1; } + } + } + } + { + const int t_slot = (slot == 0) ? CD_SF_SLOTS - 1 : slot - 1; + const int t_phase = (slot == 0) ? free_phase : (free_phase ^ 1); + ptx::mbar_wait_parity(&smem.hdr.sf_free[t_slot], + uint32_t(t_phase)); + if (m_pair == 0) + ptx::mbar_arrive_expect_tx_cluster( + &smem.hdr.sf_ready[t_slot], SF_TX); + } + } else if (pw == 0 && m_pair == 0 && ptx::lane_id() == 0) { + // ================================================================ + // Warp 4: the MMA warp — issued once per cluster, on the even + // CTA. The odd CTA's warp 4 does nothing all kernel (blog: the warp table). + // The mainloop needs no tile coordinates: operands arrive through + // the rings; this warp only needs its share of the trip count. + // ================================================================ + const uint32_t smbase = ptx::to_shared(&smem); + const uint16_t commit_mask = 0x3; // both CTAs of this pair + MainloopState st{1u, 0u, 0u, 0u, 0u}; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + st = mainloop_tile(taddr, smbase, st, commit_mask, + uint32_t(full_groups), + uint32_t(tail_cells), + uint32_t(b64), uint32_t(t_win) + ); + // The outer role branch already leaves exactly one issuer. + asm volatile( + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;" + :: "r"(smbase + uint32_t(MB_ACC_READY)), + "h"(uint16_t(0x3)) : "memory"); + } + // Drain the final tile's buffer-free arrive before exiting. + mbar_wait_relaxed(smbase + uint32_t(MB_ACC_FREE), st.d_parity); + } + } else { + // ==================================================================== + // Warps 0-3: the epilogue (blog: the epilogue). Each warp owns a 32-row band + // and drains a tile as eight 32-column tcgen05.ld loads, converts + // FP32->FP16 in packed pairs, and stores straight to global — no + // shared-memory bounce. + // ==================================================================== + const int epi_warp = warp_id; + const uint32_t taddr_lane = uint32_t(epi_warp * 32) << 16; + const uint32_t acc_free_rank = crank & ~1u; // the pair's even CTA + int d_db = 0; + int acc_ready_phase = 0; + for (int t = cluster_id; t < total_tiles; t += num_clusters) { + const int tile = tile_of(t); + const int m_row = tile % cluster_grid_m; + const int n_grp = tile / cluster_grid_m; + const uint32_t tmem_base = taddr + uint32_t(d_db) * D_STRIDE; + mbar_wait_relaxed(ptx::to_shared(&smem.hdr.acc_ready), + uint32_t(acc_ready_phase)); + acc_ready_phase ^= 1; + + const int off_m = (m_row * CTA_GROUP + int(m_pair)) * BLOCK_M; + const int off_n = n_grp * BLOCK_N; + const int row = off_m + epi_warp * 32 + lane_id; + const bool row_in = (row < M); + constexpr int EPI_BANDS = BLOCK_N / EPI_CHUNK; // 8 + #pragma unroll 1 + for (int k = 0; k < EPI_BANDS; ++k) { + // Buffer 0 drains high-to-low and buffer 1 low-to-high, so + // their shared columns [220,256) drain first for both. + const int parity = d_db & 1; + const int band = (parity == 0) ? (EPI_BANDS - 1 - k) : k; + const uint32_t taddr_n = + tmem_base + uint32_t(band) * 32 + taddr_lane; + uint32_t w[32]; + ptx::tcgen05_ld_32x32b_x32(taddr_n, w); + // Return the old buffer as soon as its shared edge is out. + if (k == 1) { + ptx::tcgen05_wait_ld_then_arrive_release( + &smem.hdr.acc_free, acc_free_rank, + /*do_arrive=*/ptx::lane_id() == 0); + } + if (!row_in) continue; + const int col_base = off_n + band * 32; + uint32_t h[16]; + #pragma unroll + for (int j = 0; j < 16; ++j) { + h[j] = ptx::cvt_pack_f16x2(__int_as_float(w[2 * j + 1]), + __int_as_float(w[2 * j])); + } + const size_t row_base = size_t(row) * size_t(N); + const bool aligned32 = ((row_base & 15u) == 0u) && + ((uintptr_t(C_ptr) & 31u) == 0u); + if (aligned32 && (col_base & 15) == 0) { + // Precondition here: row_base is 16-element (32-byte) + // aligned — the aligned32 guard owns it. Restating it to + // ptxas via __builtin_assume compiled this file + // SASS-identical (and measured perf-null on the + // production kernel), so the guard stands alone. + #pragma unroll + for (int q = 0; q < 2; ++q) { + const int c0 = col_base + q * 16; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 16 <= N) { + uint4 lo{h[q*8+0], h[q*8+1], h[q*8+2], h[q*8+3]}; + uint4 hi{h[q*8+4], h[q*8+5], h[q*8+6], h[q*8+7]}; + ptx::st_global_na_ef_b256(cdst, lo, hi); + } else { + // N tail: 128-bit then scalar. + #pragma unroll + for (int p = 0; p < 2; ++p) { + const int cq = c0 + p * 8; + if (cq + 8 <= N) { + uint4 v{h[q*8+p*4+0], h[q*8+p*4+1], + h[q*8+p*4+2], h[q*8+p*4+3]}; + ptx::st_global_na_b128( + reinterpret_cast(cdst + p * 8), v); + } else { + const uint16_t* hb = + reinterpret_cast( + &h[q * 8 + p * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (cq + c < N) cdst[p * 8 + c] = hb[c]; + } + } + } + } + } else { + #pragma unroll + for (int q = 0; q < 4; ++q) { + const int c0 = col_base + q * 8; + uint16_t* cdst = + reinterpret_cast(C_ptr) + row_base + c0; + if (c0 + 8 <= N) { + uint4 v{h[q*4+0], h[q*4+1], h[q*4+2], h[q*4+3]}; + ptx::st_global_na_b128( + reinterpret_cast(cdst), v); + } else { + const uint16_t* hb = + reinterpret_cast(&h[q * 4]); + #pragma unroll + for (int c = 0; c < 8; ++c) + if (c0 + c < N) cdst[c] = hb[c]; + } + } + } + } + if (++d_db == D_DB) d_db = 0; + } + } + + // Teardown: only the epilogue warps rendezvous per CTA; warp 0 then runs + // the cross-peer dealloc handshake. Producer and MMA warps exit freely — + // their drains above already order their work. + if (wg == 0) { + ptx::bar_sync(3, 128); + if (warp_id == 0) { + ptx::tcgen05_relinquish_2sm(); + const uint32_t peer_rank = crank ^ 1u; + // All 32 lanes arrive the PEER's count-32 barrier = one flip. + ptx::mbar_arrive_cluster_release(&smem.hdr.dealloc, peer_rank); + ptx::mbar_wait_parity(&smem.hdr.dealloc, 0); + ptx::tcgen05_dealloc_2sm(taddr, N_ALLOC_COLS); + } + } +} + +// Cluster->SM census probe: same cluster shape, block shape, and full smem +// footprint as the GEMM kernel, so the persistent grid lands on the same +// physical SM assignment the timed launches will get. +__global__ +NVFP4_BLOCK_ATTR +__cluster_dims__(2, 1, 1) +__launch_bounds__(TB_SIZE, 1) +void l2a_cluster_probe(unsigned* smids) { + extern __shared__ char pad[]; + if (threadIdx.x != 0) return; + pad[0] = 0; + unsigned rank, smid; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + asm volatile("mov.u32 %0, %%smid;" : "=r"(smid)); + const int cluster = gridDim.z > 1 ? int(blockIdx.z) : int(blockIdx.x) / 2; + smids[cluster * CTA_GROUP + int(rank)] = smid; +} + +// Keep the host launch convention paired with the kernel's block attribute. +static inline dim3 launch_grid(int clusters) { + return GRID_IN_CLUSTERS ? dim3(1, 1, unsigned(clusters)) + : dim3(2, 1, unsigned(clusters)); +} + +} // namespace nvfp4 + +// ============================================================================ +// Host-side TMA descriptor support. The shared input and benchmark code lives +// in main.cu. +// ============================================================================ +namespace tmap { + +// Thin wrappers over cuTensorMapEncodeTiled. Dim 0 is the innermost +// (stride-1) axis; strides for dim 1+ are in bytes. +inline CUtensorMap encode_tiled_3d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_depth, uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t depth_stride_bytes, + uint32_t box_depth, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[3] = { global_cols, global_rows, global_depth }; + cuuint64_t global_strides[2] = { row_stride_bytes, depth_stride_bytes }; + cuuint32_t box_dim[3] = { box_cols, box_rows, box_depth }; + cuuint32_t element_strides[3] = { 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 3, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +inline CUtensorMap encode_tiled_4d_box(void* global_ptr, CUtensorMapDataType dtype, + uint64_t global_slots, uint64_t global_tiles, + uint64_t global_rows, uint64_t global_cols, + uint64_t row_stride_bytes, uint64_t tile_stride_bytes, + uint64_t slot_stride_bytes, + uint32_t box_tiles, uint32_t box_rows, uint32_t box_cols, + CUtensorMapSwizzle swizzle) { + cuuint64_t global_dim[4] = { global_cols, global_rows, global_tiles, + global_slots }; + cuuint64_t global_strides[3] = { row_stride_bytes, tile_stride_bytes, + slot_stride_bytes }; + cuuint32_t box_dim[4] = { box_cols, box_rows, box_tiles, 1u }; + cuuint32_t element_strides[4] = { 1, 1, 1, 1 }; + CUtensorMap m{}; + CU_CHECK(cuTensorMapEncodeTiled( + &m, dtype, 4, global_ptr, global_dim, global_strides, box_dim, + element_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); + return m; +} + +} // namespace tmap + +namespace host { + +using namespace nvfp4; + +// ---- implementation-specific tensor maps --------------------------------- + +inline CUtensorMap make_ab_tmap(void* ptr, int rows, int K, + int row_stride) { + return tmap::encode_tiled_3d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*depth=*/1, uint64_t(rows), /*cols=*/uint64_t(K / 2), + /*row_stride=*/uint64_t(row_stride), + /*depth_stride=*/uint64_t(rows) * uint64_t(row_stride), + /*box=*/1, BLOCK_M, AB_LOAD_BYTES, CU_TENSOR_MAP_SWIZZLE_128B); +} + +inline CUtensorMap make_sf_tmap(uint8_t* ptr, int outer, int K) { + const uint64_t sf_inner = (uint64_t(K) + 63) / 64 * 4; + const uint64_t outer_blocks = (uint64_t(outer) + 127) / 128; + return tmap::encode_tiled_4d_box( + ptr, CU_TENSOR_MAP_DATA_TYPE_UINT8, + /*slots=*/outer_blocks, /*tiles=*/sf_inner / 4, + /*rows=*/4, /*cols=*/128, + /*row_stride=*/128, /*tile_stride=*/512, + /*slot_stride=*/sf_inner * 128, + /*box_tiles=*/3, /*box_rows=*/4, /*box_cols=*/128, + CU_TENSOR_MAP_SWIZZLE_NONE); +} +} // namespace host + +// ============================================================================ +// Runtime L2-side map (blog: "Interlude: your L2 is two L2s"). Everything +// here is measured at runtime — nothing hardcodes a split or trusts the +// address hash without checking it. The probing approach follows ademeure's +// QuickRunCUDA side_aware experiment (credited in the post). +// P1 census: one dependent-load atomic round-trip per SM against a single +// probe line; the bimodal split classifies each SM's side. +// P2 address hash: walk single address bits; on GB300 the model is +// parity(phys_addr & 0x1EF000) XOR a per-2MiB-slab phase. +// P3 model check: predicted vs measured side over representative slots — +// any mismatch fails closed (the harness then runs natural raster). +// ============================================================================ +namespace l2side { + +inline constexpr uint64_t kSlotBytes = 4096; // hash granule (bit 12 is live) +inline constexpr uint64_t kSlabBytes = 2ull << 20; +inline constexpr uint32_t kExpectedHash = 0x1EF000u; +inline constexpr double kTimingResolutionNs = 1.0; + +inline int popcount64(uint64_t v) { return __builtin_popcountll(v); } + +struct RuntimeMap { + int nsm = 0; + std::array side_counts{}; + std::array sm_side{}; + std::vector slab_phase; + double threshold = 0.0; + int near_smid = -1; + double near_ns = 0.0; + double far_ns = 0.0; + double classification_gap_ns = 0.0; + double minimum_sample_margin_ns = 0.0; + double maximum_effective_jitter_ns = 0.0; + double minimum_confidence_ns = 0.0; + int stability_repeats = 1; + uint32_t hash = 0; + int model_checks = 0; + int model_mismatches = 0; + // Every side verdict and its distance from the threshold, in probe + // order — the stability pass requires each margin to beat its own + // observed jitter. + std::vector classification_bits; + std::vector classification_margins_ns; +}; + +namespace detail { + +inline constexpr size_t kFatSmem = 130u << 10; // one resident block per SM + +[[noreturn]] inline void fail(const char* phase, const std::string& receipt) { + throw std::runtime_error(std::string("l2side ") + phase + " failed: " + receipt); +} + +inline void cuda_check(cudaError_t error, const char* operation) { + if (error != cudaSuccess) + fail("CUDA", std::string(operation) + ": " + cudaGetErrorString(error)); +} + +template +class DeviceBuffer { +public: + explicit DeviceBuffer(size_t count) { + cuda_check(cudaMalloc(&ptr_, count * sizeof(T)), "cudaMalloc scratch"); + } + ~DeviceBuffer() { if (ptr_ != nullptr) cudaFree(ptr_); } + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + T* get() const { return ptr_; } +private: + T* ptr_ = nullptr; +}; + +static __device__ __forceinline__ unsigned probe_smid() { + unsigned value; + asm volatile("mov.u32 %0, %%smid;" : "=r"(value)); + return value; +} + +static __device__ __forceinline__ unsigned long long globaltimer_ns() { + unsigned long long value; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(value)); + return value; +} + +// Dependent-load chain of atomics: each result selects the next address, so +// latency can't hide behind parallelism. Probe words are pre-cleared, so +// value>>31 stays zero while still defeating unrolling. +static __device__ __forceinline__ unsigned long long +rtt_chain(unsigned* probe, int reps) { + unsigned value = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) value = atomicAdd(probe + (value >> 31), 1u); + const unsigned long long begin = globaltimer_ns(); + for (int i = 0; i < reps; ++i) + value = atomicAdd(probe + (value >> 31), 1u); + const unsigned long long end = globaltimer_ns(); + return (end - begin) / static_cast(reps) + (value >> 31); +} + +// Fat dynamic smem keeps exactly one block resident per SM; the turn counter +// serializes them so every SM times the same line in isolation. +static __global__ void census_kernel(unsigned* probe, volatile unsigned* turn, + unsigned long long* out_ns, + unsigned* out_smid, int reps) { + extern __shared__ char pad[]; + if (threadIdx.x != 0) return; + pad[0] = 0; + const int block = static_cast(blockIdx.x); + while (*turn != static_cast(block)) {} + out_ns[block] = rtt_chain(probe, reps); + out_smid[block] = probe_smid(); + __threadfence(); + *turn = static_cast(block + 1); +} + +static __global__ void clear_probe_words(char* base, + const uint64_t* offsets, int n) { + const int i = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (i >= n) return; + unsigned* probe = reinterpret_cast(base + offsets[i]); + probe[0] = 0; + probe[1] = 0; +} + +// Exactly one block, pinned to the near-side SM, sweeps the offset list. +static __global__ void probe_kernel(char* base, const uint64_t* offsets, + int n, unsigned long long* out_ns, + int target_smid, int reps, int* hit) { + extern __shared__ char pad[]; + if (threadIdx.x != 0) return; + pad[0] = 0; + if (probe_smid() != static_cast(target_smid)) return; + if (atomicCAS(hit, 0, 1) != 0) return; + for (int i = 0; i < n; ++i) + out_ns[i] = rtt_chain(reinterpret_cast(base + offsets[i]), reps); +} + +inline double median(std::vector values) { + std::sort(values.begin(), values.end()); + return values[values.size() / 2]; +} + +inline std::vector +probe_offsets(char* arena, const std::vector& offsets, + int nsm, int near_smid) { + const int n = static_cast(offsets.size()); + DeviceBuffer d_offsets(offsets.size()); + DeviceBuffer d_ns(offsets.size()); + DeviceBuffer d_hit(1); + cuda_check(cudaMemcpy(d_offsets.get(), offsets.data(), + offsets.size() * sizeof(offsets[0]), + cudaMemcpyHostToDevice), "upload probe offsets"); + cuda_check(cudaMemset(d_ns.get(), 0, + offsets.size() * sizeof(unsigned long long)), + "clear probe timings"); + cuda_check(cudaMemset(d_hit.get(), 0, sizeof(int)), "clear probe hit"); + clear_probe_words<<<(n + 255) / 256, 256>>>(arena, d_offsets.get(), n); + cuda_check(cudaGetLastError(), "launch clear_probe_words"); + probe_kernel<<>>(arena, d_offsets.get(), n, + d_ns.get(), near_smid, 48, d_hit.get()); + cuda_check(cudaGetLastError(), "launch probe_kernel"); + cuda_check(cudaDeviceSynchronize(), "synchronize probe_kernel"); + + int hit = 0; + cuda_check(cudaMemcpy(&hit, d_hit.get(), sizeof(hit), + cudaMemcpyDeviceToHost), "copy probe hit"); + if (hit != 1) fail("probe", "target smid " + std::to_string(near_smid) + + " was not scheduled"); + std::vector timings(offsets.size()); + cuda_check(cudaMemcpy(timings.data(), d_ns.get(), + timings.size() * sizeof(timings[0]), + cudaMemcpyDeviceToHost), "copy probe timings"); + return timings; +} + +} // namespace detail + +// One full probe pass: census + address-hash walk + model check. Mutates +// probe words throughout `arena` (2 MiB-aligned, a multiple of 2 MiB long). +inline RuntimeMap probe(char* arena, uint64_t bytes) { + if (arena == nullptr) detail::fail("input", "arena is null"); + if ((reinterpret_cast(arena) & (kSlabBytes - 1)) != 0) + detail::fail("input", "arena is not 2 MiB aligned"); + if (bytes == 0 || (bytes & (kSlabBytes - 1)) != 0) + detail::fail("input", "byte size is not a nonzero multiple of 2 MiB"); + const uint64_t nslabs64 = bytes / kSlabBytes; + if (nslabs64 > 0x7fffffffu) detail::fail("input", "arena has too many slabs"); + + int device = 0; + detail::cuda_check(cudaGetDevice(&device), "cudaGetDevice"); + cudaDeviceProp prop{}; + detail::cuda_check(cudaGetDeviceProperties(&prop, device), + "cudaGetDeviceProperties"); + if (prop.major != 10 || prop.minor != 3) + detail::fail("input", "device is not an sm_103 (GB300/B300) part"); + if (prop.multiProcessorCount <= 0 || prop.multiProcessorCount > 256) + detail::fail("input", "nsm=" + std::to_string(prop.multiProcessorCount)); + if (prop.sharedMemPerBlockOptin < detail::kFatSmem) + detail::fail("input", "opt-in shared memory is below 130 KiB"); + + RuntimeMap result; + result.nsm = prop.multiProcessorCount; + result.sm_side.fill(-1); + + detail::cuda_check(cudaMemset(arena, 0, 2 * sizeof(unsigned)), + "clear census probe"); + detail::DeviceBuffer d_turn(1); + detail::DeviceBuffer d_ns(result.nsm); + detail::DeviceBuffer d_smid(result.nsm); + detail::cuda_check(cudaMemset(d_turn.get(), 0, sizeof(unsigned)), + "clear census turn"); + detail::cuda_check(cudaFuncSetAttribute( + detail::census_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(detail::kFatSmem)), + "set census dynamic shared memory"); + detail::cuda_check(cudaFuncSetAttribute( + detail::probe_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(detail::kFatSmem)), + "set probe dynamic shared memory"); + detail::census_kernel<<>>( + reinterpret_cast(arena), d_turn.get(), d_ns.get(), + d_smid.get(), 64); + detail::cuda_check(cudaGetLastError(), "launch census_kernel"); + detail::cuda_check(cudaDeviceSynchronize(), "synchronize census_kernel"); + + std::vector census_ns(result.nsm); + std::vector census_smid(result.nsm); + detail::cuda_check(cudaMemcpy(census_ns.data(), d_ns.get(), + census_ns.size() * sizeof(census_ns[0]), + cudaMemcpyDeviceToHost), "copy census timings"); + detail::cuda_check(cudaMemcpy(census_smid.data(), d_smid.get(), + census_smid.size() * sizeof(census_smid[0]), + cudaMemcpyDeviceToHost), "copy census smids"); + + // Split the bimodal distribution at the largest gap near the middle. + std::vector sorted(census_ns.begin(), census_ns.end()); + std::sort(sorted.begin(), sorted.end()); + int cut = result.nsm / 2 - 1; + double largest_gap = -1.0; + for (int i = result.nsm / 4; i < 3 * result.nsm / 4 && i + 1 < result.nsm; + ++i) { + const double gap = sorted[i + 1] - sorted[i]; + if (gap > largest_gap) { + largest_gap = gap; + cut = i; + } + } + result.threshold = (sorted[cut] + sorted[cut + 1]) * 0.5; + result.classification_gap_ns = sorted[cut + 1] - sorted[cut]; + result.near_ns = detail::median( + std::vector(sorted.begin(), sorted.begin() + cut + 1)); + result.far_ns = detail::median( + std::vector(sorted.begin() + cut + 1, sorted.end())); + + const auto record_classification = [&](double timing_ns) { + const int side = timing_ns > result.threshold; + const double margin = side ? timing_ns - result.threshold + : result.threshold - timing_ns; + result.classification_bits.push_back(static_cast(side)); + result.classification_margins_ns.push_back(margin); + return side; + }; + + unsigned long long best_near = ~0ull; + std::array sm_ns{}; + for (int block = 0; block < result.nsm; ++block) { + const unsigned sm = census_smid[block]; + if (sm >= result.sm_side.size()) + detail::fail("P1", "smid=" + std::to_string(sm) + " exceeds 255"); + if (result.sm_side[sm] != -1) + detail::fail("P1", "duplicate smid=" + std::to_string(sm)); + const int side = census_ns[block] > result.threshold; + result.sm_side[sm] = side; + sm_ns[sm] = static_cast(census_ns[block]); + ++result.side_counts[side]; + if (side == 0 && census_ns[block] < best_near) { + best_near = census_ns[block]; + result.near_smid = static_cast(sm); + } + } + for (int sm = 0; sm < result.nsm; ++sm) { + if (record_classification(sm_ns[sm]) != result.sm_side[sm]) + detail::fail("P1", "SM classification receipt mismatch"); + } + // Splits are not always even — 72/80 and 74/78 boots exist — but a split + // outside these bounds means the probe read something else entirely. + const double ratio = result.far_ns / result.near_ns; + const bool p1_ok = ratio >= 1.3 && result.side_counts[0] >= 60 && + result.side_counts[0] <= 92 && + result.side_counts[1] >= 60 && + result.side_counts[1] <= 92 && result.near_ns >= 80.0 && + result.near_ns <= 500.0 && result.near_smid >= 0; + if (!p1_ok) { + detail::fail("P1", "nsm=" + std::to_string(result.nsm) + + " split=" + std::to_string(result.side_counts[0]) + + "/" + std::to_string(result.side_counts[1]) + + " near_ns=" + std::to_string(result.near_ns) + + " far_ns=" + std::to_string(result.far_ns) + + " ratio=" + std::to_string(ratio)); + } + + // P2: which address bits flip the side? Walk bits 4..20 one at a time. + std::vector offsets{0}; + for (int bit = 4; bit <= 20; ++bit) offsets.push_back(1ull << bit); + const auto hash_ns = + detail::probe_offsets(arena, offsets, result.nsm, result.near_smid); + const bool base_far = record_classification(hash_ns[0]); + for (int bit = 4; bit <= 20; ++bit) { + if (record_classification(hash_ns[bit - 3]) != base_far) + result.hash |= 1u << bit; + } + if (result.hash != kExpectedHash) { + detail::fail("P2", "hash=" + std::to_string(result.hash) + + " expected=" + std::to_string(kExpectedHash)); + } + + // Per-slab phase: allocator slab placement changes per boot, so measure. + const int nslabs = static_cast(nslabs64); + offsets.clear(); + offsets.reserve(nslabs); + for (int slab = 0; slab < nslabs; ++slab) + offsets.push_back(static_cast(slab) * kSlabBytes); + const auto slab_ns = + detail::probe_offsets(arena, offsets, result.nsm, result.near_smid); + result.slab_phase.resize(nslabs); + for (int slab = 0; slab < nslabs; ++slab) + result.slab_phase[slab] = record_classification(slab_ns[slab]); + + // P3: the model must predict every representative slot. One miss = fail. + std::vector representative_slabs; + const auto add_slab = [&](int slab) { + if (std::find(representative_slabs.begin(), representative_slabs.end(), + slab) == representative_slabs.end()) + representative_slabs.push_back(slab); + }; + add_slab(0); + add_slab(std::min(1, nslabs - 1)); + add_slab(nslabs / 2); + add_slab(nslabs - 1); + constexpr std::array slots{1, 2, 3, 5, 8, 17}; + offsets.clear(); + for (int slab : representative_slabs) + for (int slot : slots) + offsets.push_back(static_cast(slab) * kSlabBytes + + static_cast(slot) * kSlotBytes); + const auto model_ns = + detail::probe_offsets(arena, offsets, result.nsm, result.near_smid); + size_t i = 0; + for (int slab : representative_slabs) { + for (int slot : slots) { + const auto within = static_cast(slot) * kSlotBytes; + const int expected = (popcount64(within & kExpectedHash) & 1) ^ + result.slab_phase[slab]; + const int observed = record_classification(model_ns[i++]); + ++result.model_checks; + result.model_mismatches += observed != expected; + } + } + if (result.model_mismatches != 0) { + detail::fail("P3", "model mismatches=" + + std::to_string(result.model_mismatches) + "/" + + std::to_string(result.model_checks)); + } + result.minimum_sample_margin_ns = *std::min_element( + result.classification_margins_ns.begin(), + result.classification_margins_ns.end()); + result.minimum_confidence_ns = result.minimum_sample_margin_ns; + return result; +} + +// Repeat the whole probe; every SM/hash/slab/model verdict must hold, and +// each sample's minimum margin must exceed its own observed jitter (floored +// at the timer resolution). Any disagreement fails closed — no majority vote. +inline RuntimeMap probe_stable(char* arena, uint64_t bytes, int repeats = 3) { + if (repeats < 2) detail::fail("stability", "repeats must be at least two"); + RuntimeMap result = probe(arena, bytes); + double min_gap = result.classification_gap_ns; + std::vector minimum_margins = result.classification_margins_ns; + std::vector maximum_margins = result.classification_margins_ns; + for (int repeat = 1; repeat < repeats; ++repeat) { + RuntimeMap observed = probe(arena, bytes); + if (observed.nsm != result.nsm || + observed.side_counts != result.side_counts || + observed.sm_side != result.sm_side || + observed.hash != result.hash || + observed.slab_phase != result.slab_phase || + observed.classification_bits != result.classification_bits || + observed.classification_margins_ns.size() != + result.classification_margins_ns.size() || + observed.model_mismatches != 0) { + detail::fail("stability", + "SM or address-side classification changed on repeat " + + std::to_string(repeat + 1)); + } + min_gap = std::min(min_gap, observed.classification_gap_ns); + for (size_t i = 0; i < minimum_margins.size(); ++i) { + minimum_margins[i] = std::min( + minimum_margins[i], observed.classification_margins_ns[i]); + maximum_margins[i] = std::max( + maximum_margins[i], observed.classification_margins_ns[i]); + } + } + double minimum_margin = minimum_margins[0]; + double maximum_effective_jitter = 0.0; + double minimum_confidence = minimum_margins[0]; + for (size_t i = 0; i < minimum_margins.size(); ++i) { + const double observed_jitter = maximum_margins[i] - minimum_margins[i]; + const double effective_jitter = + std::max(observed_jitter, kTimingResolutionNs); + const double confidence = minimum_margins[i] - effective_jitter; + if (!(confidence > 0.0)) { + detail::fail("stability", + "classification sample " + std::to_string(i) + + " margin_ns=" + std::to_string(minimum_margins[i]) + + " observed_jitter_ns=" + std::to_string(observed_jitter) + + " effective_jitter_ns=" + std::to_string(effective_jitter)); + } + minimum_margin = std::min(minimum_margin, minimum_margins[i]); + maximum_effective_jitter = + std::max(maximum_effective_jitter, effective_jitter); + minimum_confidence = std::min(minimum_confidence, confidence); + } + result.classification_gap_ns = min_gap; + result.minimum_sample_margin_ns = minimum_margin; + result.maximum_effective_jitter_ns = maximum_effective_jitter; + result.minimum_confidence_ns = minimum_confidence; + result.stability_repeats = repeats; + return result; +} + +} // namespace l2side + +// ============================================================================ +// The schedule is data (blog Levers 2-3). The host builds a per-shape table +// mapping each persistent work slot to an output tile, proves it covers every +// tile exactly once and honors the ownership contract, uploads it to constant +// memory, and reads it back. Tiles encode m_row + n_col * mblocks (M fast). +// ============================================================================ +namespace sched { + +using nvfp4::L2A_ROUTE_WORK_CAP; + +struct Census { + int clusters = 0; + int side_clusters[2]{}; + std::vector cluster_side; + std::array, 2> clusters_by_side; +}; + +// Classify the cluster->SM census against the probed side map: every SM seen +// exactly once, both CTAs of each cluster on one side, per-side SM counts +// consistent. A failure here is a bug or a mid-run reschedule — abort loudly. +inline Census classify_census(const std::vector& smids, + const l2side::RuntimeMap& map, int sms) { + Census c; + c.clusters = int(smids.size()) / 2; + c.cluster_side.assign(c.clusters, -1); + std::vector seen(sms, 0); + for (int cluster = 0; cluster < c.clusters; ++cluster) { + int side = -1; + for (int rank = 0; rank < 2; ++rank) { + const unsigned smid = smids[size_t(cluster) * 2 + rank]; + if (smid >= unsigned(sms) || seen[smid]++) { + std::fprintf(stderr, "cluster census coverage failed\n"); + std::abort(); + } + const int actual = map.sm_side[smid]; + if (side < 0) side = actual; + if (actual != side) { + // Every census we've taken places both CTAs of a cluster on + // one side; a mixed pair would invalidate the whole plan. + std::fprintf(stderr, "mixed-side physical cluster %d\n", cluster); + std::abort(); + } + } + c.cluster_side[cluster] = side; + ++c.side_clusters[side]; + c.clusters_by_side[side].push_back(cluster); + } + if (std::find(seen.begin(), seen.end(), 0) != seen.end() || + 2 * c.side_clusters[0] != map.side_counts[0] || + 2 * c.side_clusters[1] != map.side_counts[1]) { + std::fprintf(stderr, "cluster/SM side-count gate failed\n"); + std::abort(); + } + return c; +} + +// Per-cluster trip counts under the persistent stride, plus the proof that +// they spread by at most one (the table preserves each cluster's trip count, +// so a wrong spread would make quota assignment impossible). +inline std::vector cluster_trips(int clusters, int total_work) { + std::vector trips(clusters, 0); + for (int cluster = 0; cluster < clusters; ++cluster) + for (int work = cluster; work < total_work; work += clusters) + ++trips[cluster]; + const int trip_floor = total_work / clusters; + const int long_trips = total_work % clusters; + const int trip_ceil = trip_floor + (long_trips != 0); + const int observed_long = trip_ceil == trip_floor + ? 0 : int(std::count(trips.begin(), trips.end(), trip_ceil)); + const int observed_short = + int(std::count(trips.begin(), trips.end(), trip_floor)); + if (trip_ceil - trip_floor > 1 || observed_long != long_trips || + observed_short != clusters - long_trips) { + std::fprintf(stderr, "one-trip-spread proof failed\n"); + std::abort(); + } + return trips; +} + +// The owned M-row partition: split the M rows between the two sides in +// proportion to each side's work share, with at most one boundary row shared +// (split along N). Every re-read of a given A row then comes from one side +// only — you can't make the buffer side-pure (the hash reaches bit 12), but +// you CAN own the reuse. +struct OwnedPartition { + int side_work[2]{}; + int boundary_row = -1; + int boundary_n = 0; + std::array, 2> pool_tiles; // M-fast order within pools +}; + +inline bool make_owned_partition(const Census& census, + const std::vector& trips, + int mblocks, int nblocks, OwnedPartition* out) { + OwnedPartition p; + for (int side = 0; side < 2; ++side) + for (const int cluster : census.clusters_by_side[side]) + p.side_work[side] += trips[cluster]; + const int total_work = mblocks * nblocks; + if (p.side_work[0] + p.side_work[1] != total_work) std::abort(); + p.boundary_row = p.side_work[0] / nblocks; + p.boundary_n = p.side_work[0] % nblocks; + if (p.boundary_row <= 0 || p.boundary_row >= mblocks || + (p.boundary_n != 0 && p.boundary_row + 1 >= mblocks)) { + return false; // grid too small to own — the caller falls back + } + for (int n = 0; n < nblocks; ++n) { + for (int mrow = 0; mrow < p.boundary_row; ++mrow) + p.pool_tiles[0].push_back(mrow + n * mblocks); + if (p.boundary_n != 0 && n < p.boundary_n) + p.pool_tiles[0].push_back(p.boundary_row + n * mblocks); + if (p.boundary_n != 0 && n >= p.boundary_n) + p.pool_tiles[1].push_back(p.boundary_row + n * mblocks); + for (int mrow = p.boundary_row + (p.boundary_n != 0); mrow < mblocks; + ++mrow) + p.pool_tiles[1].push_back(mrow + n * mblocks); + } + if (int(p.pool_tiles[0].size()) != p.side_work[0] || + int(p.pool_tiles[1].size()) != p.side_work[1]) { + std::fprintf(stderr, "owned pool-size mismatch\n"); + std::abort(); + } + *out = std::move(p); + return true; +} + +// Fill one side of the table from its pool, preserving every cluster's trip +// count exactly (clusters sorted by trips, pool quotas by rank). +inline void assign_owned_pool(std::vector& table, int side, + const Census& census, + const std::vector& trips, + const std::vector& pool) { + const int clusters = census.clusters; + const int total_work = int(table.size()); + const int count = int(census.clusters_by_side[side].size()); + const int work_count = int(pool.size()); + std::vector ranks(count); + std::iota(ranks.begin(), ranks.end(), 0); + std::sort(ranks.begin(), ranks.end(), [&](int a, int b) { + const int na = (work_count + count - 1 - a) / count; + const int nb = (work_count + count - 1 - b) / count; + return na != nb ? na > nb : a < b; + }); + std::vector ids = census.clusters_by_side[side]; + std::sort(ids.begin(), ids.end(), [&](int a, int b) { + return trips[a] != trips[b] ? trips[a] > trips[b] : a < b; + }); + for (int i = 0; i < count; ++i) { + const int cluster = ids[i]; + const int rank = ranks[i]; + const int need = (work_count + count - 1 - rank) / count; + if (trips[cluster] != need) { + std::fprintf(stderr, "cannot preserve cluster trip count\n"); + std::abort(); + } + int trip = 0; + for (int work = cluster; work < total_work; + work += clusters, ++trip) { + const int pool_index = rank + trip * count; + if (pool_index >= work_count) std::abort(); + table[work] = pool[pool_index]; + } + } +} + +// Exactly-once coverage (every table, every time). +inline bool coverage_ok(const std::vector& t, int total_work) { + std::vector seen(total_work, 0); + for (int w = 0; w < total_work; ++w) + if (t[w] < 0 || t[w] >= total_work || seen[t[w]]++) return false; + return true; +} + +inline bool verify_identity(const std::vector& table) { + if (!coverage_ok(table, int(table.size()))) return false; + for (int w = 0; w < int(table.size()); ++w) + if (table[w] != w) return false; + return true; +} + +// The ownership contract: exactly-once coverage; every M row's requesters +// side-pure (split at the boundary row); every N column visited by both +// sides (B stays shared by everyone). +inline bool verify_owned(const std::vector& table, const Census& census, + int mblocks, int nblocks, const OwnedPartition& part) { + const int total_work = mblocks * nblocks; + std::vector visits(total_work, 0); + std::vector m_masks(mblocks, 0), n_masks(nblocks, 0); + for (int cluster = 0; cluster < census.clusters; ++cluster) { + const unsigned side_bit = 1u << census.cluster_side[cluster]; + for (int work = cluster; work < total_work; work += census.clusters) { + const int tile = table[work]; + if (tile < 0 || tile >= total_work || visits[tile]++) return false; + m_masks[tile % mblocks] |= side_bit; + n_masks[tile / mblocks] |= side_bit; + } + } + if (std::find(visits.begin(), visits.end(), 0) != visits.end()) + return false; + for (int mrow = 0; mrow < mblocks; ++mrow) { + unsigned want = mrow < part.boundary_row ? 1u : 2u; + if (part.boundary_n != 0 && mrow == part.boundary_row) want = 3u; + if (m_masks[mrow] != want) return false; + } + return std::find_if(n_masks.begin(), n_masks.end(), + [](unsigned mask) { return mask != 3u; }) == + n_masks.end(); +} + +// Matched ownership control: every A row must be requested from both sides, +// while every B column remains shared exactly as in the owned table. +inline bool verify_blind(const std::vector& table, const Census& census, + int mblocks, int nblocks) { + const int total_work = mblocks * nblocks; + std::vector visits(total_work, 0); + std::vector m_masks(mblocks, 0), n_masks(nblocks, 0); + for (int cluster = 0; cluster < census.clusters; ++cluster) { + const unsigned side_bit = 1u << census.cluster_side[cluster]; + for (int work = cluster; work < total_work; work += census.clusters) { + const int tile = table[work]; + if (tile < 0 || tile >= total_work || visits[tile]++) return false; + m_masks[tile % mblocks] |= side_bit; + n_masks[tile / mblocks] |= side_bit; + } + } + if (std::find(visits.begin(), visits.end(), 0) != visits.end()) + return false; + return std::find_if(m_masks.begin(), m_masks.end(), + [](unsigned mask) { return mask != 3u; }) == + m_masks.end() && + std::find_if(n_masks.begin(), n_masks.end(), + [](unsigned mask) { return mask != 3u; }) == + n_masks.end(); +} + +// Derive a blind twin from an owned table. Rotating M by an N-dependent +// amount preserves each work slot's N coordinate, so the B/SFB request stream, +// cluster trip counts, kernel instructions, and table lookup are unchanged. +inline std::vector build_blind_twin(const std::vector& owned, + const Census& census, int mblocks, + int nblocks) { + const int total_work = mblocks * nblocks; + std::vector blind(total_work, -1); + for (int stride = 1; stride < mblocks; ++stride) { + for (int work = 0; work < total_work; ++work) { + const int tile = owned[work]; + const int mrow = tile % mblocks; + const int n = tile / mblocks; + blind[work] = (mrow + n * stride) % mblocks + n * mblocks; + } + if (!verify_blind(blind, census, mblocks, nblocks)) continue; + for (int work = 0; work < total_work; ++work) + if (blind[work] / mblocks != owned[work] / mblocks) + std::abort(); + std::printf("SCHEDULE_BLIND stride=%d same_n_slot=PASS\n", stride); + return blind; + } + std::fprintf(stderr, "found no row rotation for a blind schedule twin\n"); + std::abort(); +} + +// ---- visit orders (blog: the L2 trick, part two) -------------------------------------------- + +// Standard iterative Hilbert d2xy over an npow2 x npow2 grid. +inline void hilbert_d2xy(int npow2, int d, int& x, int& y) { + x = 0; + y = 0; + for (int s = 1; s < npow2; s *= 2) { + const int rx = 1 & (d / 2); + const int ry = 1 & (d ^ rx); + if (ry == 0) { + if (rx == 1) { + x = s - 1 - x; + y = s - 1 - y; + } + const int t = x; + x = y; + y = t; + } + x += s * rx; + y += s * ry; + d /= 4; + } +} + +// Hilbert order over a rows x cols sub-grid whose first M row is row0; the +// curve covers the next power of two and out-of-grid points are skipped. +inline std::vector hilbert_order(int rows, int cols, int row0, int mblocks) { + int npow2 = 1; + while (npow2 < rows || npow2 < cols) npow2 *= 2; + std::vector order; + order.reserve(size_t(rows) * cols); + for (int d = 0; d < npow2 * npow2; ++d) { + int x, y; + hilbert_d2xy(npow2, d, x, y); + if (x < rows && y < cols) + order.push_back(row0 + x + y * mblocks); + } + return order; +} + +// Hilbert WITHIN each side's owned row range: the curve provides locality, +// the pools keep every A re-read on the requester's own side. This is the +// order that beats every other reorder at every deep-K shape in the post. +inline std::array, 2> hilbertown_pools( + const OwnedPartition& part, int mblocks, int nblocks) { + const int rows0 = part.boundary_row + (part.boundary_n != 0); + std::array, 2> hpool; + for (const int tile : hilbert_order(rows0, nblocks, 0, mblocks)) { + const int mrow = tile % mblocks, n = tile / mblocks; + if (mrow < part.boundary_row || + (part.boundary_n != 0 && n < part.boundary_n)) + hpool[0].push_back(tile); + } + for (const int tile : hilbert_order(mblocks - part.boundary_row, nblocks, + part.boundary_row, mblocks)) { + const int mrow = tile % mblocks, n = tile / mblocks; + if (mrow > part.boundary_row || part.boundary_n == 0 || + n >= part.boundary_n) + hpool[1].push_back(tile); + } + if (int(hpool[0].size()) != part.side_work[0] || + int(hpool[1].size()) != part.side_work[1]) { + std::fprintf(stderr, "hilbertown pool-size mismatch\n"); + std::abort(); + } + return hpool; +} + +// Pocket order: walk the tile grid in R x C rectangles, finishing each +// pocket before the next, pocket bands serpentine. 8x4 won the shape sweep +// on big squares (and bought the ~10% traffic cut the power governor repaid +// in clock — the blog's strangest section). +inline std::vector pocket_order(int R, int C, int mblocks, int nblocks) { + std::vector order; + order.reserve(size_t(mblocks) * nblocks); + const int prows = (mblocks + R - 1) / R; + const int pcols = (nblocks + C - 1) / C; + for (int pi = 0; pi < prows; ++pi) + for (int t = 0; t < pcols; ++t) { + const int pj = pi % 2 ? pcols - 1 - t : t; + const int m1 = std::min((pi + 1) * R, mblocks); + const int n1 = std::min((pj + 1) * C, nblocks); + for (int m = pi * R; m < m1; ++m) + for (int n = pj * C; n < n1; ++n) + order.push_back(m + n * mblocks); + } + return order; +} + +// Split a full-grid visit order into the two owned pools, preserving the +// per-side order — the pocket/side composition step. +inline std::array, 2> side_filter( + const std::vector& order, const OwnedPartition& part, int mblocks) { + std::array, 2> pools; + for (const int tile : order) { + const int mrow = tile % mblocks, n = tile / mblocks; + const int side = mrow < part.boundary_row ? 0 + : mrow > part.boundary_row ? 1 + : (part.boundary_n != 0 && n < part.boundary_n) ? 0 : 1; + pools[side].push_back(tile); + } + if (int(pools[0].size()) != part.side_work[0] || + int(pools[1].size()) != part.side_work[1]) { + std::fprintf(stderr, "side-filter pool-size mismatch\n"); + std::abort(); + } + return pools; +} + +enum class ScheduleMode { + AUTO, + RASTER, + POCKET_8x4, + POCKET_8x8, + HILBERT, + BLIND_PLAIN, + BLIND_POCKET_8x4, + BLIND_POCKET_8x8, + HILBERT_IN_BLIND, + OWNED_PLAIN, + OWNED_POCKET_8x4, + OWNED_POCKET_8x8, + HILBERT_IN_OWNED, +}; + +inline const char* schedule_name(ScheduleMode mode) { + switch (mode) { + case ScheduleMode::AUTO: return "auto"; + case ScheduleMode::RASTER: return "raster"; + case ScheduleMode::POCKET_8x4: return "pocket-8x4"; + case ScheduleMode::POCKET_8x8: return "pocket-8x8"; + case ScheduleMode::HILBERT: return "hilbert"; + case ScheduleMode::BLIND_PLAIN: return "blind-plain"; + case ScheduleMode::BLIND_POCKET_8x4: return "blind-pocket-8x4"; + case ScheduleMode::BLIND_POCKET_8x8: return "blind-pocket-8x8"; + case ScheduleMode::HILBERT_IN_BLIND: return "hilbert-in-blind"; + case ScheduleMode::OWNED_PLAIN: return "owned-plain"; + case ScheduleMode::OWNED_POCKET_8x4: return "owned-pocket-8x4"; + case ScheduleMode::OWNED_POCKET_8x8: return "owned-pocket-8x8"; + case ScheduleMode::HILBERT_IN_OWNED: return "hilbert-in-owned"; + } + return "?"; +} + +inline bool parse_schedule_mode(const std::string& name, ScheduleMode* mode) { + static const std::array modes = { + ScheduleMode::AUTO, + ScheduleMode::RASTER, + ScheduleMode::POCKET_8x4, + ScheduleMode::POCKET_8x8, + ScheduleMode::HILBERT, + ScheduleMode::BLIND_PLAIN, + ScheduleMode::BLIND_POCKET_8x4, + ScheduleMode::BLIND_POCKET_8x8, + ScheduleMode::HILBERT_IN_BLIND, + ScheduleMode::OWNED_PLAIN, + ScheduleMode::OWNED_POCKET_8x4, + ScheduleMode::OWNED_POCKET_8x8, + ScheduleMode::HILBERT_IN_OWNED, + }; + for (const ScheduleMode candidate : modes) { + if (name == schedule_name(candidate)) { + *mode = candidate; + return true; + } + } + return false; +} + +inline bool uses_side_ownership(ScheduleMode mode) { + return mode == ScheduleMode::OWNED_PLAIN || + mode == ScheduleMode::OWNED_POCKET_8x4 || + mode == ScheduleMode::OWNED_POCKET_8x8 || + mode == ScheduleMode::HILBERT_IN_OWNED; +} + +inline bool uses_blind_twin(ScheduleMode mode) { + return mode == ScheduleMode::BLIND_PLAIN || + mode == ScheduleMode::BLIND_POCKET_8x4 || + mode == ScheduleMode::BLIND_POCKET_8x8 || + mode == ScheduleMode::HILBERT_IN_BLIND; +} + +inline bool requires_side_census(ScheduleMode mode) { + return uses_side_ownership(mode) || uses_blind_twin(mode); +} + +inline const char* requester_assignment(ScheduleMode mode) { + return uses_side_ownership(mode) ? "owned" + : uses_blind_twin(mode) ? "blind" + : "none"; +} + +// Own the sides always. The 32x32 grid uses the measured 8x8 winner. Other +// grids below the reorder knee keep the plain owned order; past it, pick by +// aspect. The knee sits between the measured ~99 MB still-loses and ~159 MB +// wins points; 128 MiB (both L2 sides) is the threshold this file ships. +inline ScheduleMode pick_schedule(size_t operand_read_bytes, int mblocks, + int nblocks) { + constexpr size_t REORDER_KNEE = size_t(128) << 20; + if (operand_read_bytes < REORDER_KNEE) { + if (mblocks == 32 && nblocks == 32) + return ScheduleMode::OWNED_POCKET_8x8; + return ScheduleMode::OWNED_PLAIN; + } + const int lo = std::min(mblocks, nblocks); + const int hi = std::max(mblocks, nblocks); + if (lo >= 48 && hi * 4 <= lo * 5) + return ScheduleMode::OWNED_POCKET_8x4; // aspect <= 1.25 + return ScheduleMode::HILBERT_IN_OWNED; +} + +// Full-grid controls: these change only tile visit order. They deliberately do +// not probe, classify, or partition the two L2 sides. +inline std::vector build_unowned_schedule(ScheduleMode mode, int mblocks, + int nblocks) { + const int total_work = mblocks * nblocks; + std::vector table; + switch (mode) { + case ScheduleMode::RASTER: + table.resize(total_work); + std::iota(table.begin(), table.end(), 0); + break; + case ScheduleMode::POCKET_8x4: + table = pocket_order(8, 4, mblocks, nblocks); + break; + case ScheduleMode::POCKET_8x8: + table = pocket_order(8, 8, mblocks, nblocks); + break; + case ScheduleMode::HILBERT: + table = hilbert_order(mblocks, nblocks, 0, mblocks); + break; + default: + std::abort(); + } + if (!coverage_ok(table, total_work)) { + std::fprintf(stderr, "unowned schedule failed exactly-once coverage\n"); + std::abort(); + } + return table; +} + +// Build + prove one schedule table. Returns empty on an unbuildable +// partition (the caller falls back to natural raster and says so). +inline std::vector build_schedule(const Census& census, int mblocks, + int nblocks, + ScheduleMode mode + ) { + const int total_work = mblocks * nblocks; + const std::vector trips = cluster_trips(census.clusters, total_work); + OwnedPartition part; + if (!make_owned_partition(census, trips, mblocks, nblocks, &part)) + return {}; + std::array, 2> pools; + switch (mode) { + case ScheduleMode::BLIND_PLAIN: + case ScheduleMode::OWNED_PLAIN: + pools = part.pool_tiles; + break; + case ScheduleMode::BLIND_POCKET_8x4: + case ScheduleMode::OWNED_POCKET_8x4: + pools = side_filter(pocket_order(8, 4, mblocks, nblocks), part, + mblocks); + break; + case ScheduleMode::BLIND_POCKET_8x8: + case ScheduleMode::OWNED_POCKET_8x8: + pools = side_filter(pocket_order(8, 8, mblocks, nblocks), part, + mblocks); + break; + case ScheduleMode::HILBERT_IN_BLIND: + case ScheduleMode::HILBERT_IN_OWNED: + pools = hilbertown_pools(part, mblocks, nblocks); + break; + default: + std::abort(); + } + std::vector table(total_work, -1); + assign_owned_pool(table, 0, census, trips, pools[0]); + assign_owned_pool(table, 1, census, trips, pools[1]); + if (!verify_owned(table, census, mblocks, nblocks, part)) { + std::fprintf(stderr, "schedule table failed the ownership contract\n"); + std::abort(); + } + if (uses_blind_twin(mode)) + return build_blind_twin(table, census, mblocks, nblocks); + return table; +} + +// Upload one table slot; read it back and require equality. +inline void upload_table(const std::vector& table, int mode) { + const int total_work = int(table.size()); + if (total_work > L2A_ROUTE_WORK_CAP) std::abort(); + const size_t offset = size_t(mode) * L2A_ROUTE_WORK_CAP * sizeof(int); + CUDA_CHECK(cudaMemcpyToSymbol(nvfp4::l2a_route_tables, table.data(), + total_work * sizeof(int), offset, + cudaMemcpyHostToDevice)); + std::vector roundtrip(total_work, -1); + CUDA_CHECK(cudaMemcpyFromSymbol(roundtrip.data(), nvfp4::l2a_route_tables, + total_work * sizeof(int), offset, + cudaMemcpyDeviceToHost)); + if (roundtrip != table) { + std::fprintf(stderr, "constant table roundtrip failed mode=%d\n", mode); + std::abort(); + } +} + +inline void upload_side_constants(const l2side::RuntimeMap& map, + const Census& census) { + CUDA_CHECK(cudaMemcpyToSymbol(nvfp4::l2a_sm_side, map.sm_side.data(), + sizeof(nvfp4::l2a_sm_side))); + CUDA_CHECK(cudaMemcpyToSymbol(nvfp4::l2a_cluster_side, + census.cluster_side.data(), + census.clusters * sizeof(int))); + const unsigned zero = 0; + CUDA_CHECK(cudaMemcpyToSymbol(nvfp4::l2a_placement_errors, &zero, + sizeof(zero))); +} + +// The fallback: identity tables, all sides zeroed (the placement audit is +// inert — every SM trivially matches side 0). +inline void upload_fallback_sides() { + static const std::array zeros_sm{}; + static const std::array zeros_cl{}; + CUDA_CHECK(cudaMemcpyToSymbol(nvfp4::l2a_sm_side, zeros_sm.data(), + sizeof(nvfp4::l2a_sm_side))); + CUDA_CHECK(cudaMemcpyToSymbol(nvfp4::l2a_cluster_side, zeros_cl.data(), + sizeof(nvfp4::l2a_cluster_side))); + const unsigned zero = 0; + CUDA_CHECK(cudaMemcpyToSymbol(nvfp4::l2a_placement_errors, &zero, + sizeof(zero))); +} + +} // namespace sched diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/main.cu b/code/labs/fast_cu/upstream/gb300/nvfp4/main.cu new file mode 100644 index 000000000..a16fa11d9 --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/main.cu @@ -0,0 +1,1895 @@ +#ifndef NVFP4_GEMM_HEADER +#define NVFP4_GEMM_HEADER "gemm9.cuh" +#endif + +#include NVFP4_GEMM_HEADER + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUBLAS_CHECK(expr) do { \ + cublasStatus_t _s = (expr); \ + if (_s != CUBLAS_STATUS_SUCCESS) { \ + std::fprintf(stderr, "cuBLASLt error %d at %s:%d\n", \ + int(_s), __FILE__, __LINE__); \ + std::abort(); \ + } \ +} while (0) + +// Shared correctness, cuBLASLt, and benchmark harness. The selected +// gemmN.cuh above contains only the implementation for that rung. +namespace host { + +// ---- NVFP4 quantization ----------------------------------------------------- +// The format: 4-bit E2M1 values, one UE4M3 scale per 16 consecutive K +// elements. Scale = E4M3-quantized amax/6 (NOT power-of-two rounded). + +constexpr float E2M1_MAX = 6.0f; + +inline float fp4_e2m1_decode(uint8_t nibble) { + static constexpr float lut[16] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f, + }; + return lut[nibble & 0xF]; +} + +// Round-to-nearest-EVEN, matching the device-side cvt.rn.satfinite.e2m1x2: +// at each midpoint pick the code whose LSB is 0. +inline uint8_t fp4_e2m1_quantize(float x) { + const float boundaries[7] = { 0.25f, 0.75f, 1.25f, 1.75f, 2.5f, 3.5f, 5.0f }; + float ax = std::fabs(x); + if (ax > E2M1_MAX) ax = E2M1_MAX; + int idx = 0; + for (int b = 0; b < 7; ++b) { + if (ax > boundaries[b]) { + ++idx; + } else if (ax == boundaries[b]) { + if ((b & 1) == 1) ++idx; // tie: round to the even code + break; + } else { + break; + } + } + uint8_t code = uint8_t(idx); + if (x < 0 && idx != 0) code |= 0x8; + return code; +} + +inline uint8_t fp8_e4m3_quantize(float x) { + if (!std::isfinite(x)) return 0x7F; + if (x > 448.0f) x = 448.0f; + if (x < -448.0f) x = -448.0f; + uint32_t bits; + std::memcpy(&bits, &x, 4); + const uint32_t sign = (bits >> 31) & 0x1; + const int32_t exp = int32_t((bits >> 23) & 0xFF) - 127; + const uint32_t man = bits & 0x7FFFFF; + if (x == 0.0f) return uint8_t(sign << 7); + if (exp < -9) return uint8_t(sign << 7); + int32_t e = exp + 7; + uint32_t m; + if (e <= 0) { + const uint32_t shift = uint32_t(1 - e); + if (shift > 23) return uint8_t(sign << 7); + m = ((man | (1u << 23)) + (1u << (shift + 19))) >> (shift + 20); + e = 0; + } else { + m = (man + (1u << 19)) >> 20; + if (m >= 8) { m = 0; ++e; } + } + if (e >= 16) { e = 15; m = 6; } // saturate at 448 + return uint8_t((sign << 7) | (uint32_t(e) << 3) | (m & 0x7)); +} + +inline float fp8_e4m3_decode(uint8_t code) { + if ((code & 0x7F) == 0x7F) return std::nanf(""); + const uint32_t sign = (code >> 7) & 0x1; + const uint32_t exp = (code >> 3) & 0xF; + const uint32_t man = code & 0x7; + float v; + if (exp == 0) v = std::ldexp(float(man), -9); + else v = std::ldexp(1.0f + float(man) / 8.0f, int(exp) - 7); + return sign ? -v : v; +} + +inline float ue4m3_byte_to_sf(uint8_t code) { + return fp8_e4m3_decode(uint8_t(code & 0x7Fu)); +} + +// (M, K) float row-major -> packed E2M1 [M][K/2] + UE4M3 scales [M][K/16]. +inline void cast_to_fp4_with_ue4m3(const float* x, int M, int K, + uint8_t* fp4_out, uint8_t* sf_out) { + const int n_chunks = K / 16; + for (int m = 0; m < M; ++m) { + for (int g = 0; g < n_chunks; ++g) { + float amax = 1e-4f; + for (int j = 0; j < 16; ++j) + amax = std::max(amax, std::fabs(x[size_t(m) * K + g * 16 + j])); + const uint8_t sf_byte = + uint8_t(fp8_e4m3_quantize(amax / E2M1_MAX) & 0x7Fu); + sf_out[size_t(m) * n_chunks + g] = sf_byte; + const float sf_rec = ue4m3_byte_to_sf(sf_byte); + const float inv_sf = sf_rec > 0.0f ? 1.0f / sf_rec : 0.0f; + for (int j = 0; j < 16; j += 2) { + const uint8_t lo = + fp4_e2m1_quantize(x[size_t(m) * K + g * 16 + j] * inv_sf); + const uint8_t hi = + fp4_e2m1_quantize(x[size_t(m) * K + g * 16 + j + 1] * inv_sf); + fp4_out[size_t(m) * (K / 2) + (g * 16 + j) / 2] = + uint8_t((lo & 0xF) | ((hi & 0xF) << 4)); + } + } + } +} + +// Dequantize one element back to float (the host reference reads this). +inline float dequant_nvfp4(const uint8_t* fp4, const uint8_t* sf, + int m, int k, int K) { + const float sf_v = ue4m3_byte_to_sf(sf[size_t(m) * (K / 16) + k / 16]); + const uint8_t byte = fp4[size_t(m) * (K / 2) + k / 2]; + const uint8_t nib = (k & 1) ? uint8_t(byte >> 4) : uint8_t(byte & 0xF); + return fp4_e2m1_decode(nib) * sf_v; +} + +// ---- shared input buffers -------------------------------------------------- +// A/B use the row-major packed-E2M1 bytes accepted by cuBLASLt, and both +// implementations receive the same standard VEC16_UE4M3 scale buffers. There +// is no kernel-specific repack. + +constexpr size_t ALLOCATION_GUARD = 4096; + +inline size_t vec16_sf_offset(int outer, int sf_inner, + int inner, int outer_dim) { + const int sf_inner_dim = int((size_t(inner) + 63) / 64 * 4); + const int outer_blocks = int((size_t(outer_dim) + 127) / 128); + const int outer_block = outer / 128; + if (outer_block >= outer_blocks || sf_inner >= sf_inner_dim) std::abort(); + const int inner_block = (sf_inner / 4) * 4; + const size_t block = + size_t(inner_block + outer_block * sf_inner_dim) * 128; + const int local_outer = outer % 128; + return block + size_t(local_outer % 32) * 16 + + size_t(local_outer / 32) * 4 + size_t(sf_inner % 4); +} + +inline std::vector pack_vec16_scales( + const std::vector& logical, int outer, int inner) { + const int logical_inner = inner / 16; + if (inner % 16 != 0 || + logical.size() != size_t(outer) * logical_inner) std::abort(); + + const int packed_inner = int((size_t(inner) + 63) / 64 * 4); + const int packed_outer = int((size_t(outer) + 127) / 128 * 128); + std::vector packed(size_t(packed_inner) * packed_outer, 0); + std::vector seen(packed.size(), 0); + for (int o = 0; o < outer; ++o) { + for (int j = 0; j < logical_inner; ++j) { + const size_t offset = vec16_sf_offset(o, j, inner, outer); + if (offset >= packed.size() || seen[offset]) std::abort(); + seen[offset] = 1; + packed[offset] = logical[size_t(o) * logical_inner + j]; + } + } + for (int o = 0; o < outer; ++o) + for (int j = 0; j < logical_inner; ++j) + if (packed[vec16_sf_offset(o, j, inner, outer)] != + logical[size_t(o) * logical_inner + j]) std::abort(); + return packed; +} + +struct InputBuffers { + std::vector A, B, SFA, SFB; + int K = 0; + int ab_row_stride = 0; +}; + +inline InputBuffers make_input_buffers( + const std::vector& A_logical, + const std::vector& B_logical, + const std::vector& SFA_logical, + const std::vector& SFB_logical, + int M, int N, int K) { + if (K <= 0 || K % 16 != 0) std::abort(); + + InputBuffers inputs; + inputs.K = K; + const int row_bytes = K / 2; + inputs.ab_row_stride = (row_bytes + 15) & ~15; + + auto materialize_rows = [&](const std::vector& logical, int rows) { + if (logical.size() != size_t(rows) * row_bytes) std::abort(); + std::vector physical( + size_t(rows) * inputs.ab_row_stride, 0); + for (int row = 0; row < rows; ++row) + std::memcpy(&physical[size_t(row) * inputs.ab_row_stride], + &logical[size_t(row) * row_bytes], row_bytes); + return physical; + }; + + // Public benchmark shapes require K%32, so the aligned TMA row stride adds + // no bytes and A/B remain exactly the cuBLASLt buffers. The private + // K%16-only correctness gate may add zero stride slack. + inputs.A = materialize_rows(A_logical, M); + inputs.B = materialize_rows(B_logical, N); + inputs.SFA = pack_vec16_scales(SFA_logical, M, K); + inputs.SFB = pack_vec16_scales(SFB_logical, N, K); + return inputs; +} +// ---- device-resident operand set (one rotation slot) ---------------------- + +struct DeviceSlot { + int M = 0, N = 0, K = 0; + uint8_t *dA = nullptr, *dB = nullptr, *dSFA = nullptr, *dSFB = nullptr; + __half* dC = nullptr; + size_t a_sz = 0, b_sz = 0, sfa_sz = 0, sfb_sz = 0; + CUtensorMap A_t, B_t, SFA_t, SFB_t; +}; + +inline DeviceSlot* make_device_slot(const InputBuffers& inputs, int M, int N) { + auto* slot = new DeviceSlot{}; + slot->M = M; + slot->N = N; + slot->K = inputs.K; + slot->a_sz = inputs.A.size(); + slot->b_sz = inputs.B.size(); + slot->sfa_sz = inputs.SFA.size(); + slot->sfb_sz = inputs.SFB.size(); + + CUDA_CHECK(cudaMalloc(&slot->dA, slot->a_sz + ALLOCATION_GUARD)); + CUDA_CHECK(cudaMalloc(&slot->dB, slot->b_sz + ALLOCATION_GUARD)); + CUDA_CHECK(cudaMalloc(&slot->dSFA, slot->sfa_sz + ALLOCATION_GUARD)); + CUDA_CHECK(cudaMalloc(&slot->dSFB, slot->sfb_sz + ALLOCATION_GUARD)); + CUDA_CHECK(cudaMalloc(&slot->dC, size_t(M) * N * sizeof(__half))); + CUDA_CHECK(cudaMemcpy(slot->dA, inputs.A.data(), slot->a_sz, + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(slot->dB, inputs.B.data(), slot->b_sz, + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(slot->dSFA, inputs.SFA.data(), slot->sfa_sz, + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(slot->dSFB, inputs.SFB.data(), slot->sfb_sz, + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemset(slot->dC, 0, size_t(M) * N * sizeof(__half))); + + // Poisoning the guard catches any TMA read beyond the declared buffers. + CUDA_CHECK(cudaMemset(slot->dA + slot->a_sz, 0xA5, ALLOCATION_GUARD)); + CUDA_CHECK(cudaMemset(slot->dB + slot->b_sz, 0xA5, ALLOCATION_GUARD)); + CUDA_CHECK(cudaMemset(slot->dSFA + slot->sfa_sz, 0xA5, ALLOCATION_GUARD)); + CUDA_CHECK(cudaMemset(slot->dSFB + slot->sfb_sz, 0xA5, ALLOCATION_GUARD)); + + slot->A_t = make_ab_tmap( + slot->dA, M, inputs.K, inputs.ab_row_stride); + slot->B_t = make_ab_tmap( + slot->dB, N, inputs.K, inputs.ab_row_stride); + slot->SFA_t = make_sf_tmap(slot->dSFA, M, inputs.K); + slot->SFB_t = make_sf_tmap(slot->dSFB, N, inputs.K); + return slot; +} + +inline void free_device_slot(DeviceSlot* slot) { + cudaFree(slot->dC); + cudaFree(slot->dSFB); + cudaFree(slot->dSFA); + cudaFree(slot->dB); + cudaFree(slot->dA); + delete slot; +} +// ---- the input fixture ----------------------------------------------------- + +struct Fixture { + int M = 0, N = 0, K = 0; + std::vector A_logical, B_logical; + std::vector SFA_logical, SFB_logical; + std::vector A_deq, B_deq; + InputBuffers inputs; +}; + +inline Fixture make_fixture(int M, int N, int K, uint64_t seed, bool keep_deq) { + Fixture fixture; + fixture.M = M; + fixture.N = N; + fixture.K = K; + + std::mt19937_64 rng(seed); + std::normal_distribution dist(0.0f, 1.0f); + std::vector A_src(size_t(M) * K), B_src(size_t(N) * K); + for (auto& value : A_src) value = dist(rng); + for (auto& value : B_src) value = dist(rng); + + fixture.A_logical.resize(size_t(M) * (K / 2)); + fixture.B_logical.resize(size_t(N) * (K / 2)); + fixture.SFA_logical.resize(size_t(M) * (K / 16)); + fixture.SFB_logical.resize(size_t(N) * (K / 16)); + cast_to_fp4_with_ue4m3(A_src.data(), M, K, + fixture.A_logical.data(), + fixture.SFA_logical.data()); + cast_to_fp4_with_ue4m3(B_src.data(), N, K, + fixture.B_logical.data(), + fixture.SFB_logical.data()); + + if (keep_deq) { + fixture.A_deq.resize(size_t(M) * K); + fixture.B_deq.resize(size_t(N) * K); + for (int m = 0; m < M; ++m) + for (int k = 0; k < K; ++k) + fixture.A_deq[size_t(m) * K + k] = dequant_nvfp4( + fixture.A_logical.data(), + fixture.SFA_logical.data(), m, k, K); + for (int n = 0; n < N; ++n) + for (int k = 0; k < K; ++k) + fixture.B_deq[size_t(n) * K + k] = dequant_nvfp4( + fixture.B_logical.data(), + fixture.SFB_logical.data(), n, k, K); + } + + fixture.inputs = make_input_buffers( + fixture.A_logical, fixture.B_logical, + fixture.SFA_logical, fixture.SFB_logical, M, N, K); + return fixture; +} +} // namespace host + +// ============================================================================ +// The cuBLASLt baseline. Request one recommendation and use it for correctness +// and timing at every shape. +// ============================================================================ +namespace bench { + +constexpr size_t MIB = size_t(1) << 20; +constexpr size_t CUBLAS_WORKSPACE_LIMIT = 64 * MIB; +constexpr int SACRIFICIAL_PASSES = 2; +constexpr int TIMED_PASSES = 5; +// FP16 round-to-nearest is <= 2^-11 relative; allow four ulps plus 0.5 for +// reduction-order differences in the FP32 accumulation. +constexpr double FP16_ABS_TOL = 0.5; +constexpr double FP16_REL_TOL = 2e-3; + +struct LtSlot { + void *a = nullptr, *b = nullptr, *c = nullptr, *d = nullptr; + void *sfa = nullptr, *sfb = nullptr, *workspace = nullptr; + cublasLtMatmulDesc_t op = nullptr; +}; + +struct LtPlan { + int M, N, K; + cublasLtHandle_t handle = nullptr; + cublasLtMatrixLayout_t a_layout = nullptr, b_layout = nullptr; + cublasLtMatrixLayout_t out_layout = nullptr; + cublasLtMatmulPreference_t preference = nullptr; + cublasLtMatmulAlgo_t algo{}; + size_t workspace_bytes = 0; + + LtPlan(int m, int n, int k) : M(m), N(n), K(k) { + CUBLAS_CHECK(cublasLtCreate(&handle)); + CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&a_layout, CUDA_R_4F_E2M1, K, M, K)); + CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&b_layout, CUDA_R_4F_E2M1, K, N, K)); + const int initial_ld = std::max(M, N); + CUBLAS_CHECK(cublasLtMatrixLayoutCreate(&out_layout, CUDA_R_16F, M, N, + initial_ld)); + cublasLtOrder_t row = CUBLASLT_ORDER_ROW; + int64_t ld = N; + CUBLAS_CHECK(cublasLtMatrixLayoutSetAttribute( + out_layout, CUBLASLT_MATRIX_LAYOUT_ORDER, &row, sizeof(row))); + CUBLAS_CHECK(cublasLtMatrixLayoutSetAttribute( + out_layout, CUBLASLT_MATRIX_LAYOUT_LD, &ld, sizeof(ld))); + CUBLAS_CHECK(cublasLtMatmulPreferenceCreate(&preference)); + CUBLAS_CHECK(cublasLtMatmulPreferenceSetAttribute( + preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &CUBLAS_WORKSPACE_LIMIT, sizeof(CUBLAS_WORKSPACE_LIMIT))); + } + + ~LtPlan() { + cublasLtMatmulPreferenceDestroy(preference); + cublasLtMatrixLayoutDestroy(out_layout); + cublasLtMatrixLayoutDestroy(b_layout); + cublasLtMatrixLayoutDestroy(a_layout); + cublasLtDestroy(handle); + } + + LtSlot make_slot(const host::Fixture& f) { + const host::InputBuffers& inputs = f.inputs; + LtSlot slot; + CUDA_CHECK(cudaMalloc(&slot.a, inputs.A.size())); + CUDA_CHECK(cudaMalloc(&slot.b, inputs.B.size())); + CUDA_CHECK(cudaMalloc(&slot.sfa, inputs.SFA.size())); + CUDA_CHECK(cudaMalloc(&slot.sfb, inputs.SFB.size())); + CUDA_CHECK(cudaMalloc(&slot.d, size_t(M) * N * sizeof(__half))); + slot.c = slot.d; // beta = 0 permits C == D + CUDA_CHECK(cudaMemcpy(slot.a, inputs.A.data(), inputs.A.size(), + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(slot.b, inputs.B.data(), inputs.B.size(), + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(slot.sfa, inputs.SFA.data(), inputs.SFA.size(), + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(slot.sfb, inputs.SFB.data(), inputs.SFB.size(), + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemset(slot.d, 0, size_t(M) * N * sizeof(__half))); + + CUBLAS_CHECK(cublasLtMatmulDescCreate(&slot.op, CUBLAS_COMPUTE_32F, + CUDA_R_32F)); + cublasOperation_t transa = CUBLAS_OP_T, transb = CUBLAS_OP_N; + int32_t sf_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; + int8_t fast_accum = 0; + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + slot.op, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + slot.op, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transb))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + slot.op, CUBLASLT_MATMUL_DESC_FAST_ACCUM, &fast_accum, + sizeof(fast_accum))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + slot.op, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, &sf_mode, sizeof(sf_mode))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + slot.op, CUBLASLT_MATMUL_DESC_B_SCALE_MODE, &sf_mode, sizeof(sf_mode))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + slot.op, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &slot.sfa, + sizeof(slot.sfa))); + CUBLAS_CHECK(cublasLtMatmulDescSetAttribute( + slot.op, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &slot.sfb, + sizeof(slot.sfb))); + return slot; + } + + void initialize_algorithm(const LtSlot& slot) { + cublasLtMatmulHeuristicResult_t first{}; + int returned = 0; + CUBLAS_CHECK(cublasLtMatmulAlgoGetHeuristic( + handle, slot.op, a_layout, b_layout, out_layout, out_layout, + preference, 1, &first, &returned)); + if (returned != 1 || first.state != CUBLAS_STATUS_SUCCESS || + first.workspaceSize > CUBLAS_WORKSPACE_LIMIT) { + std::fprintf(stderr, + "cuBLASLt returned no usable first algorithm\n"); + std::abort(); + } + algo = first.algo; + workspace_bytes = first.workspaceSize; + std::printf("CUBLAS_ALGO index=0 workspace=%zu\n", workspace_bytes); + } + + void add_workspace(LtSlot& slot) const { + if (workspace_bytes) + CUDA_CHECK(cudaMalloc(&slot.workspace, workspace_bytes)); + if ((reinterpret_cast(slot.sfa) & 15u) != 0 || + (reinterpret_cast(slot.sfb) & 15u) != 0 || + (slot.workspace && + (reinterpret_cast(slot.workspace) & 255u) != 0)) { + std::fprintf(stderr, "cuBLASLt slot alignment contract failed\n"); + std::abort(); + } + } + + cublasStatus_t launch_status(const LtSlot& slot) const { + float alpha = 1.0f, beta = 0.0f; + return cublasLtMatmul( + handle, slot.op, &alpha, slot.a, a_layout, slot.b, b_layout, + &beta, slot.c, out_layout, slot.d, out_layout, &algo, + workspace_bytes ? slot.workspace : nullptr, + workspace_bytes, nullptr); + } + + void launch(const LtSlot& slot) const { CUBLAS_CHECK(launch_status(slot)); } + + void free_slot(LtSlot& slot) const { + if (slot.op) cublasLtMatmulDescDestroy(slot.op); + if (slot.workspace) cudaFree(slot.workspace); + if (slot.sfb) cudaFree(slot.sfb); + if (slot.sfa) cudaFree(slot.sfa); + if (slot.d) cudaFree(slot.d); + if (slot.b) cudaFree(slot.b); + if (slot.a) cudaFree(slot.a); + slot = {}; + } +}; + +} // namespace bench + +// ============================================================================ +// Harness: correctness gates before ANY timing, then a rotating cache-cold +// benchmark against cuBLASLt's first correctness-gated algorithm (blog: +// "How to benchmark a cache trick without lying to yourself"). +// ============================================================================ +namespace bench { + +enum class BenchMode { + CURRENT, + CLOCK_PINNED, + SUSTAINED, + TRITON, +}; + +inline const char* bench_mode_name(BenchMode mode) { + switch (mode) { + case BenchMode::CURRENT: return "current"; + case BenchMode::CLOCK_PINNED: return "clock-pinned"; + case BenchMode::SUSTAINED: return "sustained"; + case BenchMode::TRITON: return "triton"; + } + return "unknown"; +} + +inline bool parse_bench_mode(const std::string& text, BenchMode* mode) { + if (text == "current") *mode = BenchMode::CURRENT; + else if (text == "clock-pinned") *mode = BenchMode::CLOCK_PINNED; + else if (text == "sustained") *mode = BenchMode::SUSTAINED; + else if (text == "triton") *mode = BenchMode::TRITON; + else return false; + return true; +} + +struct Ctx { + int sms = 0; + int clusters = 0; + int l2_bytes = 0; + dim3 grid{}; + int cooldown_sec = 0; + int reference_clock_mhz = -1; + BenchMode mode = BenchMode::CURRENT; +}; +static Ctx g; + +constexpr int ACTIVE_IDLE_CLOCK_MAX_MHZ = 2100; +constexpr int CLOCK_MATCH_TOL_MHZ = 60; +constexpr int CLOCK_COOLDOWN_MAX_SEC = 20; +constexpr int PINNED_CLOCK_MHZ = 1305; +constexpr int PINNED_CLOCK_TOL_MHZ = 15; +constexpr size_t TRITON_FLUSH_BYTES = size_t(256) << 20; +constexpr double SUSTAINED_BLOCK_SECONDS = 60.0; +constexpr double SUSTAINED_FIRST_SECONDS = 10.0; +constexpr double SUSTAINED_TAIL_FROM_SECONDS = 20.0; + +inline const char* gemm_source_file() { +#ifdef NVFP4_GEMM_HEADER + return NVFP4_GEMM_HEADER; +#else + return "nvfp4_gemm.cu"; +#endif +} + +inline void print_result_box(const std::string& heading, + const std::string& result) { + const size_t width = std::max(heading.size(), result.size()); + const std::string border(width + 2, '-'); + std::printf("\n+%s+\n", border.c_str()); + std::printf("| %-*s |\n", int(width), heading.c_str()); + std::printf("| %-*s |\n", int(width), result.c_str()); + std::printf("+%s+\n", border.c_str()); +} + +#if NVFP4_HAS_SCHEDULE +// Table slot 0 carries the gate shape's identity map; slot 1 carries the user +// shape's ownership schedule (or identity when probing fails closed). +constexpr int GATE_TABLE = 0; +constexpr int MAIN_TABLE = 1; +#endif + +inline void launch_ours(const host::DeviceSlot& s +#if NVFP4_HAS_SCHEDULE + , int table_mode +#endif + ) { + nvfp4::nvfp4_gemm_kernel + <<>>( + s.A_t, s.B_t, s.SFA_t, s.SFB_t, + s.dC, s.M, s.N, s.K +#if NVFP4_HAS_SCHEDULE + , table_mode * nvfp4::L2A_ROUTE_WORK_CAP +#endif + ); + CUDA_CHECK(cudaGetLastError()); +} + +#if NVFP4_HAS_SCHEDULE +// A nonzero placement counter voids whatever was just measured. +inline void require_clean_placement(const char* what) { + unsigned v = 0; + CUDA_CHECK(cudaMemcpyFromSymbol(&v, nvfp4::l2a_placement_errors, sizeof(v))); + if (v != 0) { + std::fprintf(stderr, + "placement audit failed (%s): %u mismatches, result void\n", + what, v); + std::abort(); + } +} +#endif + +inline void setup_device() { + cudaDeviceProp props{}; + CUDA_CHECK(cudaGetDeviceProperties(&props, 0)); + g.sms = props.multiProcessorCount; + g.clusters = g.sms / 2; + g.l2_bytes = props.l2CacheSize; + g.grid = nvfp4::launch_grid(g.clusters); + auto* fn = nvfp4::nvfp4_gemm_kernel; + CUDA_CHECK(cudaFuncSetAttribute( + fn, cudaFuncAttributeMaxDynamicSharedMemorySize, + int(sizeof(nvfp4::SmemCD)))); +#if NVFP4_HAS_SCHEDULE + CUDA_CHECK(cudaFuncSetAttribute( + nvfp4::l2a_cluster_probe, cudaFuncAttributeMaxDynamicSharedMemorySize, + int(sizeof(nvfp4::SmemCD)))); + // The whole schedule rests on all clusters being co-resident, one per + // SM pair — verify instead of assuming. + cudaLaunchConfig_t occ{}; + cudaLaunchAttribute attr{}; + attr.id = cudaLaunchAttributeClusterDimension; + attr.val.clusterDim.x = 2; + attr.val.clusterDim.y = 1; + attr.val.clusterDim.z = 1; + occ.blockDim = dim3(nvfp4::TB_SIZE, 1, 1); + occ.dynamicSmemBytes = sizeof(nvfp4::SmemCD); + occ.attrs = &attr; + occ.numAttrs = 1; + int active = 0; + CUDA_CHECK(cudaOccupancyMaxActiveClusters(&active, fn, &occ)); + if (active != g.clusters) { + std::fprintf(stderr, "active-cluster gate failed: %d of %d resident\n", + active, g.clusters); + std::abort(); + } +#endif + std::printf("DEVICE %s sm_%d%d sms=%d clusters=%d l2_bytes=%d smem/CTA=%zu\n", + props.name, props.major, props.minor, g.sms, g.clusters, + props.l2CacheSize, sizeof(nvfp4::SmemCD)); +} + +#if NVFP4_HAS_SCHEDULE +inline void upload_identity_table(int total_work, int mode) { + std::vector t(total_work); + std::iota(t.begin(), t.end(), 0); + if (!sched::verify_identity(t)) std::abort(); + sched::upload_table(t, mode); +} + +static std::vector g_schedule_duel_table; +static sched::ScheduleMode g_schedule_duel_mode = sched::ScheduleMode::AUTO; + +inline bool activate_schedule_duel() { + if (g_schedule_duel_table.empty()) { + std::fprintf(stderr, "schedule duel table was not built\n"); + return false; + } + sched::upload_table(g_schedule_duel_table, GATE_TABLE); + std::printf("SCHEDULE_DUEL activated=%s table_slot=%d\n", + sched::schedule_name(g_schedule_duel_mode), GATE_TABLE); + return true; +} + +// Census the SM sides, verify the address-hash model, classify the +// persistent grid's clusters, build + prove + upload the user shape's +// schedule. Fails CLOSED: any probe/model/shape gate miss drops the whole +// schedule to natural raster (identity table) and says so. +inline void setup_schedule(int M, int N + , size_t operand_read_bytes, + sched::ScheduleMode requested_mode, + bool prepare_duel, + sched::ScheduleMode requested_duel_mode + ) { + const int mblocks = (M + nvfp4::CLUSTER_M - 1) / nvfp4::CLUSTER_M; + const int nblocks = (N + nvfp4::BLOCK_N - 1) / nvfp4::BLOCK_N; + const int total_work = mblocks * nblocks; + const sched::ScheduleMode mode = + requested_mode == sched::ScheduleMode::AUTO + ? sched::pick_schedule(operand_read_bytes, mblocks, nblocks) + : requested_mode; + const sched::ScheduleMode duel_mode = + requested_duel_mode == sched::ScheduleMode::AUTO + ? sched::pick_schedule(operand_read_bytes, mblocks, nblocks) + : requested_duel_mode; + g_schedule_duel_table.clear(); + g_schedule_duel_mode = duel_mode; + if (prepare_duel && sched::requires_side_census(mode) != + sched::requires_side_census(duel_mode)) { + std::fprintf(stderr, + "schedule duel requires both modes to use the census " + "or both to skip it\n"); + std::abort(); + } + + // The pure controls use the same route-table lookup as the final kernel, + // but leave the topology census and ownership partition completely out. + if (!sched::requires_side_census(mode)) { + if (total_work > nvfp4::L2A_ROUTE_WORK_CAP) { + std::fprintf(stderr, + "schedule table too small: %d tiles exceed %d entries\n", + total_work, nvfp4::L2A_ROUTE_WORK_CAP); + std::abort(); + } + sched::upload_fallback_sides(); + const std::vector table = + sched::build_unowned_schedule(mode, mblocks, nblocks); + sched::upload_table(table, MAIN_TABLE); + if (prepare_duel) + g_schedule_duel_table = + sched::build_unowned_schedule(duel_mode, mblocks, nblocks); + std::printf("SCHEDULE order=%s requested=%s ownership=off " + "grid=%dx%d tiles=%d operand_read=%.1fMiB\n", + sched::schedule_name(mode), + sched::schedule_name(requested_mode), mblocks, nblocks, + total_work, double(operand_read_bytes) / double(MIB)); + return; + } + try { + constexpr uint64_t PROBE_BYTES = uint64_t(2) << 20; + void* raw = nullptr; + CUDA_CHECK(cudaMalloc(&raw, 2 * PROBE_BYTES)); + char* arena = reinterpret_cast( + (reinterpret_cast(raw) + PROBE_BYTES - 1) & + ~(PROBE_BYTES - 1)); + const l2side::RuntimeMap map = + l2side::probe_stable(arena, PROBE_BYTES, /*repeats=*/3); + CUDA_CHECK(cudaFree(raw)); + if (map.nsm != g.sms || map.hash != l2side::kExpectedHash || + map.model_mismatches != 0) + throw std::runtime_error("runtime map inconsistent with device"); + std::printf("CENSUS sm_split=%d/%d near=%.0fns far=%.0fns hash=0x%X " + "gap=%.1fns min_margin=%.1fns jitter=%.1fns conf=%.1fns " + "repeats=%d\n", + map.side_counts[0], map.side_counts[1], map.near_ns, + map.far_ns, map.hash, map.classification_gap_ns, + map.minimum_sample_margin_ns, + map.maximum_effective_jitter_ns, + map.minimum_confidence_ns, map.stability_repeats); + + // Cluster -> SM census on the exact persistent grid we will time. + unsigned* d_smids = nullptr; + const size_t n = size_t(g.clusters) * 2; + CUDA_CHECK(cudaMalloc(&d_smids, n * sizeof(unsigned))); + CUDA_CHECK(cudaMemset(d_smids, 0xff, n * sizeof(unsigned))); + nvfp4::l2a_cluster_probe<<>>(d_smids); + CUDA_CHECK(cudaGetLastError()); + CUDA_CHECK(cudaDeviceSynchronize()); + std::vector smids(n); + CUDA_CHECK(cudaMemcpy(smids.data(), d_smids, n * sizeof(unsigned), + cudaMemcpyDeviceToHost)); + CUDA_CHECK(cudaFree(d_smids)); + const sched::Census census = sched::classify_census(smids, map, g.sms); + std::printf("CENSUS clusters=%d cluster_split=%d/%d\n", census.clusters, + census.side_clusters[0], census.side_clusters[1]); + sched::upload_side_constants(map, census); + + if (total_work > nvfp4::L2A_ROUTE_WORK_CAP) { + std::printf("SCHEDULE fallback: %d tiles exceed the %d-entry table " + "— natural raster (audit stays armed)\n", + total_work, nvfp4::L2A_ROUTE_WORK_CAP); + upload_identity_table(total_work, MAIN_TABLE); + return; + } + std::vector table = + sched::build_schedule(census, mblocks, nblocks, mode); + if (table.empty()) { + std::printf("SCHEDULE fallback: %dx%d tile grid too small for a " + "side partition — natural raster (audit stays armed)\n", + mblocks, nblocks); + upload_identity_table(total_work, MAIN_TABLE); + return; + } + if (prepare_duel) { + g_schedule_duel_table = + sched::build_schedule(census, mblocks, nblocks, duel_mode); + if (g_schedule_duel_table.empty()) { + std::fprintf(stderr, "second duel schedule is not buildable\n"); + std::abort(); + } + } + sched::upload_table(table, MAIN_TABLE); + std::printf("SCHEDULE order=%s requested=%s requester_assignment=%s " + "grid=%dx%d tiles=%d operand_read=%.1fMiB\n", + sched::schedule_name(mode), + sched::schedule_name(requested_mode), + sched::requester_assignment(mode), mblocks, nblocks, + total_work, + double(operand_read_bytes) / double(MIB)); + if (prepare_duel) + std::printf("SCHEDULE_DUEL prepared=%s requester_assignment=%s\n", + sched::schedule_name(duel_mode), + sched::requester_assignment(duel_mode)); + } catch (const std::exception& e) { + std::printf("SCHEDULE fallback: %s — the side-hash model or census " + "gates failed closed; running natural raster with an " + "identity table and an inert placement audit\n", e.what()); + sched::upload_fallback_sides(); + upload_identity_table(total_work, MAIN_TABLE); + } +} +#endif // NVFP4_HAS_SCHEDULE + +// ---- validation ---------------------------------------------------------------- + +inline float half_bits_to_float(uint16_t bits) { + __half_raw r; + r.x = bits; + return __half2float(__half(r)); +} + +inline std::vector copy_output(const void* device, size_t elements) { + std::vector result(elements); + CUDA_CHECK(cudaMemcpy(result.data(), device, elements * sizeof(uint16_t), + cudaMemcpyDeviceToHost)); + return result; +} + +inline bool validate_against(const char* label, const std::vector& got, + const std::vector& want) { + if (got.size() != want.size()) return false; + int bad = 0, nonfinite = 0; + double max_abs = 0.0, max_rel = 0.0; + for (size_t i = 0; i < got.size(); ++i) { + const double actual = half_bits_to_float(got[i]); + const double expected = want[i]; + if (!std::isfinite(actual) || !std::isfinite(expected)) { + ++bad; + ++nonfinite; + continue; + } + const double abs_error = std::fabs(actual - expected); + max_abs = std::max(max_abs, abs_error); + max_rel = std::max(max_rel, abs_error / std::max(1.0, std::fabs(expected))); + if (abs_error > FP16_ABS_TOL + FP16_REL_TOL * std::fabs(expected)) ++bad; + } + std::printf("CORRECTNESS %s bad=%d/%zu nonfinite=%d max_abs=%.5f " + "max_rel=%.6f %s\n", + label, bad, got.size(), nonfinite, max_abs, max_rel, + bad ? "FAIL" : "PASS"); + return bad == 0; +} + +inline bool validate_pair(const char* label, const std::vector& lhs, + const std::vector& rhs) { + if (lhs.size() != rhs.size()) return false; + int bad = 0, nonfinite = 0; + double max_abs = 0.0; + for (size_t i = 0; i < lhs.size(); ++i) { + const double a = half_bits_to_float(lhs[i]); + const double b = half_bits_to_float(rhs[i]); + if (!std::isfinite(a) || !std::isfinite(b)) { + ++bad; + ++nonfinite; + continue; + } + const double abs_error = std::fabs(a - b); + max_abs = std::max(max_abs, abs_error); + if (abs_error > FP16_ABS_TOL + FP16_REL_TOL * std::fabs(b)) ++bad; + } + std::printf("PAIR_GATE %s bad=%d/%zu nonfinite=%d max_abs=%.5f %s\n", + label, bad, lhs.size(), nonfinite, max_abs, + bad ? "FAIL" : "PASS"); + return bad == 0; +} + +inline double median(std::vector values) { + std::sort(values.begin(), values.end()); + return values[values.size() / 2]; +} + +inline int read_sm_clock_mhz() { + int device = 0; + char pci_bus_id[32]{}; + CUDA_CHECK(cudaGetDevice(&device)); + CUDA_CHECK(cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), device)); + + const std::string bus_id(pci_bus_id); + const bool valid_bus_id = !bus_id.empty() && + std::all_of(bus_id.begin(), bus_id.end(), [](char c) { + return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || + (c >= 'A' && c <= 'F') || c == ':' || c == '.'; + }); + if (!valid_bus_id) { + std::fprintf(stderr, "invalid CUDA PCI bus id for clock gate: %s\n", + bus_id.c_str()); + std::abort(); + } + + const std::string command = "nvidia-smi --id=" + bus_id + + " --query-gpu=clocks.sm --format=csv,noheader,nounits"; + FILE* pipe = popen(command.c_str(), "r"); + int clock_mhz = -1; + const int scanned = pipe ? std::fscanf(pipe, "%d", &clock_mhz) : 0; + const int status = pipe ? pclose(pipe) : -1; + if (scanned != 1 || status != 0) { + std::fprintf(stderr, + "SM clock query failed for %s: scanned=%d status=%d\n", + bus_id.c_str(), scanned, status); + std::abort(); + } + return clock_mhz; +} + +inline void require_pinned_clock(int clock_mhz) { + if (g.mode != BenchMode::CLOCK_PINNED) return; + if (std::abs(clock_mhz - PINNED_CLOCK_MHZ) > PINNED_CLOCK_TOL_MHZ) { + std::fprintf(stderr, + "clock-pinned mode requires %d MHz (+/-%d), got %d MHz; " + "use gb300/nvfp4/run_bench_modes.sh or lock the clock first\n", + PINNED_CLOCK_MHZ, PINNED_CLOCK_TOL_MHZ, clock_mhz); + std::abort(); + } +} + +inline void establish_clock_reference() { + if (g.cooldown_sec <= 0 || g.reference_clock_mhz >= 0) return; + CUDA_CHECK(cudaDeviceSynchronize()); + + int elapsed_sec = 0; + int previous_mhz = read_sm_clock_mhz(); + int clock_mhz = previous_mhz; + while (elapsed_sec < CLOCK_COOLDOWN_MAX_SEC) { + std::this_thread::sleep_for(std::chrono::seconds(1)); + ++elapsed_sec; + clock_mhz = read_sm_clock_mhz(); + if (previous_mhz <= ACTIVE_IDLE_CLOCK_MAX_MHZ && + clock_mhz <= ACTIVE_IDLE_CLOCK_MAX_MHZ && + std::abs(clock_mhz - previous_mhz) <= CLOCK_MATCH_TOL_MHZ) { + g.reference_clock_mhz = clock_mhz; + require_pinned_clock(clock_mhz); + std::printf("CLOCK_REFERENCE waited=%ds clock_mhz=%d tolerance_mhz=%d PASS\n", + elapsed_sec, clock_mhz, CLOCK_MATCH_TOL_MHZ); + if (g.mode == BenchMode::CLOCK_PINNED) { + std::printf("CLOCK_PIN target_mhz=%d clock_mhz=%d tolerance_mhz=%d PASS\n", + PINNED_CLOCK_MHZ, clock_mhz, + PINNED_CLOCK_TOL_MHZ); + } + return; + } + previous_mhz = clock_mhz; + } + + std::fprintf(stderr, + "clock reference failed: last_clock_mhz=%d after %ds\n", + clock_mhz, elapsed_sec); + std::abort(); +} + +inline void cooldown() { + if (g.cooldown_sec <= 0) return; + if (g.reference_clock_mhz < 0) { + std::fprintf(stderr, "clock reference was not established before timing\n"); + std::abort(); + } + CUDA_CHECK(cudaDeviceSynchronize()); + int elapsed_sec = g.cooldown_sec; + std::this_thread::sleep_for(std::chrono::seconds(elapsed_sec)); + int clock_mhz = read_sm_clock_mhz(); + while (clock_mhz > ACTIVE_IDLE_CLOCK_MAX_MHZ || + std::abs(clock_mhz - g.reference_clock_mhz) > CLOCK_MATCH_TOL_MHZ) { + if (elapsed_sec >= CLOCK_COOLDOWN_MAX_SEC) { + std::fprintf(stderr, + "cooldown clock gate failed: clock_mhz=%d " + "reference_mhz=%d tolerance_mhz=%d after %ds\n", + clock_mhz, g.reference_clock_mhz, + CLOCK_MATCH_TOL_MHZ, elapsed_sec); + std::abort(); + } + std::this_thread::sleep_for(std::chrono::seconds(1)); + ++elapsed_sec; + clock_mhz = read_sm_clock_mhz(); + } + require_pinned_clock(clock_mhz); + std::printf("COOLDOWN seconds=%d clock_mhz=%d reference_mhz=%d PASS\n", + elapsed_sec, clock_mhz, g.reference_clock_mhz); +} + +// ---- the small correctness gates --------------------------------------------------- +// Two small shapes, each checked against an O(MNK) host reference over the +// dequantized operands, both with partial M and N tiles, random N(0,1) data +// on principle (several classes of scale-layout bugs pass with constant +// inputs). cuBLASLt's FP4 path only has kernels for K % 32 == 0, so K=1136 +// is ours-only while K=1152 can also cross-check the vendor result. +// The first shape exercises a 16-mod-32 ragged tail; the second exercises a +// 32-aligned partial 768-element group. Rounded-mainloop rungs rely on TMA +// zero fill, while exact-tail rungs stop issuing dead batches. + +// Shared gate: fixture, poisoned output, host reference, and determinism. It +// hands the fixture, reference, and output back so the aligned case can also +// cross-validate cuBLASLt. +inline bool small_gate_ours(int M, int N, int K, uint64_t seed, + const char* tag, host::Fixture* f_out, + std::vector* ref_out, + std::vector* out) { + *f_out = host::make_fixture(M, N, K, seed, /*deq=*/true); + const host::Fixture& f = *f_out; +#if NVFP4_HAS_SCHEDULE + { + const int mblocks = (M + nvfp4::CLUSTER_M - 1) / nvfp4::CLUSTER_M; + const int nblocks = (N + nvfp4::BLOCK_N - 1) / nvfp4::BLOCK_N; + upload_identity_table(mblocks * nblocks, GATE_TABLE); + } +#endif + host::DeviceSlot* ours = host::make_device_slot(f.inputs, M, N); + const size_t elems = size_t(M) * N; + // Output poisoned before every gated run: a kernel that silently skipped + // work would be caught by the leftover NaNs. + CUDA_CHECK(cudaMemset(ours->dC, 0xFF, elems * sizeof(uint16_t))); + launch_ours(*ours +#if NVFP4_HAS_SCHEDULE + , GATE_TABLE +#endif + ); + CUDA_CHECK(cudaDeviceSynchronize()); +#if NVFP4_HAS_SCHEDULE + require_clean_placement(tag); +#endif + *out = copy_output(ours->dC, elems); + + ref_out->assign(elems, 0.0f); + for (int m = 0; m < M; ++m) + for (int k = 0; k < K; ++k) { + const float a = f.A_deq[size_t(m) * K + k]; + for (int n = 0; n < N; ++n) + (*ref_out)[size_t(m) * N + n] += a * f.B_deq[size_t(n) * K + k]; + } + const std::string label = std::string("OURS-") + tag + "-vs-host-ref"; + bool ok = validate_against(label.c_str(), *out, *ref_out); + + CUDA_CHECK(cudaMemset(ours->dC, 0xFF, elems * sizeof(uint16_t))); + launch_ours(*ours +#if NVFP4_HAS_SCHEDULE + , GATE_TABLE +#endif + ); + CUDA_CHECK(cudaDeviceSynchronize()); + const auto rerun = copy_output(ours->dC, elems); + const bool det = (*out == rerun); + std::printf("DETERMINISM OURS-%s %s\n", tag, det ? "PASS" : "FAIL"); + host::free_device_slot(ours); + return ok && det; +} + +inline bool small_correctness_gates() { + const int M = 768, N = 512; + host::Fixture f; + std::vector ref; + std::vector ours_out; + + // The ragged logical K, checked against the host reference. + bool ok = small_gate_ours(M, N, /*K=*/1136, 0xC011AB1EULL, + "small-ragged", &f, &ref, &ours_out); + +#if NVFP4_HAS_K64_TAIL + // (a2) the K64-dispatch classes, against the host reference: K=1280 + // (r=512: 4xK96 + 2xK64 — both plain SF words, two live A/B windows) and + // K=1024 (r=256: 2xK96 + 1xK64 — word-0 SF, ONE live window, the + // t_win==1 cursor path; the K=16384 tail class). K=1120 (r=352, odd a) + // must CLASSIFY as inactive and run the original tail — a regression + // gate on the dispatch itself. The K=1136/1152 gates above/below never + // dispatch (r=368 is 16 mod 32; r=384 is 96-aligned). + ok = small_gate_ours(M, N, /*K=*/1280, 0xC011AB20ULL, + "small-k64-pair", &f, &ref, &ours_out) && ok; + ok = small_gate_ours(M, N, /*K=*/1024, 0xC011AB22ULL, + "small-k64-single", &f, &ref, &ours_out) && ok; + ok = small_gate_ours(M, N, /*K=*/1120, 0xC011AB21ULL, + "small-k64-fallback", &f, &ref, &ours_out) && ok; +#endif + + // (b) the 32-aligned twin, cross-validated against cuBLASLt. + const int K = 1152; + ok = small_gate_ours(M, N, K, 0xC011AB1FULL, "small-aligned", + &f, &ref, &ours_out) && ok; + LtPlan plan(M, N, K); + LtSlot vendor = plan.make_slot(f); + plan.initialize_algorithm(vendor); + plan.add_workspace(vendor); + const size_t elems = size_t(M) * N; + CUDA_CHECK(cudaMemset(vendor.d, 0xFF, elems * sizeof(uint16_t))); + plan.launch(vendor); + CUDA_CHECK(cudaDeviceSynchronize()); + const auto vout = copy_output(vendor.d, elems); + ok = validate_against("CUBLAS-h0-small-vs-host-ref", vout, ref) && ok; + ok = validate_pair("OURS-vs-CUBLAS-h0-small", ours_out, vout) && ok; + + plan.free_slot(vendor); + return ok; +} + +// ---- user-shape gates ------------------------------------------------------------- +// Our kernel's deterministic user-shape output is the anchor. The first +// cuBLASLt recommendation must launch, agree with the anchor within FP16 +// tolerance, and be deterministic across two launches. + +inline bool user_shape_gates(host::DeviceSlot& ours, LtPlan& plan, + LtSlot& vendor) { + const size_t elems = size_t(ours.M) * ours.N; + CUDA_CHECK(cudaMemset(ours.dC, 0xFF, elems * sizeof(uint16_t))); + launch_ours(ours +#if NVFP4_HAS_SCHEDULE + , MAIN_TABLE +#endif + ); + CUDA_CHECK(cudaDeviceSynchronize()); +#if NVFP4_HAS_SCHEDULE + require_clean_placement("user-shape gate"); +#endif + auto anchor = copy_output(ours.dC, elems); + CUDA_CHECK(cudaMemset(ours.dC, 0xFF, elems * sizeof(uint16_t))); + launch_ours(ours +#if NVFP4_HAS_SCHEDULE + , MAIN_TABLE +#endif + ); + CUDA_CHECK(cudaDeviceSynchronize()); + const auto rerun = copy_output(ours.dC, elems); + bool all_ok = (anchor == rerun); + std::printf("DETERMINISM OURS-user %s\n", all_ok ? "PASS" : "FAIL"); + + CUDA_CHECK(cudaMemset(vendor.d, 0xFF, elems * sizeof(uint16_t))); + const cublasStatus_t status = plan.launch_status(vendor); + if (status != CUBLAS_STATUS_SUCCESS) { + std::printf("CUBLAS_GATE index=0 launch_status=%d FAIL\n", int(status)); + return false; + } + CUDA_CHECK(cudaDeviceSynchronize()); + const auto first = copy_output(vendor.d, elems); + bool ok = validate_pair("OURS-vs-CUBLAS-h0", anchor, first); + CUDA_CHECK(cudaMemset(vendor.d, 0xFF, elems * sizeof(uint16_t))); + plan.launch(vendor); + CUDA_CHECK(cudaDeviceSynchronize()); + const auto second = copy_output(vendor.d, elems); + const bool deterministic = (first == second); + ok = ok && deterministic; + std::printf("CUBLAS_GATE index=0 workspace=%zu determinism=%s %s\n", + plan.workspace_bytes, deterministic ? "PASS" : "FAIL", + ok ? "PASS" : "FAIL"); + all_ok = all_ok && ok; + return all_ok; +} +#if NVFP4_HAS_SCHEDULE + +inline bool schedule_duel_gate(host::DeviceSlot& ours) { + const size_t elems = size_t(ours.M) * ours.N; + CUDA_CHECK(cudaMemset(ours.dC, 0xFF, elems * sizeof(uint16_t))); + launch_ours(ours, MAIN_TABLE); + CUDA_CHECK(cudaDeviceSynchronize()); + const auto primary = copy_output(ours.dC, elems); + + CUDA_CHECK(cudaMemset(ours.dC, 0xFF, elems * sizeof(uint16_t))); + launch_ours(ours, GATE_TABLE); + CUDA_CHECK(cudaDeviceSynchronize()); + require_clean_placement("schedule-duel gate"); + const auto duel = copy_output(ours.dC, elems); + bool ok = validate_pair("SCHEDULE-DUEL-vs-primary", duel, primary); + + CUDA_CHECK(cudaMemset(ours.dC, 0xFF, elems * sizeof(uint16_t))); + launch_ours(ours, GATE_TABLE); + CUDA_CHECK(cudaDeviceSynchronize()); + const auto rerun = copy_output(ours.dC, elems); + const bool deterministic = (duel == rerun); + std::printf("DETERMINISM SCHEDULE-DUEL %s\n", + deterministic ? "PASS" : "FAIL"); + return ok && deterministic; +} +#endif + +// ---- the rotating cache-cold benchmark ---------------------------------------------- +// Every measured launch reads a DIFFERENT copy of A/B/scales, and the copies +// cycle so that by the time one is reused, far more data than the L2 holds +// has passed through the cache. Sacrificial slots (never timed) flush the +// previous arm's residue before the events start. + +template +double measure_rotating(int sacrificial, int timed, int iters, Launch launch) { + for (int i = 0; i < sacrificial * SACRIFICIAL_PASSES; ++i) + launch(i % sacrificial); + CUDA_CHECK(cudaDeviceSynchronize()); + cudaEvent_t start, end; + CUDA_CHECK(cudaEventCreate(&start)); + CUDA_CHECK(cudaEventCreate(&end)); + CUDA_CHECK(cudaEventRecord(start)); + for (int i = 0; i < iters; ++i) + launch(sacrificial + i % timed); + CUDA_CHECK(cudaEventRecord(end)); + CUDA_CHECK(cudaEventSynchronize(end)); + float elapsed_ms = 0.0f; + CUDA_CHECK(cudaEventElapsedTime(&elapsed_ms, start, end)); + cudaEventDestroy(end); + cudaEventDestroy(start); + return double(elapsed_ms) * 1000.0 / iters; // microseconds per launch +} + +struct TritonTiming { + double median_us = 0.0; + double q1_us = 0.0; + double q3_us = 0.0; + int warmup_iters = 0; + int timed_iters = 0; +}; + +template +TritonTiming measure_triton_style(Launch launch, ClearL2 clear_l2, + double warmup_ms = 25.0, + double measure_ms = 100.0) { + launch(); + CUDA_CHECK(cudaDeviceSynchronize()); + + cudaEvent_t estimate_start, estimate_end; + CUDA_CHECK(cudaEventCreate(&estimate_start)); + CUDA_CHECK(cudaEventCreate(&estimate_end)); + CUDA_CHECK(cudaEventRecord(estimate_start)); + for (int i = 0; i < 5; ++i) launch(); + CUDA_CHECK(cudaEventRecord(estimate_end)); + CUDA_CHECK(cudaEventSynchronize(estimate_end)); + float estimate_ms = 0.0f; + CUDA_CHECK(cudaEventElapsedTime(&estimate_ms, estimate_start, estimate_end)); + cudaEventDestroy(estimate_end); + cudaEventDestroy(estimate_start); + + const double per_launch_ms = double(estimate_ms) / 5.0; + TritonTiming timing; + timing.warmup_iters = std::max(1, int(warmup_ms / per_launch_ms)); + timing.timed_iters = std::max(1, int(measure_ms / per_launch_ms)); + for (int i = 0; i < timing.warmup_iters; ++i) launch(); + + std::vector starts(size_t(timing.timed_iters)); + std::vector ends(size_t(timing.timed_iters)); + for (int i = 0; i < timing.timed_iters; ++i) { + CUDA_CHECK(cudaEventCreate(&starts[size_t(i)])); + CUDA_CHECK(cudaEventCreate(&ends[size_t(i)])); + clear_l2(); + CUDA_CHECK(cudaEventRecord(starts[size_t(i)])); + launch(); + CUDA_CHECK(cudaEventRecord(ends[size_t(i)])); + } + CUDA_CHECK(cudaDeviceSynchronize()); + CUDA_CHECK(cudaGetLastError()); + + std::vector samples; + samples.reserve(size_t(timing.timed_iters)); + for (int i = 0; i < timing.timed_iters; ++i) { + float elapsed_ms = 0.0f; + CUDA_CHECK(cudaEventElapsedTime( + &elapsed_ms, starts[size_t(i)], ends[size_t(i)])); + samples.push_back(double(elapsed_ms) * 1000.0); + cudaEventDestroy(ends[size_t(i)]); + cudaEventDestroy(starts[size_t(i)]); + } + std::sort(samples.begin(), samples.end()); + auto quantile = [&](double q) { + const double position = q * double(samples.size() - 1); + const size_t low = size_t(position); + const size_t high = std::min(low + 1, samples.size() - 1); + return samples[low] + (position - double(low)) * + (samples[high] - samples[low]); + }; + timing.q1_us = quantile(0.25); + timing.median_us = quantile(0.50); + timing.q3_us = quantile(0.75); + return timing; +} + +struct SustainedTiming { + double elapsed_seconds = 0.0; + double first_us = 0.0; + double tail_us = 0.0; + int chunks = 0; + int first_chunks = 0; + int tail_chunks = 0; +}; + +template +SustainedTiming measure_sustained_block(int timed_slot_offset, int timed_slots, + int chunk_iters, Launch launch) { + long long launch_index = 0; + auto launch_chunk = [&] { + for (int i = 0; i < chunk_iters; ++i) { + launch(timed_slot_offset + + int(launch_index++ % timed_slots)); + } + }; + + cudaEvent_t boundaries[3]; + for (auto& event : boundaries) CUDA_CHECK(cudaEventCreate(&event)); + CUDA_CHECK(cudaDeviceSynchronize()); + CUDA_CHECK(cudaEventRecord(boundaries[0])); + launch_chunk(); + CUDA_CHECK(cudaEventRecord(boundaries[1])); + launch_chunk(); + CUDA_CHECK(cudaEventRecord(boundaries[2])); + + SustainedTiming timing; + double first_sum_us = 0.0; + double tail_sum_us = 0.0; + int in_flight = 2; + bool feeding = true; + for (int chunk = 0; in_flight > 0; ++chunk, --in_flight) { + CUDA_CHECK(cudaEventSynchronize(boundaries[(chunk + 1) % 3])); + CUDA_CHECK(cudaGetLastError()); + float elapsed_ms = 0.0f; + CUDA_CHECK(cudaEventElapsedTime( + &elapsed_ms, boundaries[chunk % 3], + boundaries[(chunk + 1) % 3])); + const double chunk_start_seconds = timing.elapsed_seconds; + const double launch_us = double(elapsed_ms) * 1000.0 / chunk_iters; + timing.elapsed_seconds += double(elapsed_ms) / 1000.0; + ++timing.chunks; + if (chunk_start_seconds < SUSTAINED_FIRST_SECONDS) { + first_sum_us += launch_us; + ++timing.first_chunks; + } + if (chunk_start_seconds >= SUSTAINED_TAIL_FROM_SECONDS) { + tail_sum_us += launch_us; + ++timing.tail_chunks; + } + if (feeding && timing.elapsed_seconds >= SUSTAINED_BLOCK_SECONDS) + feeding = false; + if (feeding) { + launch_chunk(); + CUDA_CHECK(cudaEventRecord(boundaries[chunk % 3])); + ++in_flight; + } + } + for (auto& event : boundaries) cudaEventDestroy(event); + if (timing.first_chunks == 0 || timing.tail_chunks == 0) { + std::fprintf(stderr, "sustained block was too short for tail sampling\n"); + std::abort(); + } + timing.first_us = first_sum_us / timing.first_chunks; + timing.tail_us = tail_sum_us / timing.tail_chunks; + return timing; +} + +inline double mean(const std::vector& values) { + return std::accumulate(values.begin(), values.end(), 0.0) / + double(values.size()); +} + +} // namespace bench + +int main(int argc, char** argv) { + if (argc < 4) { +#if NVFP4_HAS_SCHEDULE + std::fprintf(stderr, + "usage: %s M N K [--rounds R] [--cooldown-sec S] " + "[--bench-mode MODE] [--check-only] [--ours-only] [--schedule MODE] " + "[--schedule-duel MODE_A MODE_B]\n", + argv[0]); + std::fprintf(stderr, + "benchmark modes: current, clock-pinned, sustained, triton\n"); + std::fprintf(stderr, + "schedule modes: auto, raster, pocket-8x4, pocket-8x8, hilbert, " + "blind-plain, blind-pocket-8x4, blind-pocket-8x8, " + "hilbert-in-blind, " + "owned-plain, owned-pocket-8x4, owned-pocket-8x8, " + "hilbert-in-owned\n"); +#else + std::fprintf(stderr, + "usage: %s M N K [--rounds R] [--cooldown-sec S] " + "[--bench-mode MODE] [--check-only] [--ours-only]\n", + argv[0]); + std::fprintf(stderr, + "benchmark modes: current, clock-pinned, sustained, triton\n"); +#endif + return 2; + } + const int M = std::atoi(argv[1]); + const int N = std::atoi(argv[2]); + const int K = std::atoi(argv[3]); + int rounds = -1; + bool check_only = false; + // --ours-only: the quick-signal lane — skip timed cuBLASLt vendor legs; + // absolute paired old-vs-new reads come from two + // alternated invocations of this mode (gates still run first). + bool ours_only = false; +#if NVFP4_HAS_SCHEDULE + sched::ScheduleMode schedule_mode = sched::ScheduleMode::AUTO; + sched::ScheduleMode duel_mode = sched::ScheduleMode::AUTO; + bool schedule_explicit = false; + bool schedule_duel = false; +#endif + for (int i = 4; i < argc; ++i) { + const std::string a = argv[i]; + if (a == "--rounds" && i + 1 < argc) rounds = std::atoi(argv[++i]); + else if (a == "--cooldown-sec" && i + 1 < argc) + bench::g.cooldown_sec = std::atoi(argv[++i]); + else if (a == "--bench-mode" && i + 1 < argc) { + if (!bench::parse_bench_mode(argv[++i], &bench::g.mode)) { + std::fprintf(stderr, + "--bench-mode needs current, clock-pinned, " + "sustained, or triton\n"); + return 2; + } + } + else if (a == "--check-only") check_only = true; + else if (a == "--ours-only") ours_only = true; +#if NVFP4_HAS_SCHEDULE + else if (a == "--schedule") { + if (schedule_explicit || schedule_duel || i + 1 >= argc || + !sched::parse_schedule_mode(argv[++i], &schedule_mode)) { + std::fprintf(stderr, + "--schedule may appear once and needs one of: auto, " + "raster, pocket-8x4, " + "pocket-8x8, hilbert, blind-plain, blind-pocket-8x4, " + "blind-pocket-8x8, hilbert-in-blind, owned-plain, " + "owned-pocket-8x4, " + "owned-pocket-8x8, hilbert-in-owned\n"); + return 2; + } + schedule_explicit = true; + } + else if (a == "--schedule-duel") { + if (schedule_explicit || schedule_duel || i + 2 >= argc || + !sched::parse_schedule_mode(argv[i + 1], &schedule_mode) || + !sched::parse_schedule_mode(argv[i + 2], &duel_mode) || + schedule_mode == sched::ScheduleMode::AUTO || + duel_mode == sched::ScheduleMode::AUTO) { + std::fprintf(stderr, + "--schedule-duel needs two explicit schedule modes from " + "the same ownership class\n"); + return 2; + } + i += 2; + schedule_duel = true; + } +#endif + else { + std::fprintf(stderr, "unknown argument: %s\n", a.c_str()); + return 2; + } + } + if (rounds < 0) { + rounds = bench::g.mode == bench::BenchMode::SUSTAINED ? 3 + : bench::g.mode == bench::BenchMode::TRITON ? 1 + : 5; + } + if (M < 1 || N < 1 || K < 16 || K % 16 != 0 || rounds < 1) { + std::fprintf(stderr, + "shape gate: need M,N >= 1 and K >= 16 with K %% 16 == 0 " + "(one UE4M3 scale per 16 K elements)\n"); + return 2; + } + if (bench::g.mode != bench::BenchMode::CURRENT && ours_only) { + std::fprintf(stderr, + "--ours-only is supported only by --bench-mode current\n"); + return 2; + } +#if NVFP4_HAS_SCHEDULE + if (bench::g.mode != bench::BenchMode::CURRENT && schedule_duel) { + std::fprintf(stderr, + "--schedule-duel is supported only by --bench-mode current\n"); + return 2; + } +#endif + if (bench::g.mode == bench::BenchMode::CLOCK_PINNED && + bench::g.cooldown_sec <= 0) { + std::fprintf(stderr, + "clock-pinned mode requires a positive --cooldown-sec\n"); + return 2; + } + if (bench::g.mode == bench::BenchMode::TRITON && rounds != 1) { + std::fprintf(stderr, "triton mode uses one do_bench-style sample; " + "set --rounds 1 or omit it\n"); + return 2; + } + if (bench::g.cooldown_sec < 0 || + bench::g.cooldown_sec > bench::CLOCK_COOLDOWN_MAX_SEC) { + std::fprintf(stderr, "--cooldown-sec must be between 0 and %d\n", + bench::CLOCK_COOLDOWN_MAX_SEC); + return 2; + } + // The baseline's constraint, not ours: cuBLASLt's FP4 path has kernels + // only for K % 32 == 0 (its algorithm query returns NOT_SUPPORTED + // otherwise). This kernel takes any K % 16 == 0 — the built-in small + // gate covers the ragged 16-mod-32 case against a host reference — but + // this harness exists to race cuBLASLt, so refuse rather than skip. + // (M and N stay unconstrained: the FP4 packing constraint rides the + // K-major leading dimension, and a shape cuBLASLt still refuses fails + // loudly at the gated algorithm query, never silently.) + if (K % 32 != 0) { + std::fprintf(stderr, + "cuBLASLt comparison refused: its FP4 path requires " + "K %% 32 == 0 and K=%d gives no supporting kernel. " + "The kernel itself accepts this K; pick a 32-aligned K " + "to race the baseline.\n", K); + return 2; + } + std::printf("NVFP4 GEMM build=%s shape=%dx%dx%d bench_mode=%s " + "rounds=%d cooldown=%ds\n", + NVFP4_BUILD_NAME, M, N, K, + bench::bench_mode_name(bench::g.mode), rounds, + bench::g.cooldown_sec); + + bench::setup_device(); + + // The user-shape fixture. Deterministic seed, mixed with the shape so + // different shapes get different data. + const uint64_t seed = 0xCAFEF00DULL ^ uint64_t(M) ^ (uint64_t(K) << 17); + host::Fixture f = host::make_fixture(M, N, K, seed, /*deq=*/false); + // Both launch paths copy these same host buffers. At public K%32 shapes, + // the TMA stride adds no A/B padding, so they also equal the logical + // packed-E2M1 inputs byte for byte. + const bool identical_inputs = + f.inputs.A == f.A_logical && f.inputs.B == f.B_logical; + std::printf("INPUT_GATE ours-vs-cublas source buffers=%s\n", + identical_inputs ? "IDENTICAL" : "DIFFERENT"); + if (!identical_inputs) return 1; + const size_t ours_read = f.inputs.A.size() + f.inputs.B.size() + + f.inputs.SFA.size() + f.inputs.SFB.size(); + const size_t vendor_read = ours_read; +#if NVFP4_HAS_SCHEDULE + bench::setup_schedule(M, N + , ours_read, schedule_mode, schedule_duel, duel_mode + ); +#endif + + // Gate 1: the small clipped problems against an O(MNK) host reference. + if (!bench::small_correctness_gates()) { + std::fprintf(stderr, "small correctness gates FAILED\n"); + return 1; + } +#if NVFP4_HAS_SCHEDULE + if (schedule_duel && !bench::activate_schedule_duel()) return 1; +#endif + + // Gate 2: use the first cuBLASLt recommendation for every shape. + bench::LtPlan plan(M, N, K); + host::DeviceSlot* ours0 = host::make_device_slot(f.inputs, M, N); + bench::LtSlot vendor0 = plan.make_slot(f); + plan.initialize_algorithm(vendor0); + plan.add_workspace(vendor0); + if (!bench::user_shape_gates(*ours0, plan, vendor0)) { + std::fprintf(stderr, "user-shape correctness gates FAILED\n"); + return 1; + } +#if NVFP4_HAS_SCHEDULE + if (schedule_duel && !bench::schedule_duel_gate(*ours0)) { + std::fprintf(stderr, "schedule-duel correctness gate FAILED\n"); + return 1; + } +#endif + if (check_only) { + std::printf("CHECK-ONLY all gates PASS shape=%dx%dx%d\n", M, N, K); + return 0; + } + + // ---- benchmark sizing: enough independent copies that nothing a kernel + // is timed on can still be resident from its previous use. + const size_t required_reuse = std::max( + size_t(256) * bench::MIB, size_t(2) * size_t(bench::g.l2_bytes)); + const size_t min_read = std::min(ours_read, vendor_read); + const int timed_slots = std::max( + 3, int((required_reuse + min_read - 1) / min_read) + 1); + const int sacrificial_slots = timed_slots - 1; + const int total_slots = timed_slots + sacrificial_slots; + const int iters = timed_slots * bench::TIMED_PASSES; + if (total_slots > 40) { + std::fprintf(stderr, + "shape too small for the rotating cache-cold protocol " + "(%d slots needed; operand read is only %zu bytes)\n", + total_slots, min_read); + return 1; + } + std::printf("CACHE_COLD required_reuse=%zu ours_read=%zu cublas_read=%zu " + "sacrificial=%d timed=%d iters=%d\n", + required_reuse, ours_read, vendor_read, sacrificial_slots, + timed_slots, iters); + + std::vector ours_slots{ours0}; + std::vector vendor_slots{vendor0}; + for (int i = 1; i < total_slots; ++i) { + ours_slots.push_back(host::make_device_slot(f.inputs, M, N)); + vendor_slots.push_back(plan.make_slot(f)); + plan.add_workspace(vendor_slots.back()); + } + + auto launch_ours_slot = [&](int s) { + bench::launch_ours(*ours_slots[size_t(s)] +#if NVFP4_HAS_SCHEDULE + , bench::MAIN_TABLE +#endif + ); + }; +#if NVFP4_HAS_SCHEDULE + auto launch_duel_slot = [&](int s) { + bench::launch_ours(*ours_slots[size_t(s)], bench::GATE_TABLE); + }; +#endif + auto launch_vendor_slot = [&](int s) { + plan.launch(vendor_slots[size_t(s)]); + }; + auto free_benchmark_slots = [&] { + for (auto& slot : vendor_slots) plan.free_slot(slot); + for (auto* slot : ours_slots) host::free_device_slot(slot); + }; + bench::establish_clock_reference(); +#if NVFP4_HAS_SCHEDULE + + if (schedule_duel) { + const double flop = 2.0 * double(M) * double(N) * double(K); + auto pflops = [&](double us) { return flop / (us * 1e9); }; + std::vector primary_us_all, duel_us_all, duel_over_primary; + int duel_wins = 0; + for (int r = 0; r < rounds; ++r) { + double primary_us = 0.0, duel_us = 0.0; + const bool primary_first = (r % 2 == 0); + for (int arm = 0; arm < 2; ++arm) { + const bool run_primary = (arm == 0) == primary_first; + bench::cooldown(); + const double us = run_primary + ? bench::measure_rotating(sacrificial_slots, timed_slots, + iters, launch_ours_slot) + : bench::measure_rotating(sacrificial_slots, timed_slots, + iters, launch_duel_slot); + bench::require_clean_placement("schedule-duel timed round"); + if (run_primary) primary_us = us; + else duel_us = us; + } + const double ratio = primary_us / duel_us; + primary_us_all.push_back(primary_us); + duel_us_all.push_back(duel_us); + duel_over_primary.push_back(ratio); + duel_wins += duel_us < primary_us; + std::printf("SCHEDULE_DUEL_ROUND %d/%d order=%s " + "a=%s %.3fus (%.3f PF) b=%s %.3fus (%.3f PF) " + "b_over_a=%.6f %s\n", + r + 1, rounds, primary_first ? "a-first" : "b-first", + sched::schedule_name(schedule_mode), primary_us, + pflops(primary_us), sched::schedule_name(duel_mode), + duel_us, pflops(duel_us), ratio, + duel_us < primary_us ? "B_WIN" : "A_WIN"); + } + const double primary_med = bench::median(primary_us_all); + const double duel_med = bench::median(duel_us_all); + const double pair_med = bench::median(duel_over_primary); + double log_ratio_sum = 0.0; + for (const double ratio : duel_over_primary) + log_ratio_sum += std::log(ratio); + const double pair_geomean = + std::exp(log_ratio_sum / duel_over_primary.size()); + std::printf("SCHEDULE_DUEL_RESULT a=%s %.3fus %.3fPF b=%s %.3fus " + "%.3fPF paired_median=%.6f paired_geomean=%.6f " + "b_wins=%d/%d\n", + sched::schedule_name(schedule_mode), primary_med, + pflops(primary_med), sched::schedule_name(duel_mode), + duel_med, pflops(duel_med), pair_med, pair_geomean, + duel_wins, rounds); + free_benchmark_slots(); + return 0; + } +#endif + + const double flop = 2.0 * double(M) * double(N) * double(K); + auto pflops = [&](double us) { return flop / (us * 1e9); }; + + if (bench::g.mode == bench::BenchMode::TRITON) { + void* l2_clear = nullptr; + CUDA_CHECK(cudaMalloc(&l2_clear, bench::TRITON_FLUSH_BYTES)); + auto clear_l2 = [&] { + CUDA_CHECK(cudaMemset(l2_clear, 0, bench::TRITON_FLUSH_BYTES)); + }; + const auto ours = bench::measure_triton_style( + [&] { launch_ours_slot(0); }, clear_l2); +#if NVFP4_HAS_SCHEDULE + bench::require_clean_placement("triton-style timing"); +#endif + const auto vendor = bench::measure_triton_style( + [&] { launch_vendor_slot(0); }, clear_l2); + CUDA_CHECK(cudaFree(l2_clear)); + + std::printf("TRITON_RESULT flush_bytes=%zu warmup_ms=25 measure_ms=100 " + "ours_us=%.3f ours_iqr=%.3f..%.3f ours_n=%d " + "cublas_us=%.3f cublas_iqr=%.3f..%.3f cublas_n=%d " + "ratio=%.5f\n", + bench::TRITON_FLUSH_BYTES, ours.median_us, + ours.q1_us, ours.q3_us, ours.timed_iters, + vendor.median_us, vendor.q1_us, vendor.q3_us, + vendor.timed_iters, vendor.median_us / ours.median_us); + char heading[256]; + char result[256]; + std::snprintf(heading, sizeof(heading), + "%s (%s) | shape %dx%dx%d | mode triton", + bench::gemm_source_file(), NVFP4_BUILD_NAME, M, N, K); + std::snprintf(result, sizeof(result), + "RESULT | ours %.3f PFLOP/s | cuBLASLt %.3f PFLOP/s | " + "%.2f%% of cuBLASLt", + pflops(ours.median_us), pflops(vendor.median_us), + 100.0 * vendor.median_us / ours.median_us); + bench::print_result_box(heading, result); + free_benchmark_slots(); + return 0; + } + + if (bench::g.mode == bench::BenchMode::SUSTAINED) { + const double ours_calibration = bench::measure_rotating( + sacrificial_slots, timed_slots, iters, launch_ours_slot); +#if NVFP4_HAS_SCHEDULE + bench::require_clean_placement("sustained calibration"); +#endif + const double vendor_calibration = bench::measure_rotating( + sacrificial_slots, timed_slots, iters, launch_vendor_slot); + std::printf("SUSTAINED_CALIBRATION ours_us=%.3f cublas_us=%.3f " + "block_seconds=%.0f tail_from_seconds=%.0f\n", + ours_calibration, vendor_calibration, + bench::SUSTAINED_BLOCK_SECONDS, + bench::SUSTAINED_TAIL_FROM_SECONDS); + + std::vector ours_tail_us; + std::vector vendor_tail_us; + int wins = 0; + for (int round = 0; round < rounds; ++round) { + double ours_tail = 0.0; + double vendor_tail = 0.0; + const bool ours_first = (round % 2) == 0; + for (int arm = 0; arm < 2; ++arm) { + const bool run_ours = (arm == 0) == ours_first; + const double estimate = run_ours ? ours_calibration + : vendor_calibration; + const int chunk_iters = + std::max(1, int(500000.0 / estimate)); + const auto timing = run_ours + ? bench::measure_sustained_block( + sacrificial_slots, timed_slots, chunk_iters, + launch_ours_slot) + : bench::measure_sustained_block( + sacrificial_slots, timed_slots, chunk_iters, + launch_vendor_slot); +#if NVFP4_HAS_SCHEDULE + if (run_ours) + bench::require_clean_placement("sustained block"); +#endif + std::printf("SUSTAINED_BLOCK round=%d/%d arm=%s " + "elapsed_s=%.2f chunks=%d chunk_iters=%d " + "first_us=%.3f tail_us=%.3f tail_pf=%.3f\n", + round + 1, rounds, + run_ours ? "ours" : "cublasLt", + timing.elapsed_seconds, timing.chunks, chunk_iters, + timing.first_us, timing.tail_us, + pflops(timing.tail_us)); + if (run_ours) ours_tail = timing.tail_us; + else vendor_tail = timing.tail_us; + } + ours_tail_us.push_back(ours_tail); + vendor_tail_us.push_back(vendor_tail); + const bool win = ours_tail < vendor_tail; + wins += win; + std::printf("SUSTAINED_ROUND %d/%d order=%s ours_tail_us=%.3f " + "cublas_tail_us=%.3f ratio=%.5f %s\n", + round + 1, rounds, + ours_first ? "ours-first" : "cublas-first", + ours_tail, vendor_tail, vendor_tail / ours_tail, + win ? "WIN" : "LOSS"); + } + const double ours_us = bench::mean(ours_tail_us); + const double vendor_us = bench::mean(vendor_tail_us); + std::printf("SUSTAINED_RESULT blocks_per_arm=%d block_seconds=%.0f " + "tail_from_seconds=%.0f ours_tail_us=%.3f " + "cublas_tail_us=%.3f ratio=%.5f wins=%d/%d\n", + rounds, bench::SUSTAINED_BLOCK_SECONDS, + bench::SUSTAINED_TAIL_FROM_SECONDS, ours_us, vendor_us, + vendor_us / ours_us, wins, rounds); + char heading[256]; + char result[256]; + std::snprintf(heading, sizeof(heading), + "%s (%s) | shape %dx%dx%d | mode sustained", + bench::gemm_source_file(), NVFP4_BUILD_NAME, M, N, K); + std::snprintf(result, sizeof(result), + "RESULT | ours %.3f PFLOP/s | cuBLASLt %.3f PFLOP/s | " + "%.2f%% of cuBLASLt | wins %d/%d", + pflops(ours_us), pflops(vendor_us), + 100.0 * vendor_us / ours_us, wins, rounds); + bench::print_result_box(heading, result); + free_benchmark_slots(); + return 0; + } + + // The current and clock-pinned modes use R rounds with alternating arm + // order. FLOPs use logical K, so padded work receives no extra credit. + std::vector ours_us_all, vendor_us_all; + int wins = 0; + if (ours_only) { + for (int r = 0; r < rounds; ++r) { + bench::cooldown(); + const double us = bench::measure_rotating( + sacrificial_slots, timed_slots, iters, launch_ours_slot); +#if NVFP4_HAS_SCHEDULE + bench::require_clean_placement("timed round"); +#endif + ours_us_all.push_back(us); + std::printf("OURS_ONLY_ROUND %d/%d ours=%.3fus (%.3f PF)\n", + r + 1, rounds, us, pflops(us)); + } + const double med = bench::median(ours_us_all); + char heading[256]; + char result[256]; + std::snprintf(heading, sizeof(heading), + "%s (%s) | shape %dx%dx%d | rounds %d", + bench::gemm_source_file(), NVFP4_BUILD_NAME, + M, N, K, rounds); + std::snprintf(result, sizeof(result), + "RESULT | ours %.3f PFLOP/s | median %.3f us | ours-only", + pflops(med), med); + bench::print_result_box(heading, result); + return 0; + } + for (int r = 0; r < rounds; ++r) { + double ours_us = 0.0, vendor_us = 0.0; + const bool ours_first = (r % 2 == 0); + for (int arm = 0; arm < 2; ++arm) { + const bool run_ours = (arm == 0) == ours_first; + bench::cooldown(); + if (run_ours) { + ours_us = bench::measure_rotating( + sacrificial_slots, timed_slots, iters, launch_ours_slot); +#if NVFP4_HAS_SCHEDULE + bench::require_clean_placement("timed round"); +#endif + } else { + vendor_us = bench::measure_rotating( + sacrificial_slots, timed_slots, iters, + launch_vendor_slot); + } + } + ours_us_all.push_back(ours_us); + vendor_us_all.push_back(vendor_us); + const bool win = ours_us < vendor_us; + wins += win; + std::printf("ROUND %d/%d order=%s ours=%.3fus (%.3f PF) " + "cublasLt=%.3fus (%.3f PF) ratio=%.4f %s\n", + r + 1, rounds, ours_first ? "ours-first" : "cublas-first", + ours_us, pflops(ours_us), vendor_us, pflops(vendor_us), + vendor_us / ours_us, win ? "WIN" : "LOSS"); + } + + // Comparator-symmetry receipt: both arms' round dispersion, side by side. + // The comparator is a measured object — its power-governor state can mix + // regimes across rounds and windows, and a mixed window can fake or erase + // a small win. An arm's CoV beyond 3x the other's (above the ~0.3% CoV a + // quiet window reads) or beyond 3x that quiet class outright marks the + // window regime-suspect: explain it before trusting the ratio. + auto cov_pct = [](const std::vector& v) { + if (v.size() < 2) return 0.0; + double mean = 0.0; + for (double x : v) mean += x; + mean /= double(v.size()); + if (mean == 0.0) return 0.0; + double ss = 0.0; + for (double x : v) ss += (x - mean) * (x - mean); + return std::sqrt(ss / double(v.size() - 1)) / mean * 100.0; + }; + const double ours_cov = cov_pct(ours_us_all); + const double vendor_cov = cov_pct(vendor_us_all); + const double asym = std::max(ours_cov, vendor_cov) / + std::max(std::min(ours_cov, vendor_cov), 1e-9); + constexpr double kAsymFlag = 3.0; // flag threshold, both clauses + constexpr double kQuietCovPct = 0.3; // a quiet window's round-CoV class + auto unstable = [&](double cov, double sibling) { + return cov > kAsymFlag * kQuietCovPct || + (cov > kAsymFlag * sibling && cov > kQuietCovPct); + }; + std::printf("ARMS %dx%dx%d ours_med=%.3fus cov=%.3f%% | comp_med=%.3fus " + "cov=%.3f%% | asym=%.1fx %s\n", + M, N, K, bench::median(ours_us_all), ours_cov, + bench::median(vendor_us_all), vendor_cov, asym, + unstable(vendor_cov, ours_cov) ? "REGIME=comparator-unstable" + : unstable(ours_cov, vendor_cov) ? "REGIME=ours-unstable" + : "REGIME=quiet"); + + const double ours_med = bench::median(ours_us_all); + const double vendor_med = bench::median(vendor_us_all); + char heading[256]; + char result[256]; + std::snprintf(heading, sizeof(heading), + "%s (%s) | shape %dx%dx%d | rounds %d", + bench::gemm_source_file(), NVFP4_BUILD_NAME, + M, N, K, rounds); + std::snprintf(result, sizeof(result), + "RESULT | ours %.3f PFLOP/s | cuBLASLt %.3f PFLOP/s | " + "%.2f%% of cuBLASLt | wins %d/%d", + pflops(ours_med), pflops(vendor_med), + 100.0 * vendor_med / ours_med, wins, rounds); + bench::print_result_box(heading, result); + + free_benchmark_slots(); + return 0; +} diff --git a/code/labs/fast_cu/upstream/gb300/nvfp4/run_bench_modes.sh b/code/labs/fast_cu/upstream/gb300/nvfp4/run_bench_modes.sh new file mode 100755 index 000000000..df36916ab --- /dev/null +++ b/code/labs/fast_cu/upstream/gb300/nvfp4/run_bench_modes.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +set -euo pipefail + +usage() { + cat <<'EOF' +usage: run_bench_modes.sh MODE [M N K] + +MODE: current | clock-pinned | sustained | triton | all +Defaults: M=N=K=8192, CUDA_VISIBLE_DEVICES=0 +EOF +} + +mode=${1:-all} +if [[ $# -gt 0 ]]; then shift; fi +M=${1:-8192} +N=${2:-8192} +K=${3:-8192} +if [[ $# -gt 3 ]]; then usage >&2; exit 2; fi + +case "$mode" in + current|clock-pinned|sustained|triton|all) ;; + *) usage >&2; exit 2 ;; +esac + +export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0} +gpu_id=${CUDA_VISIBLE_DEVICES%%,*} +if [[ ! "$gpu_id" =~ ^[0-9A-Za-z:.-]+$ ]]; then + echo "invalid CUDA_VISIBLE_DEVICES entry: $gpu_id" >&2 + exit 2 +fi + +root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) +cd "$root" +make nvfp4-r9 +bin=./out/nvfp4-r9 + +pin_active=0 +reset_clock() { + if [[ $pin_active == 1 ]]; then + sudo -n nvidia-smi --id="$gpu_id" -rgc >/dev/null + pin_active=0 + fi +} +trap reset_clock EXIT +trap 'exit 130' INT TERM + +run_current() { + "$bin" "$M" "$N" "$K" \ + --bench-mode current --rounds 5 --cooldown-sec 8 +} + +run_clock_pinned() { + sudo -n nvidia-smi --id="$gpu_id" -lgc 1305,1305 >/dev/null + pin_active=1 + local rc=0 + "$bin" "$M" "$N" "$K" \ + --bench-mode clock-pinned --rounds 5 --cooldown-sec 8 || rc=$? + reset_clock + return "$rc" +} + +run_sustained() { + "$bin" "$M" "$N" "$K" \ + --bench-mode sustained --rounds 3 --cooldown-sec 0 +} + +run_triton() { + "$bin" "$M" "$N" "$K" \ + --bench-mode triton --rounds 1 --cooldown-sec 0 +} + +case "$mode" in + current) run_current ;; + clock-pinned) run_clock_pinned ;; + sustained) run_sustained ;; + triton) run_triton ;; + all) + run_current + run_clock_pinned + run_sustained + run_triton + ;; +esac diff --git a/code/labs/fast_cu/upstream/h100/logs.txt b/code/labs/fast_cu/upstream/h100/logs.txt new file mode 100644 index 000000000..14fca11e1 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/logs.txt @@ -0,0 +1,8 @@ +nvcc -std=c++17 -O3 -DNDEBUG --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math -Xcompiler=-fPIE -Xcompiler=-Wno-psabi -Xcompiler=-fno-strict-aliasing --generate-code arch=compute_90a,code=[compute_90a,sm_90a] -lcublas -lcuda -lineinfo -I/usr/local/cuda-12.6/include matmul.cu -o out/matmul +examples/matmul/matmul_8.cuh(488): error: identifier "c" is undefined + init_barrier(&fullB[c], 0, 1); + ^ + +1 error detected in the compilation of "matmul.cu". +make: *** [Makefile:32: matmul] Error 2 +Build failed. diff --git a/code/labs/fast_cu/upstream/h100/matmul.cu b/code/labs/fast_cu/upstream/h100/matmul.cu new file mode 100644 index 000000000..eac907f14 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul.cu @@ -0,0 +1,244 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +typedef __nv_bfloat16 bf16; +#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) + +void cudaCheck(cudaError_t error, const char *file, int line) { + if (error != cudaSuccess) { + printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line, + cudaGetErrorString(error)); + exit(1); + } +} +#define cudaCheck(err) (cudaCheck(err, __FILE__, __LINE__)) + +#include "matmul/matmul_1.cuh" +#include "matmul/matmul_2.cuh" +#include "matmul/matmul_3.cuh" +#include "matmul/matmul_4.cuh" +#include "matmul/matmul_5.cuh" +#include "matmul/matmul_6.cuh" +#include "matmul/matmul_7.cuh" +#include "matmul/matmul_8.cuh" +#include "matmul/matmul_9.cuh" +#include "matmul/matmul_10.cuh" +#include "matmul/matmul_11.cuh" +#include "matmul/matmul_12.cuh" + +std::default_random_engine generator(69); +cublasHandle_t cublas_handle; +void runCublasGemmBF16(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C) { + float alpha = 1, beta = 0; + // C(column major) = A(row major) * B(column major) + cublasStatus_t status = cublasGemmEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, M, N, K, &alpha, A, CUDA_R_16BF, + N, B, CUDA_R_16BF, K, &beta, C, CUDA_R_16BF, N, CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT); + + if (status != CUBLAS_STATUS_SUCCESS) { + std::cout << "CUBLAS error: " << status << std::endl; + exit(1); + } +} + +void run_kernel(int kernel_num, int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB = nullptr) { + switch (kernel_num) { + case 0: + runCublasGemmBF16(M, N, K, A, B, C); + break; + case 1: + runKernel1(M, N, K, A, B, C); + break; + case 2: + runKernel2(M, N, K, A, B, C); + break; + case 3: + runKernel3(M, N, K, A, B, C, DB); + break; + case 4: + runKernel4(M, N, K, A, B, C, DB); + break; + case 5: + runKernel5(M, N, K, A, B, C, DB); + break; + case 6: + runKernel6(M, N, K, A, B, C, DB); + break; + case 7: + runKernel7(M, N, K, A, B, C, DB); + break; + case 8: + runKernel8(M, N, K, A, B, C, DB); + break; + case 9: + runKernel9(M, N, K, A, B, C, DB); + break; + case 10: + runKernel10(M, N, K, A, B, C, DB); + break; + case 11: + runKernel11(M, N, K, A, B, C, DB); + break; + case 12: + runKernel12(M, N, K, A, B, C, DB); + break; + } +} +int yo = 0; +void randomize_matrix(bf16 *mat, int N) { + std::normal_distribution distribution(0, 1); + for (int i = 0; i < N; i++) { + mat[i] = distribution(generator); + } + ++yo; +} + +bool verify_matrix(bf16 *matRef, bf16 *matOut, int N) { + double diff = 0.0; + int i; + for (i = 0; i < N; i++) { + int r = i / 8192, c = i % 8192; + int it = c*8192+r; + diff = std::fabs(__bfloat162float(matRef[i] - matOut[i])); + if (diff > 0.1) { + printf("Divergence! Should %5.2f, Is %5.2f (Diff %5.2f) at %d\n", + __bfloat162float(matRef[i]), __bfloat162float(matOut[i]), diff, i); + return false; + } + } + return true; +} + +__global__ void warmupKernel() { + __shared__ int s[100]; + s[0] += s[1]; +} + +int main() { + warmupKernel<<<1024, 1024>>>(); + + cublasCreate(&cublas_handle); + float elapsed_time; + cudaEvent_t start, stop; + cudaEventCreate(&start); + cudaEventCreate(&stop); + + long max_size = 8192; + long m = max_size, n = max_size, k = max_size; + + bf16 *A = nullptr, *B = nullptr, *C = nullptr, + *C_ref = nullptr; // host matrices + bf16 *dA = nullptr, *dB = nullptr, *dC = nullptr, + *dC_ref = nullptr; // device matrices + + int *DB = nullptr; int *dDB = nullptr; + + A = (bf16 *)malloc(sizeof(bf16) * max_size * max_size); + B = (bf16 *)malloc(sizeof(bf16) * max_size * max_size); + C = (bf16 *)malloc(sizeof(bf16) * max_size * max_size); + C_ref = (bf16 *)malloc(sizeof(bf16) * max_size * max_size); + DB = (int *)malloc(sizeof(int) * max_size * 128); + cudaCheck(cudaMalloc((void **)&dDB, sizeof(int) * max_size * 128)); + + randomize_matrix(A, max_size * max_size); + randomize_matrix(B, max_size * max_size); + randomize_matrix(C, max_size * max_size); + + cudaCheck(cudaMalloc((void **)&dA, sizeof(bf16) * max_size * max_size)); + cudaCheck(cudaMalloc((void **)&dB, sizeof(bf16) * max_size * max_size)); + cudaCheck(cudaMalloc((void **)&dC, sizeof(bf16) * max_size * max_size)); + cudaCheck(cudaMalloc((void **)&dC_ref, sizeof(bf16) * max_size * max_size)); + + cudaCheck(cudaMemcpy(dA, A, sizeof(bf16) * max_size * max_size, + cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpy(dB, B, sizeof(bf16) * max_size * max_size, + cudaMemcpyHostToDevice)); + + int repeat_times = 8; + bool run_verif = true; + for (int kernel_num : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}) { + // for (int kernel_num : {0, 11}) { + // Give the GPU some rest to avoid thermal throttling + sleep(5); + std::cout << "KERNEL " << kernel_num << std::endl; + // Verify against cuBLAS. Also serves as a warmup step. + if (run_verif) { + memset(C, 0, sizeof(bf16) * max_size * max_size); + cudaCheck(cudaMemcpy(dC, C, sizeof(bf16) * max_size * max_size, cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpy(dC_ref, C, sizeof(bf16) * max_size * max_size, cudaMemcpyHostToDevice)); + memset(DB, ~0, sizeof(int) * max_size * 128); + cudaCheck(cudaMemcpy(dDB, DB, sizeof(int) * max_size * 128, + cudaMemcpyHostToDevice)); + run_kernel(0, m, n, k, dA, dB, dC_ref); // cuBLAS + run_kernel(kernel_num, m, n, k, dA, dB, dC, dDB); // Executes the kernel, modifies the result matrix + cudaCheck(cudaDeviceSynchronize()); + cudaCheck(cudaGetLastError()); // Check for async errors during kernel run + cudaMemcpy(C, dC, sizeof(bf16) * max_size * max_size, cudaMemcpyDeviceToHost); + cudaMemcpy(C_ref, dC_ref, sizeof(bf16) * max_size * max_size, cudaMemcpyDeviceToHost); + + if (kernel_num > 1 && !verify_matrix(C_ref, C, m * n)) { + std::cout << "~~~~~~~~~~~~~~~~ Failed to pass the correctness verification against cuBLAS. ~~~~~~~~~~~~~~~~" << std::endl; + printf("%f\n", __bfloat162float(C_ref[m])); + } + + cudaMemcpy(DB, dDB, sizeof(int) * max_size * 8, cudaMemcpyDeviceToHost); + + int i = 0; + long sumLoad = 0, cntLoad = 0; + long sumCompute = 0, cntCompute = 0; + long sumStore = 0, cntStore = 0; + int times = 0; + while (DB[i] != ~0) { + sumLoad += DB[i], cntLoad += DB[i + 1]; + sumCompute += DB[i + 2], cntCompute += DB[i + 3]; + sumStore += DB[i + 4], cntStore += DB[i + 5]; + i += 6; + times++; + } + if (times > 0) { + printf("Load: %f, Compute: %f, Store: %f, Datapoints: %d\n", (sumLoad + .0) / cntLoad, (sumCompute + .0) / cntCompute, (sumStore + .0) / cntStore, times); + } + + } + + // Benchmark + cudaEventRecord(start); + for (int j = 0; j < repeat_times; j++) { + run_kernel(kernel_num, m, n, k, dA, dB, dC); + } + cudaEventRecord(stop); + cudaEventSynchronize(start); + cudaEventSynchronize(stop); + cudaEventElapsedTime(&elapsed_time, start, stop); + + long flops = (2LL * m) * (n * k); + printf( + "Average elapsed time: (%7.6f) s, performance: (%7.1f) TFLOPS. size: (%ld).\n\n", + elapsed_time / 1000.0 / repeat_times, + (repeat_times * flops * 1e-9) / elapsed_time, m); + } + + // Free up CPU and GPU space + free(A); + free(B); + free(C); + free(C_ref); + cudaFree(dA); + cudaFree(dB); + cudaFree(dC); + cudaFree(dC_ref); + return 0; +}; diff --git a/code/labs/fast_cu/upstream/h100/matmul/cublaslt_matmul.cu b/code/labs/fast_cu/upstream/h100/matmul/cublaslt_matmul.cu new file mode 100644 index 000000000..b1322bbc9 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/cublaslt_matmul.cu @@ -0,0 +1,74 @@ +#include + +// Random code snippets to use while development +const size_t cublaslt_workspace_size = 32 * 1024 * 1024; +void* cublaslt_workspace = NULL; +cublasLtHandle_t cublaslt_handle; + +void initCublasLt() { + cublasLtCreate(&cublaslt_handle); + cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); +} + +void runCublasMatmulBF16(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C) { + // check alignment (some modes work unaligned but it always best to be aligned for performance) + if(((uintptr_t)A % 16) != 0 || ((uintptr_t)B % 16) != 0 || ((uintptr_t)C % 16) != 0) { + printf("All cuBLASLt pointers must be aligned!\n"); + exit(EXIT_FAILURE); + } + + // create the operation descriptor + cublasLtMatmulDesc_t operationDesc; + cublasLtMatmulDescCreate(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F); + + int returnedResults = 0; + cublasLtMatmulPreference_t preference; + cublasLtMatmulHeuristicResult_t heuristic; + + cublasOperation_t opNoTranspose = CUBLAS_OP_N; + cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opNoTranspose, sizeof(opNoTranspose)); + cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opNoTranspose, sizeof(opNoTranspose)); + + // define matrix layouts + cublasLtMatrixLayout_t ALayout; + cublasLtMatrixLayout_t BLayout; + cublasLtMatrixLayout_t CLayout; + cublasLtMatrixLayoutCreate(&ALayout, CUDA_R_16BF, M, K, M); + cublasLtMatrixLayoutCreate(&BLayout, CUDA_R_16BF, K, N, K); + + // cuBLASLt requires C in FP8 mode to be BF16 or FP32... (sigh) + cublasLtMatrixLayoutCreate(&CLayout, CUDA_R_16BF, M, N, M); + + // create a preference handle with specified max workspace + cublasLtMatmulPreferenceCreate(&preference); + cublasLtMatmulPreferenceSetAttribute(preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTE, + &cublaslt_workspace_size, sizeof(cublaslt_workspace_size))); + + // set scale type to FP32 (needs to be FP16 if and only if using CUBLAS_COMPUTE_16F, so it's FP32 even for FP8!) + cublasDataType_t scale_type = CUDA_R_32F; + cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type)); + + // find a suitable algorithm (cached internally so shouldn't take much CPU time in practice) + cublasLtMatmulAlgoGetHeuristic(cublaslt_handle, operationDesc, ALayout, BLayout, CLayout, CLayout, + preference, 1, &heuristic, &returnedResults); + if (returnedResults == 0) { + printf("No cuBLASLt algorithm: m: %d, n: %d, k: %d", N, M, K); + exit(EXIT_FAILURE); + } + + // set whether to accumulate (i.e. D += C) or not - note this isn't considered in algorithm selection (?!) + float alpha = 1, beta = 0; + + // call the matmul + cublasLtMatmul(cublaslt_handle, operationDes, + &alpha, A, ALayout, B, BLayout, &beta, C, CLayout, C, CLayout, + &heuristic.algo, cublaslt_workspace, cublaslt_workspace_size, 0)); + + // cleanups + cublasLtMatmulPreferenceDestroy(preference); + cublasLtMatmulDescDestroy(operationDesc); + cublasLtMatrixLayoutDestroy(ALayout); + cublasLtMatrixLayoutDestroy(BLayout); + cublasLtMatrixLayoutDestroy(CLayout); + cudaCheck(cudaGetLastError()); + } \ No newline at end of file diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_1.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_1.cuh new file mode 100644 index 000000000..60567db81 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_1.cuh @@ -0,0 +1,180 @@ +#include +#include + +template +__global__ void __launch_bounds__(NUM_THREADS) + sgemmWarptiling(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C) { + const uint cRow = blockIdx.y; + const uint cCol = blockIdx.x; + + // Placement of the warp in the threadblock tile + const uint warpIdx = threadIdx.x / 32; // the warp this thread is in + const uint warpCol = warpIdx % (BN / WN); + const uint warpRow = warpIdx / (BN / WN); + + // size of the warp subtile + constexpr uint WMITER = (WM * WN) / (32 * TM * TN * WNITER); + constexpr uint WSUBM = WM / WMITER; // 64/2=32 + constexpr uint WSUBN = WN / WNITER; // 32/2=16 + + // Placement of the thread in the warp subtile + const uint threadIdxInWarp = threadIdx.x % 32; // [0, 31] + const uint threadColInWarp = threadIdxInWarp % (WSUBN / TN); // i%(16/4) + const uint threadRowInWarp = threadIdxInWarp / (WSUBN / TN); // i/4 + + // allocate space for the current blocktile in SMEM + __shared__ bf16 As[BM * BK]; + __shared__ bf16 Bs[BK * BN]; + + // Move blocktile to beginning of A's row and B's column + A += cRow * BM * K; + B += cCol * BN; + // Move C_ptr to warp's output tile + C += (cRow * BM + warpRow * WM) * N + cCol * BN + warpCol * WN; + + // calculating the indices that this thread will load into SMEM + // we'll load 128bit / 32bit = 4 elements per thread at each step + const uint innerRowA = threadIdx.x / (BK / 4); + const uint innerColA = threadIdx.x % (BK / 4); + constexpr uint rowStrideA = (NUM_THREADS * 4) / BK; + const uint innerRowB = threadIdx.x / (BN / 4); + const uint innerColB = threadIdx.x % (BN / 4); + constexpr uint rowStrideB = NUM_THREADS / (BN / 4); + + // allocate thread-local cache for results in registerfile + float threadResults[WMITER * TM * WNITER * TN] = {0.0}; + // we cache into registers on the warptile level + bf16 regM[WMITER * TM] = {0.0}; + bf16 regN[WNITER * TN] = {0.0}; + + // outer-most loop over block tiles + for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) { + // Load from A + for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) { + bf16 tmp[4]; + const float2 x = *reinterpret_cast(&A[(innerRowA + offset) * K + innerColA * 4]); + memcpy(&tmp[0], &x, sizeof(bf16) * 4); + As[(innerColA * 4 + 0) * BM + innerRowA + offset] = tmp[0]; + As[(innerColA * 4 + 1) * BM + innerRowA + offset] = tmp[1]; + As[(innerColA * 4 + 2) * BM + innerRowA + offset] = tmp[2]; + As[(innerColA * 4 + 3) * BM + innerRowA + offset] = tmp[3]; + } + + // Load from B + for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) { + reinterpret_cast( + &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] = + reinterpret_cast( + &B[(innerRowB + offset) * N + innerColB * 4])[0]; + } + __syncthreads(); + + for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) { + // populate registers for whole warptile + for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) { + for (uint i = 0; i < TM; ++i) { + regM[wSubRowIdx * TM + i] = + As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM + + threadRowInWarp * TM + i]; + } + } + for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) { + for (uint i = 0; i < TN; ++i) { + regN[wSubColIdx * TN + i] = + Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN + + threadColInWarp * TN + i]; + } + } + + // execute warptile matmul + for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) { + for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) { + // calculate per-thread results + for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) { + for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) { + threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) + + (wSubColIdx * TN) + resIdxN] += + __bfloat162float(regM[wSubRowIdx * TM + resIdxM]) * + __bfloat162float(regN[wSubColIdx * TN + resIdxN]); + } + } + } + } + } + + A += BK; // move BK columns to right + B += BK * N; // move BK rows down + __syncthreads(); + } + + // write out the results + for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) { + for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) { + // move C pointer to current warp subtile + bf16 *C_interim = C + (wSubRowIdx * WSUBM) * N + wSubColIdx * WSUBN; + for (uint resIdxM = 0; resIdxM < TM; resIdxM += 1) { + for (uint resIdxN = 0; resIdxN < TN; resIdxN += 4) { + bf16 tmp[4]; + float2 x; + const int i = (wSubRowIdx * TM + resIdxM) * (WNITER * TN) + + wSubColIdx * TN + resIdxN; + tmp[0] = threadResults[i + 0]; + tmp[1] = threadResults[i + 1]; + tmp[2] = threadResults[i + 2]; + tmp[3] = threadResults[i + 3]; + memcpy(&x, &tmp[0], sizeof(bf16) * 4); + // write + *reinterpret_cast(&C_interim[(threadRowInWarp * TM + resIdxM) * N + threadColInWarp * TN + resIdxN]) = x; + } + } + } + } +} + +void runKernel1(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C) { + const uint K10_NUM_THREADS = 128; + const uint K10_BN = 128; + const uint K10_BM = 128; + const uint K10_BK = 16; + const uint K10_WN = 64; + const uint K10_WM = 64; + const uint K10_WNITER = 4; + const uint K10_TN = 4; + const uint K10_TM = 8; + dim3 blockDim(K10_NUM_THREADS); + + constexpr uint NUM_WARPS = K10_NUM_THREADS / 32; + + // warptile in threadblocktile + static_assert((K10_BN % K10_WN == 0) and (K10_BM % K10_WM == 0)); + static_assert((K10_BN / K10_WN) * (K10_BM / K10_WM) == NUM_WARPS); + + // threads in warpsubtile + static_assert((K10_WM * K10_WN) % (32 * K10_TM * K10_TN * K10_WNITER) == 0); + constexpr uint K10_WMITER = + (K10_WM * K10_WN) / (32 * K10_TM * K10_TN * K10_WNITER); + // warpsubtile in warptile + static_assert((K10_WM % K10_WMITER == 0) and (K10_WN % K10_WNITER == 0)); + + static_assert((K10_NUM_THREADS * 4) % K10_BK == 0, + "NUM_THREADS*4 must be multiple of K9_BK to avoid quantization " + "issues during GMEM->SMEM tiling (loading only parts of the " + "final row of Bs during each iteraion)"); + static_assert((K10_NUM_THREADS * 4) % K10_BN == 0, + "NUM_THREADS*4 must be multiple of K9_BN to avoid quantization " + "issues during GMEM->SMEM tiling (loading only parts of the " + "final row of As during each iteration)"); + static_assert(K10_BN % (16 * K10_TN) == 0, + "BN must be a multiple of 16*TN to avoid quantization effects"); + static_assert(K10_BM % (16 * K10_TM) == 0, + "BM must be a multiple of 16*TM to avoid quantization effects"); + static_assert((K10_BM * K10_BK) % (4 * K10_NUM_THREADS) == 0, + "BM*BK must be a multiple of 4*256 to vectorize loads"); + static_assert((K10_BN * K10_BK) % (4 * K10_NUM_THREADS) == 0, + "BN*BK must be a multiple of 4*256 to vectorize loads"); + + dim3 gridDim(CEIL_DIV(N, K10_BN), CEIL_DIV(M, K10_BM)); + sgemmWarptiling + <<>>(M, N, K, A, B, C); +} \ No newline at end of file diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_10.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_10.cuh new file mode 100644 index 000000000..4e7a31f71 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_10.cuh @@ -0,0 +1,621 @@ + +namespace M10 { + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +__device__ uint64_t make_smem_desc(bf16* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; +} + + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +template +__device__ void warpgroup_wait() { + static_assert(N >= 0 && N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +template +__host__ static inline CUtensorMap create_tensor_map(bf16* gmem_ptr, int global_height, int global_width) { + CUtensorMap tma_map; + void* gmem_address = (void*)gmem_ptr; + static_assert(BlockMinorSize >= 64); + assert(global_width % 64 == 0); + uint64_t gmem_prob_shape[5] = {64, (uint64_t)global_height, (uint64_t)global_width/64, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16) * global_width, 64*sizeof(bf16), 0, 0, 0}; + uint32_t smem_box_shape[5] = {64, uint32_t(BlockMajorSize), uint32_t(BlockMinorSize/64), 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + &tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, gmem_address, gmem_prob_shape, + gmem_prob_stride, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle ? CU_TENSOR_MAP_SWIZZLE_128B : CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); + return tma_map; +} + +CUtensorMap d_tma_map_A; +CUtensorMap d_tma_map_B; +CUtensorMap d_tma_map_C; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__device__ __forceinline__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + + +template +__device__ __forceinline__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma64(float d[4][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " %34, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma32(float d[2][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}, " + " %16," + " %17," + " %18, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 32 || WGMMA_N == 64 || WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 208 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128(d, sA, sB); + if constexpr (WGMMA_N == 64) + wgmma64(d, sA, sB); + if constexpr (WGMMA_N == 32) + wgmma32(d, sA, sB); +} + +template +__device__ void warpgroup_reg_alloc() { + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +__device__ void warpgroup_reg_dealloc() { + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +__device__ static __forceinline__ void init_barrier(uint64_t* bar, int thread_count, int transaction_count) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); +} + +__device__ static __forceinline__ void expect_bytes(uint64_t* bar, uint32_t bytes) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + :: "r"(bar_ptr), "r"(bytes)); +} + +__device__ static inline void load_async(bf16 *dst, void const* src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64) + : "memory" + ); +} + +__device__ static inline void store_async(void const* dst_tma_map, bf16 *src, int global_col_idx, int global_row_idx) { + uint64_t tma_ptr = reinterpret_cast(dst_tma_map); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(src)); + + asm volatile ( + "cp.async.bulk.tensor.3d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx / 64) + : "memory" + ); +} + +__device__ static __forceinline__ void wait(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static __forceinline__ void arrive(uint64_t* bar, uint32_t count=1) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); +} + +__device__ static __forceinline__ void wait_cluster(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static inline void load_async_multicast(bf16 *dst, void const* src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx, uint16_t cluster_mask) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5}], [%2], %6;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64), "h"(cluster_mask) + : "memory" + ); +} + +__device__ void arrive_cluster(uint64_t* bar, uint32_t cta_id, uint32_t count=1) { + uint32_t smem_addr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + : + : "r"(smem_addr), "r"(cta_id), "r"(count)); +} + +template +struct Schedule; + +template +struct Schedule<1, NUM_SM, BM, BN, TM, TN> { + int block; + int it; + int total_blocks_m, total_blocks_n; + + __device__ __forceinline__ Schedule(int M, int N, int _block) { + block = _block; + it = 0; + total_blocks_m = CEIL_DIV(M, BM); + total_blocks_n = CEIL_DIV(N, BN); + assert(CEIL_DIV(M, BM)%TM == 0 && total_blocks_n%TN == 0); + } + + __device__ __forceinline__ bool next(int &block_m, int& block_n) { + int num = it*NUM_SM + block; + if (num >= total_blocks_m*total_blocks_n) {return false;} + + int cur_tile = num / (TM*TN); + int cur_tile_pos = num % (TM*TN); + block_m = TM*(cur_tile / (total_blocks_n/TN)); + block_n = TN*(cur_tile % (total_blocks_n/TN)); + block_m += cur_tile_pos / TN; + block_n += cur_tile_pos % TN; + ++it; + return true; + } +}; + +template +struct SMem { + alignas(128) bf16 A[BM*BK*QSIZE]; + alignas(128) bf16 B[BK*BN*QSIZE]; + alignas(128) bf16 C[BN*BM]; + alignas(8) uint64_t full[QSIZE], empty[QSIZE]; +}; + +template +__global__ __launch_bounds__(NUM_THREADS) void __cluster_dims__(CLUSTER_M * CLUSTER_N, 1, 1) matmulKernel10(int M, int N, int K, const __grid_constant__ CUtensorMap tensorMapC, const __grid_constant__ CUtensorMap tensorMapA, const __grid_constant__ CUtensorMap tensorMapB) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int num_consumers = (NUM_THREADS / 128) - 1; + constexpr int B_WG_M = BM / num_consumers; + constexpr int CLUSTERS = CLUSTER_M * CLUSTER_N; + assert((M / BM) % CLUSTER_M == 0); + assert((N / BN) % CLUSTER_N == 0); + + extern __shared__ __align__(128) uint8_t smem[]; + SMem &s = *reinterpret_cast*>(smem); + bf16 *sA = s.A, *sB = s.B, *sC = s.C; + uint64_t *full = s.full, *empty = s.empty; + + uint32_t rank; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(rank) :); + + const int num_blocks_k = K / BK; + int wg_idx = threadIdx.x / 128; + int tid = threadIdx.x % 128; + + if (threadIdx.x == 0) { + for (int i = 0; i < QSIZE; ++i) { + init_barrier(&full[i], 0, 1); + init_barrier(&empty[i], 0, num_consumers*CLUSTERS); + } + } + asm volatile("barrier.cluster.arrive;\n" : :); + asm volatile("barrier.cluster.wait;\n" : :); + + Schedule<1, NUM_SM/CLUSTERS, BM*CLUSTER_M, BN*CLUSTER_N, 16/CLUSTER_M, 8/CLUSTER_N> schedule(M, N, rank); + + asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(rank) :); + uint32_t rank_m = rank / CLUSTER_N; + uint32_t rank_n = rank % CLUSTER_N; + + // Producer + if (wg_idx == 0) { + constexpr int num_regs = (num_consumers <= 2 ? 24 : 32); + warpgroup_reg_dealloc(); + if (tid == 0) { + int p = 0; + int qidx = 0; + uint32_t col_mask = 0; + for (int i = 0; i < CLUSTER_M; ++i) { + col_mask |= (1 << (i * CLUSTER_N)); + } + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) { qidx = 0; p ^= 1;} + wait(&empty[qidx], p); + + expect_bytes(&full[qidx], (BK*BN+BK*BM)*sizeof(bf16)); + if constexpr (CLUSTER_N > 1) { + uint32_t mask = ((1 << CLUSTER_N) - 1) << (rank_m * CLUSTER_N); + if (rank_n == 0) { + load_async_multicast(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM, mask); + } + } else { + load_async(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM); + } + + if constexpr (CLUSTER_M > 1) { + if (rank_m == 0) { + load_async_multicast(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN, col_mask << rank_n); + } + } else { + load_async(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN); + } + } + } + } + } else { + constexpr int num_regs = (num_consumers == 1 ? 256 : (num_consumers == 2 ? 240 : 160)); + warpgroup_reg_alloc(); + float d[B_WG_M/WGMMA_M][WGMMA_N/16][8]; + --wg_idx; + for (int qidx = 0; qidx < QSIZE; ++qidx) { + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + } + int p = 0; + int qidx = 0; + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + { + if (qidx == QSIZE) {qidx = 0; p ^= 1; }; + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + { + wgmma(d[m_it], &wgmma_sA[0], &wgmma_sB[0]); + #pragma unroll + for (int k_it = 1; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + #pragma unroll + for (int bk = 64; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + ++qidx; + } + for (int block_k_iter = 1; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) {qidx = 0; p ^= 1; }; + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + #pragma unroll + for (int bk = 0; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + } + + asm volatile("cp.async.bulk.wait_group 0;"); + + int lane = tid % 32, warp = tid / 32; + int row = warp*16 + lane / 4; + + bf16* block_sC = sC + wg_idx*B_WG_M*BN; + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + int yo = m_it*WGMMA_M; + #pragma unroll + for (int w = 0; w < WGMMA_N; w+=16) { + int col = w + 2*(tid % 4); + #define ST(i, j, v) block_sC[(j)*B_WG_M + (i) + yo] = v + + ST(row, col, d[m_it][w/16][0]); + ST(row+8, col, d[m_it][w/16][2]); + + + ST(row, col+1, d[m_it][w/16][1]); + ST(row+8, col+1, d[m_it][w/16][3]); + + + ST(row, col+8, d[m_it][w/16][4]); + ST(row+8, col+8, d[m_it][w/16][6]); + + + ST(row, col+9, d[m_it][w/16][5]); + ST(row+8, col+9, d[m_it][w/16][7]); + + + // #undef IDX + #undef ST + } + } + asm volatile("bar.sync 10, 256;\n"); + if (threadIdx.x == 128) { + store_async(&tensorMapC, (bf16*)&sC[0], num_block_m*BM, num_block_n*BN); + asm volatile("cp.async.bulk.commit_group;"); + } + } + } +} + +void runKernel10(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 128; + constexpr int BN = 256; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128*3; + constexpr int QSIZE = 3; + constexpr int CLUSTER_M = 2; + constexpr int CLUSTER_N = 1; + constexpr int NUM_SM = 128; + static_assert(NUM_SM % (CLUSTER_M*CLUSTER_N) == 0); + + if (_prev_m != M) { + d_tma_map_A = create_tensor_map(A, M, K); + d_tma_map_B = create_tensor_map(B, N, K); + d_tma_map_C = create_tensor_map(C, N, M); + _prev_m = M; + _prev_n = N; + _prev_k = K; + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = matmulKernel10; + constexpr size_t sMemSize = sizeof(SMem); + static_assert(sMemSize < 256 * 1024); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<>>(M, N, K, d_tma_map_C, d_tma_map_A, d_tma_map_B); +} + +} // namespace M10 + +using M10::runKernel10; + \ No newline at end of file diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_11.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_11.cuh new file mode 100644 index 000000000..dea974a7c --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_11.cuh @@ -0,0 +1,639 @@ + +namespace M11 { + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { + return (((x) & 0x3FFFF) >> 0x4); +} + +// Descriptor for a shared memory matrix. +// Implementation is derived from PTX guide: https://docs.nvidia.com/cuda/parallel-thread-execution/#matrix-descriptor-format +__device__ uint64_t make_smem_desc(bf16* ptr) { + // Convert shared memory pointer to integer + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; +} + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_wait() { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); +} + +template +__host__ static inline CUtensorMap create_tensor_map(bf16* gmem_ptr, int global_height, int global_width) { + CUtensorMap tma_map; + void* gmem_address = (void*)gmem_ptr; + static_assert(BlockMinorSize >= 64); + assert(global_width % 64 == 0); + uint64_t gmem_prob_shape[5] = {64, (uint64_t)global_height, (uint64_t)global_width/64, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16) * global_width, 64*sizeof(bf16), 0, 0, 0}; + uint32_t smem_box_shape[5] = {64, uint32_t(BlockMajorSize), uint32_t(BlockMinorSize/64), 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + &tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, gmem_address, gmem_prob_shape, + gmem_prob_stride, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle ? CU_TENSOR_MAP_SWIZZLE_128B : CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); + return tma_map; +} + +CUtensorMap d_tma_map_A; +CUtensorMap d_tma_map_B; +CUtensorMap d_tma_map_C; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__device__ __forceinline__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + + +template +__device__ __forceinline__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128(d, sA, sB); +} + +template +__device__ void warpgroup_reg_alloc() { + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +__device__ void warpgroup_reg_dealloc() { + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +__device__ static __forceinline__ void init_barrier(uint64_t* bar, int thread_count, int transaction_count) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); +} + +__device__ static __forceinline__ void expect_bytes(uint64_t* bar, uint32_t bytes) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + :: "r"(bar_ptr), "r"(bytes)); +} + +__device__ static inline void load_async(bf16 *dst, void const* src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + // We use 3d tiling to load from GMEM to SMEM. 2d tiling only works for tiles <= 64 columns. + // For larger tiles, we need 3d tiling: + // First dimension: 64 (columns) + // Second dimension: Row + // Third dimension: Column / 64 + // I spent a lot of time figuring this out. Hope this gets documented at some point. + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64) + : "memory" + ); +} + +__device__ static inline void store_async(void const* dst_tma_map, bf16 *src, int global_col_idx, int global_row_idx) { + uint64_t tma_ptr = reinterpret_cast(dst_tma_map); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(src)); + + asm volatile ( + "cp.async.bulk.tensor.3d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx / 64) + : "memory" + ); +} + +__device__ static __forceinline__ void wait(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + // Call mbarrier.try_wait in a while loop till it returns true. + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static __forceinline__ void arrive(uint64_t* bar, uint32_t count=1) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); +} + +__device__ void arrive_cluster(uint64_t* bar, uint32_t cta_id, uint32_t count=1) { + uint32_t smem_addr = static_cast(__cvta_generic_to_shared(bar)); + // Call arrive on barrier stored in SM number `cta_id` + // Use mapa.shared instruction to access other SM's barrier. + asm volatile( + "{\n" + ".reg .b32 remAddr32;\n" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n" + "mbarrier.arrive.shared::cluster.b64 _, [remAddr32], %2;\n" + "}" + :: "r"(smem_addr), "r"(cta_id), "r"(count)); +} + +__device__ static inline void load_async_multicast(bf16 *dst, void const* src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx, uint16_t cluster_mask) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5}], [%2], %6;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64), "h"(cluster_mask) + : "memory" + ); +} + +template +struct Schedule; + +constexpr int SPACE_LEN = 128; +int *_dspace; + +template +struct Schedule<2, NUM_SM, BM, BN, TM, TN> { + int it; + int *space; + + __device__ __forceinline__ Schedule(int M, int N, int block, int *_space) { + it = 0; + space = _space; + } + + __device__ __forceinline__ bool next(int &block_m, int& block_n) { + if (it >= SPACE_LEN) { + return false; + } + int now = space[it]; + if (now == -1) { + return false; + } + block_m = now >> 16; + block_n = (now & ((1<<16)-1)); + ++it; + return true; + } +}; + +template +struct SMem { + alignas(128) bf16 A[BM*BK*QSIZE]; + alignas(128) bf16 B[BK*BN*QSIZE]; + alignas(128) bf16 C[BN*BM]; + alignas(8) uint64_t full[QSIZE], empty[QSIZE]; + int space[SPACE_LEN]; +}; + +template +__global__ __launch_bounds__(NUM_THREADS) void __cluster_dims__(CLUSTER_M * CLUSTER_N, 1, 1) matmulKernel11(int M, int N, int K, const __grid_constant__ CUtensorMap tensorMapC, const __grid_constant__ CUtensorMap tensorMapA, const __grid_constant__ CUtensorMap tensorMapB, int* dspace) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int num_consumers = (NUM_THREADS / 128) - 1; + constexpr int B_WG_M = BM / num_consumers; + constexpr int CLUSTERS = CLUSTER_M * CLUSTER_N; + assert((M / BM) % CLUSTER_M == 0); + assert((N / BN) % CLUSTER_N == 0); + + extern __shared__ __align__(128) uint8_t smem[]; + SMem &s = *reinterpret_cast*>(smem); + bf16 *sA = s.A, *sB = s.B, *sC = s.C; + uint64_t *full = s.full, *empty = s.empty; + int *space = s.space; + + uint32_t rank; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(rank) :); + // Load schedule for this SM + if (threadIdx.x < SPACE_LEN) space[threadIdx.x] = dspace[rank*SPACE_LEN+threadIdx.x]; + + const int num_blocks_k = K / BK; + int wg_idx = threadIdx.x / 128; + int tid = threadIdx.x % 128; + + if (threadIdx.x == 0) { + for (int i = 0; i < QSIZE; ++i) { + init_barrier(&full[i], 0, 1); + init_barrier(&empty[i], 0, num_consumers*CLUSTERS); + } + } + asm volatile("barrier.cluster.arrive;\n" : :); + asm volatile("barrier.cluster.wait;\n" : :); + + + Schedule<2, NUM_SM/CLUSTERS, BM*CLUSTER_M, BN*CLUSTER_N, 16/CLUSTER_M, 8/CLUSTER_N> schedule(M, N, rank, &space[0]); + + asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(rank) :); + uint32_t rank_m = rank / CLUSTER_N; + uint32_t rank_n = rank % CLUSTER_N; + + // Producer + if (wg_idx == 0) { + constexpr int num_regs = (num_consumers <= 2 ? 24 : 32); + warpgroup_reg_dealloc(); + if (tid == 0) { + int p = 0; + int qidx = 0; + uint32_t col_mask = 0; + for (int i = 0; i < CLUSTER_M; ++i) { + col_mask |= (1 << (i * CLUSTER_N)); + } + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) { qidx = 0; p ^= 1;} + wait(&empty[qidx], p); + + expect_bytes(&full[qidx], (BK*BN+BK*BM)*sizeof(bf16)); + if constexpr (CLUSTER_N > 1) { + uint32_t mask = ((1 << CLUSTER_N) - 1) << (rank_m * CLUSTER_N); + if (rank_n == 0) { + load_async_multicast(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM, mask); + } + } else { + load_async(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM); + } + + if constexpr (CLUSTER_M > 1) { + if (rank_m == 0) { + load_async_multicast(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN, col_mask << rank_n); + } + } else { + load_async(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN); + } + } + } + } + } else { + constexpr int num_regs = (num_consumers == 1 ? 256 : (num_consumers == 2 ? 240 : 160)); + warpgroup_reg_alloc(); + float d[B_WG_M/WGMMA_M][WGMMA_N/16][8]; + --wg_idx; + for (int qidx = 0; qidx < QSIZE; ++qidx) { + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + } + int p = 0; + int qidx = 0; + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + { + if (qidx == QSIZE) {qidx = 0; p ^= 1; }; + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + { + wgmma(d[m_it], &wgmma_sA[0], &wgmma_sB[0]); + #pragma unroll + for (int k_it = 1; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + #pragma unroll + for (int bk = 64; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + ++qidx; + } + for (int block_k_iter = 1; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) {qidx = 0; p ^= 1; }; + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + #pragma unroll + for (int bk = 0; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + } + + asm volatile("cp.async.bulk.wait_group 0;"); + + int lane = tid % 32, warp = tid / 32; + int row = warp*16 + lane / 4; + + bf16* block_sC = sC + wg_idx*B_WG_M*BN; + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + int yo = m_it*WGMMA_M; + #pragma unroll + for (int w = 0; w < WGMMA_N; w+=16) { + int col = w + 2*(tid % 4); + #define ST(i, j, v) block_sC[(j)*B_WG_M + (i) + yo] = v + + ST(row, col, d[m_it][w/16][0]); + ST(row+8, col, d[m_it][w/16][2]); + + + ST(row, col+1, d[m_it][w/16][1]); + ST(row+8, col+1, d[m_it][w/16][3]); + + + ST(row, col+8, d[m_it][w/16][4]); + ST(row+8, col+8, d[m_it][w/16][6]); + + + ST(row, col+9, d[m_it][w/16][5]); + ST(row+8, col+9, d[m_it][w/16][7]); + + #undef ST + } + } + // Wait for all 256 consumer threads to reach here + asm volatile("bar.sync 1, 256;\n"); + + if (threadIdx.x == 128) { + store_async(&tensorMapC, (bf16*)&sC[0], num_block_m*BM, num_block_n*BN); + asm volatile("cp.async.bulk.commit_group;"); + } + } + } +} + +// Rotate/flip quadrant appropriately +void rot(int n, int& x, int& y, int rx, int ry) { + if (ry == 0) { + if (rx == 1) { + x = n-1 - x; + y = n-1 - y; + } + // Swap x and y + int t = x; + x = y; + y = t; + } +} + +// Convert distance along curve to (x,y) point +void d2xy(int n, int d, int& x, int& y) { + int rx, ry, s, t = d; + x = y = 0; + for (s = 1; s < n; s *= 2) { + rx = 1 & (t/2); + ry = 1 & (t ^ rx); + rot(s, x, y, rx, ry); + x += s * rx; + y += s * ry; + t /= 4; + } +} + +void createHilbert(int M, int N, int CORES, int *space) { + int dim = (1 << (32 - __builtin_clz(max(M, N) - 1))); + int core = 0; + std::vector v(dim, std::string(dim, '.')); + memset(space, -1, sizeof(int)*CORES*SPACE_LEN); + int FCORES = 64; + int total = 0; + std::vector> pos(CORES, std::vector()); + for (int i = 0; i < dim*dim; ++i) { + int x, y; + d2xy(dim, i, x, y); + if (x < M && y < N) { + assert(loc < SPACE_LEN); + assert(v[x][y] == '.'); + v[x][y] = '*'; + ++total; + pos[core].push_back((x << 16) | y); + ++core; + if (core == FCORES) {core = 0;} + } + } + core = FCORES; + for (int i = 0; i < FCORES; ++i) { + if (pos.back().size() >= pos[0].size()-1) break; + pos[core].push_back(pos[i].back()); + pos[i].pop_back(); + ++core; + if (core == CORES) {core = FCORES;} + } + for (int i = 0; i < CORES; ++i) { + for (int j = 0; j < pos[i].size(); ++j) { + space[i*SPACE_LEN + j] = pos[i][j]; + } + } + assert(total == M*N); +} + +void runKernel11(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 128; + constexpr int BN = 256; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128*3; + constexpr int QSIZE = 3; + constexpr int CLUSTER_M = 2; + constexpr int CLUSTER_N = 1; + constexpr int NUM_SM = 128; + static_assert(NUM_SM % (CLUSTER_M*CLUSTER_N) == 0); + + if (_prev_m != M) { + d_tma_map_A = create_tensor_map(A, M, K); + d_tma_map_B = create_tensor_map(B, N, K); + d_tma_map_C = create_tensor_map(C, N, M); + _prev_m = M; + _prev_n = N; + _prev_k = K; + int *space; + space = (int*)malloc(sizeof(int)*NUM_SM*SPACE_LEN); + createHilbert(CEIL_DIV(M, BM*CLUSTER_M), CEIL_DIV(N, BN*CLUSTER_N), NUM_SM/CLUSTER_M/CLUSTER_N, space); + cudaCheck(cudaMalloc((void **)&_dspace, sizeof(int)*NUM_SM*SPACE_LEN)); + cudaCheck(cudaMemcpy(_dspace, space, sizeof(int)*NUM_SM*SPACE_LEN, cudaMemcpyHostToDevice)); + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = matmulKernel11; + constexpr size_t sMemSize = sizeof(SMem); + static_assert(sMemSize < 256 * 1024); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<>>(M, N, K, d_tma_map_C, d_tma_map_A, d_tma_map_B, _dspace); +} + +} // namespace M11 + +using M11::runKernel11; + \ No newline at end of file diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_12.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_12.cuh new file mode 100644 index 000000000..c7766a7b9 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_12.cuh @@ -0,0 +1,631 @@ + +namespace M12 { + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { + return (((x) & 0x3FFFF) >> 0x4); +} + +// Descriptor for a shared memory matrix. +// Implementation is derived from PTX guide: https://docs.nvidia.com/cuda/parallel-thread-execution/#matrix-descriptor-format +__device__ uint64_t make_smem_desc(bf16* ptr) { + // Convert shared memory pointer to integer + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; +} + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_wait() { + asm volatile("wgmma.wait_group.sync.aligned 0;\n" ::: "memory"); +} + +template +__host__ static inline CUtensorMap create_tensor_map(bf16* gmem_ptr, int global_height, int global_width) { + CUtensorMap tma_map; + void* gmem_address = (void*)gmem_ptr; + static_assert(BlockMinorSize >= 64); + assert(global_width % 64 == 0); + uint64_t gmem_prob_shape[5] = {64, (uint64_t)global_height, (uint64_t)global_width/64, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16) * global_width, 64*sizeof(bf16), 0, 0, 0}; + uint32_t smem_box_shape[5] = {padding ? 72 : 64, uint32_t(BlockMajorSize), uint32_t(BlockMinorSize/64), 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + &tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, gmem_address, gmem_prob_shape, + gmem_prob_stride, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle ? CU_TENSOR_MAP_SWIZZLE_128B : CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); + return tma_map; +} + +CUtensorMap d_tma_map_A; +CUtensorMap d_tma_map_B; +CUtensorMap d_tma_map_C; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__device__ __forceinline__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + + +template +__device__ __forceinline__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128(d, sA, sB); +} + +template +__device__ void warpgroup_reg_alloc() { + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +__device__ void warpgroup_reg_dealloc() { + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +__device__ static __forceinline__ void init_barrier(uint64_t* bar, int thread_count, int transaction_count) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); +} + +__device__ static __forceinline__ void expect_bytes(uint64_t* bar, uint32_t bytes) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + :: "r"(bar_ptr), "r"(bytes)); +} + +__device__ static inline void load_async(bf16 *dst, void const* src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + // We use 3d tiling to load from GMEM to SMEM. 2d tiling only works for tiles <= 64 columns. + // For larger tiles, we need 3d tiling: + // First dimension: 64 (columns) + // Second dimension: Row + // Third dimension: Column / 64 + // I spent a lot of time figuring this out. Hope this gets documented at some point. + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64) + : "memory" + ); +} + +__device__ static inline void store_async(void const* dst_tma_map, bf16 *src, int global_col_idx, int global_row_idx) { + uint64_t tma_ptr = reinterpret_cast(dst_tma_map); + uint32_t src_ptr = static_cast(__cvta_generic_to_shared(src)); + + asm volatile ( + "cp.async.bulk.tensor.3d.global.shared::cta.tile.bulk_group" + " [%0, {%2, %3, %4}], [%1];" + : + : "l"(tma_ptr), "r"(src_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx / 64) + : "memory" + ); +} + +__device__ static __forceinline__ void wait(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + // Call mbarrier.try_wait in a while loop till it returns true. + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static __forceinline__ void arrive(uint64_t* bar, uint32_t count=1) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); +} + +__device__ void arrive_cluster(uint64_t* bar, uint32_t cta_id, uint32_t count=1) { + uint32_t smem_addr = static_cast(__cvta_generic_to_shared(bar)); + // Call arrive on barrier stored in SM number `cta_id` + // Use mapa.shared instruction to access other SM's barrier. + asm volatile( + "{\n" + ".reg .b32 remAddr32;\n" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n" + "mbarrier.arrive.shared::cluster.b64 _, [remAddr32], %2;\n" + "}" + :: "r"(smem_addr), "r"(cta_id), "r"(count)); +} + +__device__ static inline void load_async_multicast(bf16 *dst, void const* src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx, uint16_t cluster_mask) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5}], [%2], %6;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64), "h"(cluster_mask) + : "memory" + ); +} + +template +struct Schedule; + +constexpr int SPACE_LEN = 128; +int *_dspace; + +template +struct Schedule<2, NUM_SM, BM, BN, TM, TN> { + int it; + int *space; + + __device__ __forceinline__ Schedule(int M, int N, int block, int *_space) { + it = 0; + space = _space; + } + + __device__ __forceinline__ bool next(int &block_m, int& block_n) { + if (it >= SPACE_LEN) { + return false; + } + int now = space[it]; + if (now == -1) { + return false; + } + block_m = now >> 16; + block_n = (now & ((1<<16)-1)); + ++it; + return true; + } +}; + +template +struct SMem { + alignas(128) bf16 A[BM*BK*QSIZE]; + alignas(128) bf16 B[BK*BN*QSIZE]; + alignas(128) bf16 C[BN*(BM+(BM/64)*8)]; // padding of 8 elements per consumer + alignas(8) uint64_t full[QSIZE], empty[QSIZE]; + int space[SPACE_LEN]; +}; + +template +__global__ __launch_bounds__(NUM_THREADS) void __cluster_dims__(CLUSTER_M * CLUSTER_N, 1, 1) matmulKernel12(int M, int N, int K, const __grid_constant__ CUtensorMap tensorMapC, const __grid_constant__ CUtensorMap tensorMapA, const __grid_constant__ CUtensorMap tensorMapB, int* dspace) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int num_consumers = (NUM_THREADS / 128) - 1; + constexpr int B_WG_M = BM / num_consumers; + constexpr int B_WG_M_PADDED = B_WG_M + 8; + constexpr int CLUSTERS = CLUSTER_M * CLUSTER_N; + assert((M / BM) % CLUSTER_M == 0); + assert((N / BN) % CLUSTER_N == 0); + + extern __shared__ __align__(128) uint8_t smem[]; + SMem &s = *reinterpret_cast*>(smem); + bf16 *sA = s.A, *sB = s.B, *sC = s.C; + uint64_t *full = s.full, *empty = s.empty; + int *space = s.space; + + uint32_t rank; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(rank) :); + // Load schedule for this SM + if (threadIdx.x < SPACE_LEN) space[threadIdx.x] = dspace[rank*SPACE_LEN+threadIdx.x]; + + const int num_blocks_k = K / BK; + int wg_idx = threadIdx.x / 128; + int tid = threadIdx.x % 128; + + if (threadIdx.x == 0) { + for (int i = 0; i < QSIZE; ++i) { + init_barrier(&full[i], 0, 1); + init_barrier(&empty[i], 0, num_consumers*CLUSTERS); + } + } + asm volatile("barrier.cluster.arrive;\n" : :); + asm volatile("barrier.cluster.wait;\n" : :); + + + Schedule<2, NUM_SM/CLUSTERS, BM*CLUSTER_M, BN*CLUSTER_N, 16/CLUSTER_M, 8/CLUSTER_N> schedule(M, N, rank, &space[0]); + + asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(rank) :); + uint32_t rank_m = rank / CLUSTER_N; + uint32_t rank_n = rank % CLUSTER_N; + + // Producer + if (wg_idx == 0) { + constexpr int num_regs = (num_consumers <= 2 ? 24 : 32); + warpgroup_reg_dealloc(); + if (tid == 0) { + int p = 0; + int qidx = 0; + uint32_t col_mask = 0; + for (int i = 0; i < CLUSTER_M; ++i) { + col_mask |= (1 << (i * CLUSTER_N)); + } + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) { qidx = 0; p ^= 1;} + wait(&empty[qidx], p); + + expect_bytes(&full[qidx], (BK*BN+BK*BM)*sizeof(bf16)); + if constexpr (CLUSTER_N > 1) { + uint32_t mask = ((1 << CLUSTER_N) - 1) << (rank_m * CLUSTER_N); + if (rank_n == 0) { + load_async_multicast(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM, mask); + } + } else { + load_async(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM); + } + + if constexpr (CLUSTER_M > 1) { + if (rank_m == 0) { + load_async_multicast(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN, col_mask << rank_n); + } + } else { + load_async(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN); + } + } + } + } + } else { + constexpr int num_regs = (num_consumers == 1 ? 256 : (num_consumers == 2 ? 240 : 160)); + warpgroup_reg_alloc(); + float d[B_WG_M/WGMMA_M][WGMMA_N/16][8]; + --wg_idx; + for (int qidx = 0; qidx < QSIZE; ++qidx) { + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + } + int p = 0; + int qidx = 0; + int num_block_m, num_block_n; + + // stmatrix and wgmma layouts match for BF16: + // 1) 9.7.14.5.16 Warp-level matrix store instruction: stmatrix + // 2) Figure 148: WGMMA .m64nNk16 register fragment layout for accumulator matrix D + int lane = tid % 32, warp = tid / 32; + bf16* block_sC = sC + wg_idx*B_WG_M_PADDED*BN; + uint32_t tid_offset = warp*16 + (lane % 8) * B_WG_M_PADDED; // offset for x1 stmatrix + tid_offset += (lane/16)*B_WG_M_PADDED*8 + (lane & 8); // row & column offsets for .x4 on stmatrix + uint32_t base_addr = static_cast(__cvta_generic_to_shared(block_sC)) + tid_offset * sizeof(bf16); + + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + { + if (qidx == QSIZE) {qidx = 0; p ^= 1; }; + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + { + wgmma(d[m_it], &wgmma_sA[0], &wgmma_sB[0]); + #pragma unroll + for (int k_it = 1; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + #pragma unroll + for (int bk = 64; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + ++qidx; + } + for (int block_k_iter = 1; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) {qidx = 0; p ^= 1; }; + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + #pragma unroll + for (int bk = 0; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + } + + asm volatile("cp.async.bulk.wait_group 0;"); + + bf16 d_bf16[8]; + int* data_ptr = (int*)d_bf16; + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + for (int w = 0; w < WGMMA_N; w+=16) { + uint32_t addr = base_addr + (w*B_WG_M_PADDED + m_it*WGMMA_M) * sizeof(bf16); + for (int k = 0; k < 8; k++) { + d_bf16[k] = (bf16)(d[m_it][w/16][k]); + } + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared::cta.b16 [%0], {%1, %2, %3, %4};" + :: "r"(addr), "r"(data_ptr[0]), "r"(data_ptr[1]), "r"(data_ptr[2]), "r"(data_ptr[3])); + } + } + asm volatile("bar.sync %0, 128;\n" ::"r"(wg_idx + 2) : "memory"); + if (tid == 0) { + store_async(&tensorMapC, block_sC, num_block_m*BM + wg_idx*B_WG_M, num_block_n*BN); + asm volatile("cp.async.bulk.commit_group;"); + } + } + } +} + +// Rotate/flip quadrant appropriately +void rot(int n, int& x, int& y, int rx, int ry) { + if (ry == 0) { + if (rx == 1) { + x = n-1 - x; + y = n-1 - y; + } + // Swap x and y + int t = x; + x = y; + y = t; + } +} + +// Convert distance along curve to (x,y) point +void d2xy(int n, int d, int& x, int& y) { + int rx, ry, s, t = d; + x = y = 0; + for (s = 1; s < n; s *= 2) { + rx = 1 & (t/2); + ry = 1 & (t ^ rx); + rot(s, x, y, rx, ry); + x += s * rx; + y += s * ry; + t /= 4; + } +} + +void createHilbert(int M, int N, int CORES, int *space) { + int dim = (1 << (32 - __builtin_clz(max(M, N) - 1))); + int core = 0; + std::vector v(dim, std::string(dim, '.')); + memset(space, -1, sizeof(int)*CORES*SPACE_LEN); + int FCORES = 64; + int total = 0; + std::vector> pos(CORES, std::vector()); + for (int i = 0; i < dim*dim; ++i) { + int x, y; + d2xy(dim, i, x, y); + if (x < M && y < N) { + assert(loc < SPACE_LEN); + assert(v[x][y] == '.'); + v[x][y] = '*'; + ++total; + pos[core].push_back((x << 16) | y); + ++core; + if (core == FCORES) {core = 0;} + } + } + core = FCORES; + for (int i = 0; i < FCORES; ++i) { + if (pos.back().size() >= pos[0].size()-1) break; + pos[core].push_back(pos[i].back()); + pos[i].pop_back(); + ++core; + if (core == CORES) {core = FCORES;} + } + for (int i = 0; i < CORES; ++i) { + for (int j = 0; j < pos[i].size(); ++j) { + space[i*SPACE_LEN + j] = pos[i][j]; + } + } + assert(total == M*N); +} + +void runKernel12(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 128; + constexpr int BN = 256; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128*3; + constexpr int QSIZE = 3; + constexpr int CLUSTER_M = 2; + constexpr int CLUSTER_N = 1; + constexpr int NUM_SM = 128; + constexpr int NUM_CONSUMERS = (NUM_THREADS / 128) - 1; + static_assert(NUM_SM % (CLUSTER_M*CLUSTER_N) == 0); + + if (_prev_m != M) { + d_tma_map_A = create_tensor_map(A, M, K); + d_tma_map_B = create_tensor_map(B, N, K); + d_tma_map_C = create_tensor_map(C, N, M); + _prev_m = M; + _prev_n = N; + _prev_k = K; + int *space; + space = (int*)malloc(sizeof(int)*NUM_SM*SPACE_LEN); + createHilbert(CEIL_DIV(M, BM*CLUSTER_M), CEIL_DIV(N, BN*CLUSTER_N), NUM_SM/CLUSTER_M/CLUSTER_N, space); + cudaCheck(cudaMalloc((void **)&_dspace, sizeof(int)*NUM_SM*SPACE_LEN)); + cudaCheck(cudaMemcpy(_dspace, space, sizeof(int)*NUM_SM*SPACE_LEN, cudaMemcpyHostToDevice)); + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = matmulKernel12; + constexpr size_t sMemSize = sizeof(SMem); + static_assert(sMemSize < 256 * 1024); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<>>(M, N, K, d_tma_map_C, d_tma_map_A, d_tma_map_B, _dspace); +} + +} // namespace M12 + +using M12::runKernel12; + \ No newline at end of file diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_2.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_2.cuh new file mode 100644 index 000000000..4e8cefd16 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_2.cuh @@ -0,0 +1,197 @@ + +namespace M2 { + +using barrier = cuda::barrier; +namespace cde = cuda::device::experimental; + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +__device__ uint64_t make_smem_desc(bf16* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; + } + + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +template +__device__ void warpgroup_wait() { + static_assert(N >= 0 && N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +template +void create_tensor_map(CUtensorMap *tma_map, bf16* gmem_ptr, int blocks_height, int blocks_width) { + void* gmem_address = (void*)gmem_ptr; + uint64_t gmem_prob_shape[5] = {(uint64_t)BlockMinorSize*blocks_width, (uint64_t)BlockMajorSize*blocks_height, 1, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16), sizeof(bf16) * BlockMinorSize*blocks_width, 0, 0, 0}; + uint32_t smem_box_shape[5] = {uint32_t(BlockMinorSize), uint32_t(BlockMajorSize), 1, 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, gmem_address, gmem_prob_shape, + gmem_prob_stride + 1, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); +} + +CUtensorMap *d_tma_map_A = 0; +CUtensorMap *d_tma_map_B = 0; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__host__ static inline CUtensorMap* allocate_and_create_tensor_map(bf16* src, int blocks_height, int blocks_width) { + CUtensorMap *tma_map_d; + cudaMalloc(&tma_map_d, sizeof(CUtensorMap)); + CUtensorMap tma_map_host; + create_tensor_map(&tma_map_host, src, blocks_height, blocks_width); + cudaMemcpy(tma_map_d, &tma_map_host, sizeof(CUtensorMap), cudaMemcpyHostToDevice); + return tma_map_d; +} + +template +__device__ void wgmma64(float d[4][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " %34, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__global__ void __launch_bounds__(NUM_THREADS) matmulKernel2(int M, int N, int K, bf16* C, const CUtensorMap* tensorMapA, const CUtensorMap* tensorMapB) { + __shared__ alignas(128) bf16 sA[BM*BK]; + __shared__ alignas(128) bf16 sB[BK*BN]; + float d[WGMMA_N/16][8]; + static_assert(sizeof(d) * 128 == BM * BN * sizeof(float)); + memset(d, 0, sizeof(d)); + + const int num_blocks_k = K / BK; + int num_block_n = blockIdx.x % (N / BN); + int num_block_m = blockIdx.x / (N / BN); + #pragma nv_diag_suppress static_var_with_dynamic_init + __shared__ barrier barA; + __shared__ barrier barB; + + if (threadIdx.x == 0) { + init(&barA, blockDim.x); + init(&barB, blockDim.x); + cde::fence_proxy_async_shared_cta(); + } + __syncthreads(); + + barrier::arrival_token tokenA, tokenB; + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter) { + // Load + if (threadIdx.x == 0) { + cde::cp_async_bulk_tensor_2d_global_to_shared(&sA[0], tensorMapA, block_k_iter*BK, num_block_m*BM, barA); + tokenA = cuda::device::barrier_arrive_tx(barA, 1, sizeof(sA)); + cde::cp_async_bulk_tensor_2d_global_to_shared(&sB[0], tensorMapB, block_k_iter*BK, num_block_n*BN, barB); + tokenB = cuda::device::barrier_arrive_tx(barB, 1, sizeof(sB)); + } else { + tokenA = barA.arrive(); + tokenB = barB.arrive(); + } + barA.wait(std::move(tokenA)); + barB.wait(std::move(tokenB)); + __syncthreads(); + + // Compute + warpgroup_arrive(); + wgmma64<1, 1, 1, 0, 0>(d, &sA[0], &sB[0]); + wgmma64<1, 1, 1, 0, 0>(d, &sA[WGMMA_K], &sB[WGMMA_K]); + wgmma64<1, 1, 1, 0, 0>(d, &sA[2*WGMMA_K], &sB[2*WGMMA_K]); + wgmma64<1, 1, 1, 0, 0>(d, &sA[3*WGMMA_K], &sB[3*WGMMA_K]); + warpgroup_commit_batch(); + warpgroup_wait<0>(); + } + + // Store + { + int tid = threadIdx.x; + int lane = tid % 32; + int warp = tid / 32; + uint32_t row = warp*16 + lane / 4; + bf16 *block_C = C + num_block_n*BN*M + num_block_m*BM; + + for (int m_it = 0; m_it < BM/WGMMA_M; ++m_it) { + for (int n_it = 0; n_it < BN/WGMMA_N; ++n_it) { + for (int w = 0; w < WGMMA_N/16; ++w) { + int col = 16*w + 2*(tid % 4); + #define IDX(i, j) ((j + n_it*WGMMA_N)*M + ((i) + m_it*WGMMA_M)) + + block_C[IDX(row, col)] = d[w][0]; + block_C[IDX(row, col+1)] = d[w][1]; + block_C[IDX(row+8, col)] = d[w][2]; + block_C[IDX(row+8, col+1)] = d[w][3]; + + block_C[IDX(row, col+8)] = d[w][4]; + block_C[IDX(row, col+9)] = d[w][5]; + block_C[IDX(row+8, col+8)] = d[w][6]; + block_C[IDX(row+8, col+9)] = d[w][7]; + + #undef IDX + } + } + } + } +} + + +void runKernel2(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C) { + constexpr int BM = 64; + constexpr int BN = 64; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128; + + if (!d_tma_map_A) { + d_tma_map_A = allocate_and_create_tensor_map(A, M / BM, K / BK); + d_tma_map_B = allocate_and_create_tensor_map(B, N / BN, K / BK); + _prev_m = M; + _prev_n = N; + _prev_k = K; + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + matmulKernel2< + /*BM*/ BM, + /*BN*/ BN, + /*BK*/ BK, + /*WGMMA_M*/ 64, + /*WGMMA_N*/ 64, + /*WGMMA_K*/ 16, + /*NUM_THREADS*/ NUM_THREADS> + <<<(M/BM) * (N/BN), NUM_THREADS>>>(M, N, K, C, d_tma_map_A, d_tma_map_B); +} + +} // namespace M2 + +using M2::runKernel2; diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_3.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_3.cuh new file mode 100644 index 000000000..08b48802a --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_3.cuh @@ -0,0 +1,410 @@ + +namespace M3 { + +using barrier = cuda::barrier; +namespace cde = cuda::device::experimental; + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +__device__ uint64_t make_smem_desc(bf16* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; + } + + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +template +__device__ void warpgroup_wait() { + static_assert(N >= 0 && N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +template +void create_tensor_map(CUtensorMap *tma_map, bf16* gmem_ptr, int blocks_height, int blocks_width) { + void* gmem_address = (void*)gmem_ptr; + uint64_t gmem_prob_shape[5] = {(uint64_t)BlockMinorSize*blocks_width, (uint64_t)BlockMajorSize*blocks_height, 1, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16), sizeof(bf16) * BlockMinorSize*blocks_width, 0, 0, 0}; + uint32_t smem_box_shape[5] = {uint32_t(BlockMinorSize), uint32_t(BlockMajorSize), 1, 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, gmem_address, gmem_prob_shape, + gmem_prob_stride + 1, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); +} + +CUtensorMap *d_tma_map_A = 0; +CUtensorMap *d_tma_map_B = 0; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__host__ static inline CUtensorMap* allocate_and_create_tensor_map(bf16* src, int blocks_height, int blocks_width) { + CUtensorMap *tma_map_d; + cudaMalloc(&tma_map_d, sizeof(CUtensorMap)); + CUtensorMap tma_map_host; + create_tensor_map(&tma_map_host, src, blocks_height, blocks_width); + cudaMemcpy(tma_map_d, &tma_map_host, sizeof(CUtensorMap), cudaMemcpyHostToDevice); + return tma_map_d; +} + +template +__device__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ void wgmma64(float d[4][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " %34, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ void wgmma32(float d[2][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}, " + " %16," + " %17," + " %18, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ void wgmma16(float d[1][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}, " + " %8," + " %9," + " %10, %11, %12, %13, %14;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ inline void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 32 || WGMMA_N == 64 || WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 64) + wgmma64<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 32) + wgmma32<1, 1, 1, 0, 0>(d, sA, sB); +} + +template +struct SMem { + alignas(128) bf16 A[BM*BK]; + alignas(128) bf16 B[BK*BN]; +}; + +template +__global__ void __launch_bounds__(NUM_THREADS) matmulKernel3(int M, int N, int K, bf16* C, const CUtensorMap* tensorMapA, const CUtensorMap* tensorMapB, int *DB) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int B_WG_M = BM / (NUM_THREADS / 128); + extern __shared__ SMem s; + bf16 *sA = s.A; + bf16 *sB = s.B; + // Barriers cannot be in the struct and have to be declared this way + #pragma nv_diag_suppress static_var_with_dynamic_init + __shared__ barrier barA, barB; + float d[B_WG_M/WGMMA_M][WGMMA_N/16][8]; + static_assert(sizeof(d) * NUM_THREADS == BM * BN * sizeof(float)); + memset(d, 0, sizeof(d)); + + const int num_blocks_k = K / BK; + int num_block_n = blockIdx.x % (N / BN); + int num_block_m = blockIdx.x / (N / BN); + + if (threadIdx.x == 0) { + init(&barA, blockDim.x); + init(&barB, blockDim.x); + cde::fence_proxy_async_shared_cta(); + } + __syncthreads(); + int wg_idx = threadIdx.x / 128; + + barrier::arrival_token tokenA, tokenB; + int sumLoad = 0, cntLoad = 0; + int sumCompute = 0, cntCompute = 0; + int sumStore = 0, cntStore = 0; + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter) { + clock_t start = clock(); + // Load + if (threadIdx.x == 0) { + cde::cp_async_bulk_tensor_2d_global_to_shared(&sA[0], tensorMapA, block_k_iter*BK, num_block_m*BM, barA); + tokenA = cuda::device::barrier_arrive_tx(barA, 1, BK*BM*sizeof(bf16)); + cde::cp_async_bulk_tensor_2d_global_to_shared(&sB[0], tensorMapB, block_k_iter*BK, num_block_n*BN, barB); + tokenB = cuda::device::barrier_arrive_tx(barB, 1, BK*BN*sizeof(bf16)); + } else { + tokenA = barA.arrive(); + tokenB = barB.arrive(); + } + barA.wait(std::move(tokenA)); + barB.wait(std::move(tokenB)); + __syncthreads(); + if constexpr (DBG) { + sumLoad += clock() - start; + cntLoad++; + start = clock(); + } + + // Compute + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + BK*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + #pragma unroll + for (int k_it = 0; k_it < BK/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &sB[k_it*WGMMA_K]); + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + + if constexpr (DBG) { + sumCompute += clock() - start; + cntCompute++; + } + } + + // Store + { + clock_t start = clock(); + + uint32_t tid = threadIdx.x % 128; + uint32_t lane = tid & 31; + uint32_t warp = tid / 32; + uint32_t row = warp*16 + lane / 4; + + bf16 *block_C = C + num_block_n*BN*M + num_block_m*BM; + + #pragma unroll + for (uint32_t m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + int yo = m_it*WGMMA_M + wg_idx*B_WG_M; + #pragma unroll + for (uint32_t w = 0; w < WGMMA_N/16; ++w) { + int col = 16*w + 2*(tid % 4); + #define IDX(i, j) ((j)*M + ((i) + yo)) + + block_C[IDX(row, col)] = d[m_it][w][0]; + block_C[IDX(row, col+1)] = d[m_it][w][1]; + block_C[IDX(row+8, col)] = d[m_it][w][2]; + block_C[IDX(row+8, col+1)] = d[m_it][w][3]; + block_C[IDX(row, col+8)] = d[m_it][w][4]; + block_C[IDX(row, col+9)] = d[m_it][w][5]; + block_C[IDX(row+8, col+8)] = d[m_it][w][6]; + block_C[IDX(row+8, col+9)] = d[m_it][w][7]; + + #undef IDX + } + } + if constexpr (DBG) { + sumStore += clock() - start; + cntStore++; + if (threadIdx.x == 63) { + int i = blockIdx.x*6; + DB[i] = sumLoad; DB[i + 1] = cntLoad; + DB[i + 2] = sumCompute; DB[i + 3] = cntCompute; + DB[i + 4] = sumStore; DB[i + 5] = cntStore; + } + } + } +} + + +void runKernel3(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 128; + constexpr int BN = 128; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128; + + if (!d_tma_map_A) { + d_tma_map_A = allocate_and_create_tensor_map(A, M / BM, K / BK); + d_tma_map_B = allocate_and_create_tensor_map(B, N / BN, K / BK); + _prev_m = M; + _prev_n = N; + _prev_k = K; + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = DB ? matmulKernel3 + : matmulKernel3; + size_t sMemSize = sizeof(SMem); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<<(M/BM) * (N/BN), NUM_THREADS, sMemSize>>>(M, N, K, C, d_tma_map_A, d_tma_map_B, DB); +} + +} // namespace M3 + +using M3::runKernel3; diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_4.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_4.cuh new file mode 100644 index 000000000..7f8dcf406 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_4.cuh @@ -0,0 +1,393 @@ + +namespace M4 { + +using barrier = cuda::barrier; +namespace cde = cuda::device::experimental; + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +__device__ uint64_t make_smem_desc(bf16* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; + } + + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +template +__device__ void warpgroup_wait() { + static_assert(N >= 0 && N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +template +void create_tensor_map(CUtensorMap *tma_map, bf16* gmem_ptr, int blocks_height, int blocks_width) { + void* gmem_address = (void*)gmem_ptr; + uint64_t gmem_prob_shape[5] = {(uint64_t)BlockMinorSize*blocks_width, (uint64_t)BlockMajorSize*blocks_height, 1, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16), sizeof(bf16) * BlockMinorSize*blocks_width, 0, 0, 0}; + uint32_t smem_box_shape[5] = {uint32_t(BlockMinorSize), uint32_t(BlockMajorSize), 1, 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, gmem_address, gmem_prob_shape, + gmem_prob_stride + 1, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); +} + +CUtensorMap *d_tma_map_A = 0; +CUtensorMap *d_tma_map_B = 0; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__host__ static inline CUtensorMap* allocate_and_create_tensor_map(bf16* src, int blocks_height, int blocks_width) { + CUtensorMap *tma_map_d; + cudaMalloc(&tma_map_d, sizeof(CUtensorMap)); + CUtensorMap tma_map_host; + create_tensor_map(&tma_map_host, src, blocks_height, blocks_width); + cudaMemcpy(tma_map_d, &tma_map_host, sizeof(CUtensorMap), cudaMemcpyHostToDevice); + return tma_map_d; +} + +template +__device__ __forceinline__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + + +template +__device__ __forceinline__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma64(float d[4][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " %34, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma32(float d[2][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}, " + " %16," + " %17," + " %18, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma16(float d[1][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}, " + " %8," + " %9," + " %10, %11, %12, %13, %14;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 32 || WGMMA_N == 64 || WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 208 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 64) + wgmma64<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 32) + wgmma32<1, 1, 1, 0, 0>(d, sA, sB); +} + +template +struct SMem { + alignas(128) bf16 A[BM*BK*QSIZE]; + alignas(128) bf16 B[BK*BN*QSIZE]; +}; + +template +__global__ __launch_bounds__(NUM_THREADS) void matmulKernel4(int M, int N, int K, bf16* C, const CUtensorMap* tensorMapA, const CUtensorMap* tensorMapB) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int num_consumers = (NUM_THREADS / 128) - 1; + constexpr int B_WG_M = BM / num_consumers; + + extern __shared__ __align__(128) uint8_t smem[]; + SMem &s = *reinterpret_cast*>(smem); + bf16 *sA = s.A; + bf16 *sB = s.B; + // Barriers cannot be in the struct and have to be declared this way + #pragma nv_diag_suppress static_var_with_dynamic_init + __shared__ barrier full[QSIZE], empty[QSIZE]; + + const int num_blocks_k = K / BK; + int num_block_n = blockIdx.x % (N / BN); + int num_block_m = blockIdx.x / (N / BN); + int wg_idx = threadIdx.x / 128; + int tid = threadIdx.x % 128; + + if (threadIdx.x == 0) { + for (int i = 0; i < QSIZE; ++i) { + init(&full[i], num_consumers * 128 + 1); + init(&empty[i], num_consumers * 128 + 1); + } + cde::fence_proxy_async_shared_cta(); + } + __syncthreads(); + + // Producer + if (wg_idx == 0) { + constexpr int num_regs = (num_consumers <= 2 ? 24 : 32); + if (tid == 0) { + int qidx = 0; + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) qidx = 0; + empty[qidx].wait(empty[qidx].arrive()); + cde::cp_async_bulk_tensor_2d_global_to_shared(&sA[qidx*BK*BM], tensorMapA, block_k_iter*BK, num_block_m*BM, full[qidx]); + cde::cp_async_bulk_tensor_2d_global_to_shared(&sB[qidx*BK*BN], tensorMapB, block_k_iter*BK, num_block_n*BN, full[qidx]); + barrier::arrival_token _ = cuda::device::barrier_arrive_tx(full[qidx], 1, (BK*BN+BK*BM)*sizeof(bf16)); + } + } + } else { + for (int i = 0; i < QSIZE; ++i) { + barrier::arrival_token _ = empty[i].arrive(); + } + float d[BM/WGMMA_M][WGMMA_N/16][8]; + memset(d, 0, sizeof(d)); + int qidx = 0; + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) qidx = 0; + full[qidx].wait(full[qidx].arrive()); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < BM/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + BK*m_it*WGMMA_M; + #pragma unroll + for (int k_it = 0; k_it < BK/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &sB[qidx*BK*BN + k_it*WGMMA_K]); + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + barrier::arrival_token _ = empty[qidx].arrive(); + } + + tid = threadIdx.x % 128; + int lane = tid % 32; + int warp = tid / 32; + int row = warp*16 + lane / 4; + bf16 *block_C = C + num_block_n*BN*M + num_block_m*BM; + + #pragma unroll + for (int m_it = 0; m_it < BM/WGMMA_M; ++m_it) { + int yo = m_it*WGMMA_M; + #pragma unroll + for (int w = 0; w < WGMMA_N/16; ++w) { + int col = 16*w + 2*(tid % 4); + #define IDX(i, j) ((j)*M + ((i) + yo)) + + block_C[IDX(row, col)] = d[m_it][w][0]; + block_C[IDX(row, col+1)] = d[m_it][w][1]; + block_C[IDX(row+8, col)] = d[m_it][w][2]; + block_C[IDX(row+8, col+1)] = d[m_it][w][3]; + + block_C[IDX(row, col+8)] = d[m_it][w][4]; + block_C[IDX(row, col+9)] = d[m_it][w][5]; + block_C[IDX(row+8, col+8)] = d[m_it][w][6]; + block_C[IDX(row+8, col+9)] = d[m_it][w][7]; + #undef IDX + } + } + } +} + + +void runKernel4(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 128; + constexpr int BN = 128; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128*2; + constexpr int QSIZE = 5; + + if (!d_tma_map_A) { + d_tma_map_A = allocate_and_create_tensor_map(A, M / BM, K / BK); + d_tma_map_B = allocate_and_create_tensor_map(B, N / BN, K / BK); + _prev_m = M; + _prev_n = N; + _prev_k = K; + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = matmulKernel4; + size_t sMemSize = sizeof(SMem); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<<(M/BM) * (N/BN), NUM_THREADS, sMemSize>>>(M, N, K, C, d_tma_map_A, d_tma_map_B); +} + +} // namespace M4 + +using M4::runKernel4; + \ No newline at end of file diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_5.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_5.cuh new file mode 100644 index 000000000..c4478ec27 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_5.cuh @@ -0,0 +1,407 @@ + +namespace M5 { + +using barrier = cuda::barrier; +namespace cde = cuda::device::experimental; + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +__device__ uint64_t make_smem_desc(bf16* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; + } + + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +template +__device__ void warpgroup_wait() { + static_assert(N >= 0 && N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +template +void create_tensor_map(CUtensorMap *tma_map, bf16* gmem_ptr, int blocks_height, int blocks_width) { + void* gmem_address = (void*)gmem_ptr; + uint64_t gmem_prob_shape[5] = {(uint64_t)BlockMinorSize*blocks_width, (uint64_t)BlockMajorSize*blocks_height, 1, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16), sizeof(bf16) * BlockMinorSize*blocks_width, 0, 0, 0}; + uint32_t smem_box_shape[5] = {uint32_t(BlockMinorSize), uint32_t(BlockMajorSize), 1, 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, gmem_address, gmem_prob_shape, + gmem_prob_stride + 1, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); +} + +CUtensorMap *d_tma_map_A = 0; +CUtensorMap *d_tma_map_B = 0; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__host__ static inline CUtensorMap* allocate_and_create_tensor_map(bf16* src, int blocks_height, int blocks_width) { + CUtensorMap *tma_map_d; + cudaMalloc(&tma_map_d, sizeof(CUtensorMap)); + CUtensorMap tma_map_host; + create_tensor_map(&tma_map_host, src, blocks_height, blocks_width); + cudaMemcpy(tma_map_d, &tma_map_host, sizeof(CUtensorMap), cudaMemcpyHostToDevice); + return tma_map_d; +} + +template +__device__ __forceinline__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + + +template +__device__ __forceinline__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma64(float d[4][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " %34, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma32(float d[2][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}, " + " %16," + " %17," + " %18, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma16(float d[1][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}, " + " %8," + " %9," + " %10, %11, %12, %13, %14;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 32 || WGMMA_N == 64 || WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 208 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 64) + wgmma64<1, 1, 1, 0, 0>(d, sA, sB); + if constexpr (WGMMA_N == 32) + wgmma32<1, 1, 1, 0, 0>(d, sA, sB); +} + +template +struct SMem { + alignas(128) bf16 A[BM*BK*QSIZE]; + alignas(128) bf16 B[BK*BN*QSIZE]; +}; + +template +__device__ void warpgroup_reg_alloc() { + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +__device__ void warpgroup_reg_dealloc() { + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +__global__ __launch_bounds__(NUM_THREADS) void matmulKernel5(int M, int N, int K, bf16* C, const CUtensorMap* tensorMapA, const CUtensorMap* tensorMapB) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int num_consumers = (NUM_THREADS / 128) - 1; + constexpr int B_WG_M = BM / num_consumers; + + extern __shared__ __align__(128) uint8_t smem[]; + SMem &s = *reinterpret_cast*>(smem); + bf16 *sA = s.A; + bf16 *sB = s.B; + // Barriers cannot be in the struct and have to be declared this way + #pragma nv_diag_suppress static_var_with_dynamic_init + __shared__ barrier full[QSIZE], empty[QSIZE]; + + const int num_blocks_k = K / BK; + int num_block_n = blockIdx.x % (N / BN); + int num_block_m = blockIdx.x / (N / BN); + int wg_idx = threadIdx.x / 128; + int tid = threadIdx.x % 128; + + if (threadIdx.x == 0) { + for (int i = 0; i < QSIZE; ++i) { + init(&full[i], num_consumers * 128 + 1); + init(&empty[i], num_consumers * 128 + 1); + } + cde::fence_proxy_async_shared_cta(); + } + __syncthreads(); + + // Producer + if (wg_idx == 0) { + constexpr int num_regs = (num_consumers <= 2 ? 24 : 32); + warpgroup_reg_dealloc(); + if (tid == 0) { + int qidx = 0; + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) qidx = 0; + empty[qidx].wait(empty[qidx].arrive()); + cde::cp_async_bulk_tensor_2d_global_to_shared(&sA[qidx*BK*BM], tensorMapA, block_k_iter*BK, num_block_m*BM, full[qidx]); + cde::cp_async_bulk_tensor_2d_global_to_shared(&sB[qidx*BK*BN], tensorMapB, block_k_iter*BK, num_block_n*BN, full[qidx]); + barrier::arrival_token _ = cuda::device::barrier_arrive_tx(full[qidx], 1, (BK*BN+BK*BM)*sizeof(bf16)); + } + } + } else { + constexpr int num_regs = (num_consumers == 1 ? 256 : (num_consumers == 2 ? 240 : 160)); + warpgroup_reg_alloc(); + --wg_idx; + for (int i = 0; i < QSIZE; ++i) { + barrier::arrival_token _ = empty[i].arrive(); + } + float d[B_WG_M/WGMMA_M][WGMMA_N/16][8]; + memset(d, 0, sizeof(d)); + int qidx = 0; + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) qidx = 0; + full[qidx].wait(full[qidx].arrive()); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + BK*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + #pragma unroll + for (int k_it = 0; k_it < BK/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &sB[qidx*BK*BN + k_it*WGMMA_K]); + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + barrier::arrival_token _ = empty[qidx].arrive(); + } + + tid = threadIdx.x % 128; + int lane = tid % 32; + int warp = tid / 32; + int row = warp*16 + lane / 4; + bf16 *block_C = C + num_block_n*BN*M + num_block_m*BM; + + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + int yo = m_it*WGMMA_M + wg_idx*B_WG_M; + #pragma unroll + for (int w = 0; w < WGMMA_N/16; ++w) { + int col = 16*w + 2*(tid % 4); + #define IDX(i, j) ((j)*M + ((i) + yo)) + + block_C[IDX(row, col)] = d[m_it][w][0]; + block_C[IDX(row, col+1)] = d[m_it][w][1]; + block_C[IDX(row+8, col)] = d[m_it][w][2]; + block_C[IDX(row+8, col+1)] = d[m_it][w][3]; + + block_C[IDX(row, col+8)] = d[m_it][w][4]; + block_C[IDX(row, col+9)] = d[m_it][w][5]; + block_C[IDX(row+8, col+8)] = d[m_it][w][6]; + block_C[IDX(row+8, col+9)] = d[m_it][w][7]; + #undef IDX + } + } + } +} + + +void runKernel5(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 128; + constexpr int BN = 256; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128*3; + constexpr int QSIZE = 3; + + if (!d_tma_map_A) { + d_tma_map_A = allocate_and_create_tensor_map(A, M / BM, K / BK); + d_tma_map_B = allocate_and_create_tensor_map(B, N / BN, K / BK); + _prev_m = M; + _prev_n = N; + _prev_k = K; + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = matmulKernel5; + size_t sMemSize = sizeof(SMem); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<<(M/BM) * (N/BN), NUM_THREADS, sMemSize>>>(M, N, K, C, d_tma_map_A, d_tma_map_B); +} + +} // namespace M5 + +using M5::runKernel5; + \ No newline at end of file diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_6.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_6.cuh new file mode 100644 index 000000000..ab75d2f21 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_6.cuh @@ -0,0 +1,450 @@ + +namespace M6 { + +using barrier = cuda::barrier; +namespace cde = cuda::device::experimental; + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +__device__ uint64_t make_smem_desc(bf16* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; +} + + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +template +__device__ void warpgroup_wait() { + static_assert(N >= 0 && N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +template +void create_tensor_map(CUtensorMap *tma_map, bf16* gmem_ptr, int blocks_height, int blocks_width) { + void* gmem_address = (void*)gmem_ptr; + uint64_t gmem_prob_shape[5] = {(uint64_t)BlockMinorSize*blocks_width, (uint64_t)BlockMajorSize*blocks_height, 1, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16), sizeof(bf16) * BlockMinorSize*blocks_width, 0, 0, 0}; + uint32_t smem_box_shape[5] = {uint32_t(BlockMinorSize), uint32_t(BlockMajorSize), 1, 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, gmem_address, gmem_prob_shape, + gmem_prob_stride + 1, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); +} + +CUtensorMap d_tma_map_A; +CUtensorMap d_tma_map_B; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__host__ static inline CUtensorMap allocate_and_create_tensor_map(bf16* src, int blocks_height, int blocks_width) { + CUtensorMap tma_map_host; + create_tensor_map(&tma_map_host, src, blocks_height, blocks_width); + return tma_map_host; +} + +template +__device__ __forceinline__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + + +template +__device__ __forceinline__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma64(float d[4][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " %34, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma32(float d[2][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}, " + " %16," + " %17," + " %18, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 32 || WGMMA_N == 64 || WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 208 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128(d, sA, sB); + if constexpr (WGMMA_N == 64) + wgmma64(d, sA, sB); + if constexpr (WGMMA_N == 32) + wgmma32(d, sA, sB); +} + +template +struct SMem { + alignas(128) bf16 A[BM*BK*QSIZE]; + alignas(128) bf16 B[BK*BN*QSIZE]; +}; + +template +__device__ void warpgroup_reg_alloc() { + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +__device__ void warpgroup_reg_dealloc() { + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +struct Schedule; + +template +struct Schedule<0, NUM_SM, BM, BN, TM, TN> { + int st, en; + + __device__ __forceinline__ Schedule(int M, int N, int block) { + int total_blocks = M*N/(BM*BN); + int blocks_per_sm = total_blocks / NUM_SM; + int extra_blocks = total_blocks % NUM_SM; + if (block < extra_blocks) { + st = block*(blocks_per_sm + 1); + en = st + blocks_per_sm + 1; + } else { + st = extra_blocks + block*blocks_per_sm; + en = st + blocks_per_sm; + } + } + + __device__ __forceinline__ int next() { + if (en == st) return -1; + return st++; + } +}; + +template +struct Schedule<1, NUM_SM, BM, BN, TM, TN> { + int block; + int it; + int total_blocks_m; + int total_blocks_n; + + __device__ __forceinline__ Schedule(int M, int N, int _block) { + block = _block; + it = 0; + total_blocks_m = M/BM; + total_blocks_n = N/BN; + assert(total_blocks_m%TM == 0 && total_blocks_n%TN == 0); + } + + __device__ __forceinline__ int next() { + int num = it*NUM_SM + block; + if (num >= total_blocks_m*total_blocks_n) return -1; + + int cur_tile = num / (TM*TN); + int cur_tile_pos = num % (TM*TN); + int m = TM*(cur_tile / (total_blocks_n/TN)); + int n = TN*(cur_tile % (total_blocks_n/TN)); + m += cur_tile_pos / TN; + n += cur_tile_pos % TN; + ++it; + return m*total_blocks_n + n; + } +}; + +template +__global__ __launch_bounds__(NUM_THREADS) void matmulKernel6(int M, int N, int K, bf16* C, const __grid_constant__ CUtensorMap tensorMapA, const __grid_constant__ CUtensorMap tensorMapB) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int num_consumers = (NUM_THREADS / 128) - 1; + constexpr int B_WG_M = BM / num_consumers; + + extern __shared__ __align__(128) uint8_t smem[]; + SMem &s = *reinterpret_cast*>(smem); + bf16 *sA = s.A; + bf16 *sB = s.B; + // Barriers cannot be in the struct and have to be declared this way + #pragma nv_diag_suppress static_var_with_dynamic_init + __shared__ barrier full[QSIZE], empty[QSIZE]; + + const int num_blocks_k = K / BK; + int wg_idx = threadIdx.x / 128; + int tid = threadIdx.x % 128; + + if (threadIdx.x == 0) { + for (int i = 0; i < QSIZE; ++i) { + init(&full[i], num_consumers * 128 + 1); + init(&empty[i], num_consumers * 128 + 1); + } + cde::fence_proxy_async_shared_cta(); + } + __syncthreads(); + + Schedule<1, NUM_SM, BM, BN, 16, 8> schedule(M, N, blockIdx.x); + + // Producer + if (wg_idx == 0) { + constexpr int num_regs = (num_consumers <= 2 ? 24 : 32); + warpgroup_reg_dealloc(); + if (tid == 0) { + int qidx = 0; + for (int num_block = schedule.next(); num_block >= 0; num_block = schedule.next()) { + int num_block_n = num_block % (N / BN); + int num_block_m = num_block / (N / BN); + + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) qidx = 0; + empty[qidx].wait(empty[qidx].arrive()); + cde::cp_async_bulk_tensor_2d_global_to_shared(&sA[qidx*BK*BM], &tensorMapA, block_k_iter*BK, num_block_m*BM, full[qidx]); + cde::cp_async_bulk_tensor_2d_global_to_shared(&sB[qidx*BK*BN], &tensorMapB, block_k_iter*BK, num_block_n*BN, full[qidx]); + barrier::arrival_token _ = cuda::device::barrier_arrive_tx(full[qidx], 1, (BK*BN+BK*BM)*sizeof(bf16)); + } + } + } + } else { + constexpr int num_regs = (num_consumers == 1 ? 256 : (num_consumers == 2 ? 240 : 160)); + warpgroup_reg_alloc(); + float d[B_WG_M/WGMMA_M][WGMMA_N/16][8]; + --wg_idx; + for (int i = 0; i < QSIZE; ++i) { + barrier::arrival_token _ = empty[i].arrive(); + } + int qidx = 0; + + for (int num_block = schedule.next(); num_block >= 0; num_block = schedule.next()) { + int num_block_n = num_block % (N / BN); + int num_block_m = num_block / (N / BN); + memset(d, 0, sizeof(d)); + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) qidx = 0; + full[qidx].wait(full[qidx].arrive()); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + BK*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + #pragma unroll + for (int k_it = 0; k_it < BK/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &sB[qidx*BK*BN + k_it*WGMMA_K]); + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + barrier::arrival_token _ = empty[qidx].arrive(); + } + + int lane = tid % 32, warp = tid / 32, row = warp*16 + lane / 4; + bf16 *block_C = C + num_block_n*BN*M + num_block_m*BM; + + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + int yo = m_it*WGMMA_M + wg_idx*B_WG_M; + #pragma unroll + for (int w = 0; w < WGMMA_N/16; ++w) { + + int col = 16*w + 2*(tid % 4); + #define IDX(i, j) ((j)*M + ((i) + yo)) + + block_C[IDX(row, col)] = d[m_it][w][0]; + block_C[IDX(row, col+1)] = d[m_it][w][1]; + block_C[IDX(row+8, col)] = d[m_it][w][2]; + block_C[IDX(row+8, col+1)] = d[m_it][w][3]; + + block_C[IDX(row, col+8)] = d[m_it][w][4]; + block_C[IDX(row, col+9)] = d[m_it][w][5]; + block_C[IDX(row+8, col+8)] = d[m_it][w][6]; + block_C[IDX(row+8, col+9)] = d[m_it][w][7]; + #undef IDX + } + } + } + } +} + +void runKernel6(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 128; + constexpr int BN = 256; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128*3; + constexpr int QSIZE = 3; + constexpr int NUM_SM = 128; + + if (_prev_m != M) { + d_tma_map_A = allocate_and_create_tensor_map(A, M / BM, K / BK); + d_tma_map_B = allocate_and_create_tensor_map(B, N / BN, K / BK); + _prev_m = M; + _prev_n = N; + _prev_k = K; + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = matmulKernel6; + size_t sMemSize = sizeof(SMem); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<>>(M, N, K, C, d_tma_map_A, d_tma_map_B); +} + +} // namespace M6 + +using M6::runKernel6; + \ No newline at end of file diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_7.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_7.cuh new file mode 100644 index 000000000..d1c38d4d0 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_7.cuh @@ -0,0 +1,482 @@ + +namespace M7 { + +// using barrier = cuda::barrier; +namespace cde = cuda::device::experimental; + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +__device__ uint64_t make_smem_desc(bf16* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; +} + + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +template +__device__ void warpgroup_wait() { + static_assert(N >= 0 && N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +template +__host__ static inline CUtensorMap create_tensor_map(bf16* gmem_ptr, int global_height, int global_width) { + CUtensorMap tma_map; + void* gmem_address = (void*)gmem_ptr; + static_assert(BlockMinorSize >= 64); + assert(global_width % 64 == 0); + uint64_t gmem_prob_shape[5] = {64, (uint64_t)global_height, (uint64_t)global_width/64, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16) * global_width, 64*sizeof(bf16), 0, 0, 0}; + uint32_t smem_box_shape[5] = {64, uint32_t(BlockMajorSize), uint32_t(BlockMinorSize/64), 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + &tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 5, gmem_address, gmem_prob_shape, + gmem_prob_stride, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); + return tma_map; +} + +CUtensorMap d_tma_map_A; +CUtensorMap d_tma_map_B; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__device__ __forceinline__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + + +template +__device__ __forceinline__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma64(float d[4][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " %34, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma32(float d[2][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}, " + " %16," + " %17," + " %18, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 32 || WGMMA_N == 64 || WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 208 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128(d, sA, sB); + if constexpr (WGMMA_N == 64) + wgmma64(d, sA, sB); + if constexpr (WGMMA_N == 32) + wgmma32(d, sA, sB); +} + +template +struct SMem { + alignas(128) bf16 A[BM*BK*QSIZE]; + alignas(128) bf16 B[BK*BN*QSIZE]; +}; + +template +__device__ void warpgroup_reg_alloc() { + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +__device__ void warpgroup_reg_dealloc() { + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +__device__ static __forceinline__ void init_barrier(uint64_t* bar, int thread_count, int transaction_count) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); +} + +__device__ static __forceinline__ void expect_bytes(uint64_t* bar, uint32_t bytes) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;\n" + :: "r"(bar_ptr), "r"(bytes)); +} + +__device__ static inline void load_async(bf16 *dst, void const* const src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + // TODO(aleksa): cluster -> cta reduces scope, and should thus improve perf - but requires PTX 8.6 + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, 0, 0}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64) + : "memory" + ); +} + +__device__ static __forceinline__ void wait(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static __forceinline__ void arrive(uint64_t* bar, uint32_t count=1) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); +} + +template +struct Schedule; + +template +struct Schedule<1, NUM_SM, BM, BN, TM, TN> { + int block; + int it; + int total_blocks_m, total_blocks_n; + + __device__ __forceinline__ Schedule(int M, int N, int _block) { + block = _block; + it = 0; + total_blocks_m = CEIL_DIV(M, BM); + total_blocks_n = CEIL_DIV(N, BN); + assert(CEIL_DIV(M, BM)%TM == 0 && total_blocks_n%TN == 0); + } + + __device__ __forceinline__ bool next(int &block_m, int& block_n) { + int num = it*NUM_SM + block; + if (num >= total_blocks_m*total_blocks_n) {return false;} + + int cur_tile = num / (TM*TN); + int cur_tile_pos = num % (TM*TN); + block_m = TM*(cur_tile / (total_blocks_n/TN)); + block_n = TN*(cur_tile % (total_blocks_n/TN)); + block_m += cur_tile_pos / TN; + block_n += cur_tile_pos % TN; + ++it; + return true; + } +}; + + +template +__global__ __launch_bounds__(NUM_THREADS) void matmulKernel7(int M, int N, int K, bf16* C, const __grid_constant__ CUtensorMap tensorMapA, const __grid_constant__ CUtensorMap tensorMapB) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int num_consumers = (NUM_THREADS / 128) - 1; + constexpr int B_WG_M = BM / num_consumers; + + extern __shared__ __align__(128) uint8_t smem[]; + SMem &s = *reinterpret_cast*>(smem); + bf16 *sA = s.A; + bf16 *sB = s.B; + // Declare barriers + __shared__ __align__(8) uint64_t full[QSIZE], empty[QSIZE]; + + const int num_blocks_k = K / BK; + int wg_idx = threadIdx.x / 128; + int tid = threadIdx.x % 128; + + if (threadIdx.x == 0) { + for (int i = 0; i < QSIZE; ++i) { + init_barrier(&full[i], 1, 0); + init_barrier(&empty[i], num_consumers, 0); + } + } + __syncthreads(); + + Schedule<1, NUM_SM, BM, BN, 16, 8> schedule(M, N, blockIdx.x); + + // Producer + if (wg_idx == 0) { + constexpr int num_regs = (num_consumers <= 2 ? 24 : 32); + warpgroup_reg_dealloc(); + if (tid == 0) { + int p = 0; + int qidx = 0; + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) { qidx = 0; p ^= 1; } + wait(&empty[qidx], p); + expect_bytes(&full[qidx], (BK*BN+BK*BM)*sizeof(bf16)); + load_async(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM); + load_async(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN); + } + } + } + } else { + constexpr int num_regs = (num_consumers == 1 ? 256 : (num_consumers == 2 ? 240 : 160)); + warpgroup_reg_alloc(); + float d[B_WG_M/WGMMA_M][WGMMA_N/16][8]; + --wg_idx; + for (int i = 0; i < QSIZE; ++i) { + if (tid == 0) arrive(&empty[i], 1); + } + int p = 0; + int qidx = 0; + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + memset(d, 0, sizeof(d)); + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) {qidx = 0; p ^= 1; }; + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + #pragma unroll + for (int bk = 0; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + if (tid == 0) arrive(&empty[qidx], 1); + } + + int lane = tid % 32, warp = tid / 32, row = warp*16 + lane / 4; + bf16 *block_C = C + num_block_n*BN*M + num_block_m*BM; + + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + int yo = m_it*WGMMA_M + wg_idx*B_WG_M; + if (row + yo + num_block_m*BM >= M) continue; + #pragma unroll + for (int w = 0; w < WGMMA_N; w+=16) if (w < N-num_block_n*BN) { + int col = w + 2*(tid % 4); + #define IDX(i, j) ((j)*M + ((i) + yo)) + #define ST(i, j, v) block_C[IDX(i, j)] = v; + + ST(row, col, d[m_it][w/16][0]); + ST(row, col+1, d[m_it][w/16][1]); + ST(row+8, col, d[m_it][w/16][2]); + ST(row+8, col+1, d[m_it][w/16][3]); + ST(row, col+8, d[m_it][w/16][4]); + ST(row, col+9, d[m_it][w/16][5]); + ST(row+8, col+8, d[m_it][w/16][6]); + ST(row+8, col+9, d[m_it][w/16][7]); + #undef IDX + #undef ST + } + } + } + } +} + +void runKernel7(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 128; + constexpr int BN = 256; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128*3; + constexpr int QSIZE = 3; + constexpr int NUM_SM = 128; + + if (_prev_m != M) { + d_tma_map_A = create_tensor_map(A, M, K); + d_tma_map_B = create_tensor_map(B, N, K); + _prev_m = M; + _prev_n = N; + _prev_k = K; + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = matmulKernel7; + size_t sMemSize = sizeof(SMem); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<>>(M, N, K, C, d_tma_map_A, d_tma_map_B); +} + +} // namespace M7 + +using M7::runKernel7; diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_8.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_8.cuh new file mode 100644 index 000000000..c38583d95 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_8.cuh @@ -0,0 +1,560 @@ + +namespace M8 { + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +__device__ uint64_t make_smem_desc(bf16* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; +} + + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +template +__device__ void warpgroup_wait() { + static_assert(N >= 0 && N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +template +__host__ static inline CUtensorMap create_tensor_map(bf16* gmem_ptr, int global_height, int global_width) { + CUtensorMap tma_map; + void* gmem_address = (void*)gmem_ptr; + static_assert(BlockMinorSize >= 64); + assert(global_width % 64 == 0); + uint64_t gmem_prob_shape[5] = {64, (uint64_t)global_height, (uint64_t)global_width/64, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16) * global_width, 64*sizeof(bf16), 0, 0, 0}; + uint32_t smem_box_shape[5] = {64, uint32_t(BlockMajorSize), uint32_t(BlockMinorSize/64), 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + &tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, gmem_address, gmem_prob_shape, + gmem_prob_stride, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); + return tma_map; +} + +CUtensorMap d_tma_map_A; +CUtensorMap d_tma_map_B; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__device__ __forceinline__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + + +template +__device__ __forceinline__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma64(float d[4][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " %34, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma32(float d[2][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}, " + " %16," + " %17," + " %18, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 32 || WGMMA_N == 64 || WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 208 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128(d, sA, sB); + if constexpr (WGMMA_N == 64) + wgmma64(d, sA, sB); + if constexpr (WGMMA_N == 32) + wgmma32(d, sA, sB); +} + +template +struct SMem { + alignas(128) bf16 A[BM*BK*QSIZE]; + alignas(128) bf16 B[BK*BN*QSIZE]; +}; + +template +__device__ void warpgroup_reg_alloc() { + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +__device__ void warpgroup_reg_dealloc() { + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +__device__ static __forceinline__ void init_barrier(uint64_t* bar, int thread_count, int transaction_count) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); +} + +__device__ static __forceinline__ void expect_bytes(uint64_t* bar, uint32_t bytes) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + :: "r"(bar_ptr), "r"(bytes)); +} + +__device__ static inline void load_async(bf16 *dst, void const* const src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64) + : "memory" + ); +} + +__device__ static __forceinline__ void wait(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static __forceinline__ void arrive(uint64_t* bar, uint32_t count=1) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); +} + +__device__ static __forceinline__ void wait_cluster(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static inline void load_async_multicast(bf16 *dst, void const* const src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx, uint16_t cluster_mask) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5}], [%2], %6;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64), "h"(cluster_mask) + : "memory" + ); +} + +__device__ void arrive_cluster(uint64_t* bar, uint32_t cta_id, uint32_t count=1) { + uint32_t smem_addr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + : + : "r"(smem_addr), "r"(cta_id), "r"(count)); +} + +template +struct Schedule; + +template +struct Schedule<1, NUM_SM, BM, BN, TM, TN> { + int block; + int it; + int total_blocks_m, total_blocks_n; + + __device__ __forceinline__ Schedule(int M, int N, int _block) { + block = _block; + it = 0; + total_blocks_m = CEIL_DIV(M, BM); + total_blocks_n = CEIL_DIV(N, BN); + assert(CEIL_DIV(M, BM)%TM == 0 && total_blocks_n%TN == 0); + } + + __device__ __forceinline__ bool next(int &block_m, int& block_n) { + int num = it*NUM_SM + block; + if (num >= total_blocks_m*total_blocks_n) {return false;} + + int cur_tile = num / (TM*TN); + int cur_tile_pos = num % (TM*TN); + block_m = TM*(cur_tile / (total_blocks_n/TN)); + block_n = TN*(cur_tile % (total_blocks_n/TN)); + block_m += cur_tile_pos / TN; + block_n += cur_tile_pos % TN; + ++it; + return true; + } +}; + +template +__global__ __launch_bounds__(NUM_THREADS) void __cluster_dims__(CLUSTER_M * CLUSTER_N, 1, 1) matmulKernel8(int M, int N, int K, bf16* C, const __grid_constant__ CUtensorMap tensorMapA, const __grid_constant__ CUtensorMap tensorMapB) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int num_consumers = (NUM_THREADS / 128) - 1; + constexpr int B_WG_M = BM / num_consumers; + constexpr int CLUSTERS = CLUSTER_M * CLUSTER_N; + assert((M / BM) % CLUSTER_M == 0); + assert((N / BN) % CLUSTER_N == 0); + + extern __shared__ __align__(128) uint8_t smem[]; + SMem &s = *reinterpret_cast*>(smem); + bf16 *sA = s.A; + bf16 *sB = s.B; + // Declare barriers + __shared__ __align__(8) uint64_t full[QSIZE], empty[QSIZE]; + uint32_t rank; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(rank) :); + + const int num_blocks_k = K / BK; + int wg_idx = threadIdx.x / 128; + int tid = threadIdx.x % 128; + + if (threadIdx.x == 0) { + for (int i = 0; i < QSIZE; ++i) { + init_barrier(&full[i], 0, 1); + init_barrier(&empty[i], 0, num_consumers*CLUSTERS); + } + } + asm volatile("barrier.cluster.arrive;\n" : :); + asm volatile("barrier.cluster.wait;\n" : :); + + Schedule<1, NUM_SM/CLUSTERS, BM*CLUSTER_M, BN*CLUSTER_N, 16/CLUSTER_M, 8/CLUSTER_N> schedule(M, N, rank); + + asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(rank) :); + uint32_t rank_m = rank / CLUSTER_N; + uint32_t rank_n = rank % CLUSTER_N; + + // Producer + if (wg_idx == 0) { + constexpr int num_regs = (num_consumers <= 2 ? 24 : 32); + warpgroup_reg_dealloc(); + if (tid == 0) { + int p = 0; + int qidx = 0; + uint32_t col_mask = 0; + for (int i = 0; i < CLUSTER_M; ++i) { + col_mask |= (1 << (i * CLUSTER_N)); + } + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) { qidx = 0; p ^= 1;} + wait(&empty[qidx], p); + + expect_bytes(&full[qidx], (BK*BN+BK*BM)*sizeof(bf16)); + if constexpr (CLUSTER_N > 1) { + uint32_t mask = ((1 << CLUSTER_N) - 1) << (rank_m * CLUSTER_N); + if (rank_n == 0) { + load_async_multicast(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM, mask); + } + } else { + load_async(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM); + } + + if constexpr (CLUSTER_M > 1) { + if (rank_m == 0) { + load_async_multicast(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN, col_mask << rank_n); + } + } else { + load_async(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN); + } + } + } + } + } else { + constexpr int num_regs = (num_consumers == 1 ? 256 : (num_consumers == 2 ? 240 : 160)); + warpgroup_reg_alloc(); + float d[B_WG_M/WGMMA_M][WGMMA_N/16][8]; + --wg_idx; + for (int qidx = 0; qidx < QSIZE; ++qidx) { + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + } + int p = 0; + int qidx = 0; + int lane = tid % 32, warp = tid / 32, row = warp*16 + lane / 4; + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + memset(d, 0, sizeof(d)); + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) {qidx = 0; p ^= 1; }; + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + #pragma unroll + for (int bk = 0; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + } + + bf16 *block_C = C + num_block_n*BN*M + num_block_m*BM; + + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + int yo = m_it*WGMMA_M + wg_idx*B_WG_M; + if (row + 8 + yo + num_block_m*BM >= M) continue; + #pragma unroll + for (int w = 0; w < WGMMA_N; w+=16) if (w+num_block_n*BN < N) { + int col = w + 2*(tid % 4); + #define IDX(i, j) ((j)*M + ((i) + yo)) + #define ST(i, j, v) block_C[IDX(i, j)] = v; + + + ST(row, col, d[m_it][w/16][0]); + ST(row, col+1, d[m_it][w/16][1]); + ST(row+8, col, d[m_it][w/16][2]); + ST(row+8, col+1, d[m_it][w/16][3]); + ST(row, col+8, d[m_it][w/16][4]); + ST(row, col+9, d[m_it][w/16][5]); + ST(row+8, col+8, d[m_it][w/16][6]); + ST(row+8, col+9, d[m_it][w/16][7]); + + + #undef IDX + #undef ST + } + } + } + } +} + +void runKernel8(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 64*2; + constexpr int BN = 256; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128*3; + constexpr int QSIZE = 3; + constexpr int CLUSTER_M = 2; + constexpr int CLUSTER_N = 1; + constexpr int NUM_SM = 128; + static_assert(NUM_SM % (CLUSTER_M*CLUSTER_N) == 0); + + if (_prev_m != M) { + d_tma_map_A = create_tensor_map(A, M, K); + d_tma_map_B = create_tensor_map(B, N, K); + _prev_m = M; + _prev_n = N; + _prev_k = K; + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = matmulKernel8; + size_t sMemSize = sizeof(SMem); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<>>(M, N, K, C, d_tma_map_A, d_tma_map_B); +} + +} // namespace M8 + +using M8::runKernel8; diff --git a/code/labs/fast_cu/upstream/h100/matmul/matmul_9.cuh b/code/labs/fast_cu/upstream/h100/matmul/matmul_9.cuh new file mode 100644 index 000000000..1df7301ad --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/matmul_9.cuh @@ -0,0 +1,600 @@ + +namespace M9 { + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +__device__ uint64_t make_smem_desc(bf16* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; +} + + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +template +__device__ void warpgroup_wait() { + static_assert(N >= 0 && N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +template +__host__ static inline CUtensorMap create_tensor_map(bf16* gmem_ptr, int global_height, int global_width) { + CUtensorMap tma_map; + void* gmem_address = (void*)gmem_ptr; + static_assert(BlockMinorSize >= 64); + assert(global_width % 64 == 0); + uint64_t gmem_prob_shape[5] = {64, (uint64_t)global_height, (uint64_t)global_width/64, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16) * global_width, 64*sizeof(bf16), 0, 0, 0}; + uint32_t smem_box_shape[5] = {64, uint32_t(BlockMajorSize), uint32_t(BlockMinorSize/64), 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + &tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, gmem_address, gmem_prob_shape, + gmem_prob_stride, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); + return tma_map; +} + +CUtensorMap d_tma_map_A; +CUtensorMap d_tma_map_B; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__device__ __forceinline__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + + +template +__device__ __forceinline__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma64(float d[4][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " %34, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma32(float d[2][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}, " + " %16," + " %17," + " %18, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 32 || WGMMA_N == 64 || WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 208 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128(d, sA, sB); + if constexpr (WGMMA_N == 64) + wgmma64(d, sA, sB); + if constexpr (WGMMA_N == 32) + wgmma32(d, sA, sB); +} + +template +struct SMem { + alignas(128) bf16 A[BM*BK*QSIZE]; + alignas(128) bf16 B[BK*BN*QSIZE]; +}; + +template +__device__ void warpgroup_reg_alloc() { + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +__device__ void warpgroup_reg_dealloc() { + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +__device__ static __forceinline__ void init_barrier(uint64_t* bar, int thread_count, int transaction_count) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); +} + +__device__ static __forceinline__ void expect_bytes(uint64_t* bar, uint32_t bytes) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + :: "r"(bar_ptr), "r"(bytes)); +} + +__device__ static inline void load_async(bf16 *dst, void const* const src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64) + : "memory" + ); +} + +__device__ static __forceinline__ void wait(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static __forceinline__ void arrive(uint64_t* bar, uint32_t count=1) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); +} + +__device__ static __forceinline__ void wait_cluster(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static inline void load_async_multicast(bf16 *dst, void const* const src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx, uint16_t cluster_mask) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5}], [%2], %6;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64), "h"(cluster_mask) + : "memory" + ); +} + +__device__ void arrive_cluster(uint64_t* bar, uint32_t cta_id, uint32_t count=1) { + uint32_t smem_addr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + : + : "r"(smem_addr), "r"(cta_id), "r"(count)); +} + +template +struct Schedule; + +template +struct Schedule<1, NUM_SM, BM, BN, TM, TN> { + int block; + int it; + int total_blocks_m, total_blocks_n; + + __device__ __forceinline__ Schedule(int M, int N, int _block) { + block = _block; + it = 0; + total_blocks_m = CEIL_DIV(M, BM); + total_blocks_n = CEIL_DIV(N, BN); + assert(CEIL_DIV(M, BM)%TM == 0 && total_blocks_n%TN == 0); + } + + __device__ __forceinline__ bool next(int &block_m, int& block_n) { + int num = it*NUM_SM + block; + if (num >= total_blocks_m*total_blocks_n) {return false;} + + int cur_tile = num / (TM*TN); + int cur_tile_pos = num % (TM*TN); + block_m = TM*(cur_tile / (total_blocks_n/TN)); + block_n = TN*(cur_tile % (total_blocks_n/TN)); + block_m += cur_tile_pos / TN; + block_n += cur_tile_pos % TN; + ++it; + return true; + } +}; + +template +__global__ __launch_bounds__(NUM_THREADS) void __cluster_dims__(CLUSTER_M * CLUSTER_N, 1, 1) matmulKernel9(int M, int N, int K, bf16* C, const __grid_constant__ CUtensorMap tensorMapA, const __grid_constant__ CUtensorMap tensorMapB) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int num_consumers = (NUM_THREADS / 128) - 1; + constexpr int B_WG_M = BM / num_consumers; + constexpr int CLUSTERS = CLUSTER_M * CLUSTER_N; + assert((M / BM) % CLUSTER_M == 0); + assert((N / BN) % CLUSTER_N == 0); + + extern __shared__ __align__(128) uint8_t smem[]; + SMem &s = *reinterpret_cast*>(smem); + bf16 *sA = s.A; + bf16 *sB = s.B; + // Declare barriers + __shared__ __align__(8) uint64_t full[QSIZE], empty[QSIZE]; + uint32_t rank; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(rank) :); + + const int num_blocks_k = K / BK; + int wg_idx = threadIdx.x / 128; + int tid = threadIdx.x % 128; + + if (threadIdx.x == 0) { + for (int i = 0; i < QSIZE; ++i) { + init_barrier(&full[i], 0, 1); + init_barrier(&empty[i], 0, num_consumers*CLUSTERS); + } + } + asm volatile("barrier.cluster.arrive;\n" : :); + asm volatile("barrier.cluster.wait;\n" : :); + + Schedule<1, NUM_SM/CLUSTERS, BM*CLUSTER_M, BN*CLUSTER_N, 16/CLUSTER_M, 8/CLUSTER_N> schedule(M, N, rank); + + asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(rank) :); + uint32_t rank_m = rank / CLUSTER_N; + uint32_t rank_n = rank % CLUSTER_N; + + // Producer + if (wg_idx == 0) { + constexpr int num_regs = (num_consumers <= 2 ? 24 : 32); + warpgroup_reg_dealloc(); + if (tid == 0) { + int p = 0; + int qidx = 0; + uint32_t col_mask = 0; + for (int i = 0; i < CLUSTER_M; ++i) { + col_mask |= (1 << (i * CLUSTER_N)); + } + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) { qidx = 0; p ^= 1;} + wait(&empty[qidx], p); + + expect_bytes(&full[qidx], (BK*BN+BK*BM)*sizeof(bf16)); + if constexpr (CLUSTER_N > 1) { + uint32_t mask = ((1 << CLUSTER_N) - 1) << (rank_m * CLUSTER_N); + if (rank_n == 0) { + load_async_multicast(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM, mask); + } + } else { + load_async(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM); + } + + if constexpr (CLUSTER_M > 1) { + if (rank_m == 0) { + load_async_multicast(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN, col_mask << rank_n); + } + } else { + load_async(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN); + } + } + } + } + } else { + constexpr int num_regs = (num_consumers == 1 ? 256 : (num_consumers == 2 ? 240 : 160)); + warpgroup_reg_alloc(); + float d[B_WG_M/WGMMA_M][WGMMA_N/16][8]; + --wg_idx; + for (int qidx = 0; qidx < QSIZE; ++qidx) { + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + } + int p = 0; + int qidx = 0; + int lane = tid % 32, warp = tid / 32, row = warp*16 + lane / 4; + int num_block_m, num_block_n; + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + { + if (qidx == QSIZE) {qidx = 0; p ^= 1; }; + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + { + wgmma(d[m_it], &wgmma_sA[0], &wgmma_sB[0]); + #pragma unroll + for (int k_it = 1; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + #pragma unroll + for (int bk = 64; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + ++qidx; + } + for (int block_k_iter = 1; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) {qidx = 0; p ^= 1; }; + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*(m_it + wg_idx*B_WG_M/WGMMA_M)*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + #pragma unroll + for (int bk = 0; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + } + + // __nv_bfloat162* block_C = reinterpret_cast<__nv_bfloat162*>(C + num_block_m*BM*N + num_block_n*BN); + bf16 *block_C = C + num_block_n*BN*M + num_block_m*BM; + + #pragma unroll + for (int m_it = 0; m_it < B_WG_M/WGMMA_M; ++m_it) { + int yo = m_it*WGMMA_M + wg_idx*B_WG_M; + if (row + 8 + yo + num_block_m*BM >= M) continue; + #pragma unroll + for (int w = 0; w < WGMMA_N; w+=16) if (w+num_block_n*BN < N) { + // int col = 8*w + (tid % 4); + // #define IDX(i, j) ((((i) + yo)*N/2 + (j))) + // block_C[IDX(row, col)] = __halves2bfloat162(d[m_it][w][0], d[m_it][w][1]); + // block_C[IDX(row + 8, col)] = __halves2bfloat162(d[m_it][w][2], d[m_it][w][3]); + // block_C[IDX(row, col + 4)] = __halves2bfloat162(d[m_it][w][4], d[m_it][w][5]); + // block_C[IDX(row + 8, col + 4)] = __halves2bfloat162(d[m_it][w][6], d[m_it][w][7]); + // #undef IDX + int col = w + 2*(tid % 4); + #define IDX(i, j) ((j)*M + ((i) + yo)) + #define ST(i, j, v) __stwt(&block_C[IDX(i, j)], v); + + ST(row+8, col, d[m_it][w/16][2]); + ST(row, col, d[m_it][w/16][0]); + + ST(row+8, col+1, d[m_it][w/16][3]); + ST(row, col+1, d[m_it][w/16][1]); + + ST(row+8, col+8, d[m_it][w/16][6]); + ST(row, col+8, d[m_it][w/16][4]); + + ST(row+8, col+9, d[m_it][w/16][7]); + ST(row, col+9, d[m_it][w/16][5]); + + #undef IDX + #undef ST + } + } + } + } +} + +void runKernel9(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 64*2; + constexpr int BN = 256; + constexpr int BK = 64; + constexpr int NUM_THREADS = 128*3; + constexpr int QSIZE = 3; + constexpr int CLUSTER_M = 2; + constexpr int CLUSTER_N = 1; + constexpr int NUM_SM = 128; + static_assert(NUM_SM % (CLUSTER_M*CLUSTER_N) == 0); + + if (_prev_m != M) { + d_tma_map_A = create_tensor_map(A, M, K); + d_tma_map_B = create_tensor_map(B, N, K); + _prev_m = M; + _prev_n = N; + _prev_k = K; + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = matmulKernel9; + size_t sMemSize = sizeof(SMem); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<>>(M, N, K, C, d_tma_map_A, d_tma_map_B); +} + +} // namespace M9 + +using M9::runKernel9; diff --git a/code/labs/fast_cu/upstream/h100/matmul/pingpong_experimental.cuh b/code/labs/fast_cu/upstream/h100/matmul/pingpong_experimental.cuh new file mode 100644 index 000000000..6d95fda21 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/matmul/pingpong_experimental.cuh @@ -0,0 +1,761 @@ + +namespace M10 { + +// using barrier = cuda::barrier; +namespace cde = cuda::device::experimental; + +__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) { return (((x) & 0x3FFFF) >> 0x4); } + +__device__ uint64_t make_smem_desc(bf16* ptr) { + uint32_t addr = static_cast(__cvta_generic_to_shared(ptr)); + uint64_t desc = 0x0000000000000000; + desc |= matrix_descriptor_encode(addr); + desc |= matrix_descriptor_encode((uint64_t)16) << 16; + desc |= matrix_descriptor_encode((uint64_t)1024) << 32; + desc |= 1llu << 62; // 128B swizzle + return desc; +} + + +__device__ void warpgroup_arrive() { + asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory"); +} + +__device__ void warpgroup_commit_batch() { + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +} + +template +__device__ void warpgroup_wait() { + static_assert(N >= 0 && N <= 7, "WGMMA wait: N must be in range [0, 7]"); + asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory"); +} + +template +__host__ static inline CUtensorMap create_tensor_map(bf16* gmem_ptr, int global_height, int global_width) { + CUtensorMap tma_map; + void* gmem_address = (void*)gmem_ptr; + static_assert(BlockMinorSize >= 64); + assert(global_width % 64 == 0); + uint64_t gmem_prob_shape[5] = {64, (uint64_t)global_height, (uint64_t)global_width/64, 1, 1}; + uint64_t gmem_prob_stride[5] = {sizeof(bf16) * global_width, 64*sizeof(bf16), 0, 0, 0}; + uint32_t smem_box_shape[5] = {64, uint32_t(BlockMajorSize), uint32_t(BlockMinorSize/64), 1, 1}; + uint32_t smem_box_stride[5] = {1, 1, 1, 1, 1}; + + CUresult result = cuTensorMapEncodeTiled( + &tma_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, gmem_address, gmem_prob_shape, + gmem_prob_stride, smem_box_shape, smem_box_stride, CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_128B, CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + assert(result == CUDA_SUCCESS); + return tma_map; +} + +CUtensorMap d_tma_map_A; +CUtensorMap d_tma_map_B; +int _prev_m=0, _prev_n=0, _prev_k=0; + +template +__device__ __forceinline__ void wgmma256(float d[16][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " %130, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]), + "+f"(d[12][0]), "+f"(d[12][1]), "+f"(d[12][2]), "+f"(d[12][3]), "+f"(d[12][4]), "+f"(d[12][5]), "+f"(d[12][6]), "+f"(d[12][7]), + "+f"(d[13][0]), "+f"(d[13][1]), "+f"(d[13][2]), "+f"(d[13][3]), "+f"(d[13][4]), "+f"(d[13][5]), "+f"(d[13][6]), "+f"(d[13][7]), + "+f"(d[14][0]), "+f"(d[14][1]), "+f"(d[14][2]), "+f"(d[14][3]), "+f"(d[14][4]), "+f"(d[14][5]), "+f"(d[14][6]), "+f"(d[14][7]), + "+f"(d[15][0]), "+f"(d[15][1]), "+f"(d[15][2]), "+f"(d[15][3]), "+f"(d[15][4]), "+f"(d[15][5]), "+f"(d[15][6]), "+f"(d[15][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma192(float d[12][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}, " + " %96," + " %97," + " %98, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), "+f"(d[0][6]), "+f"(d[0][7]), + "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), + "+f"(d[2][0]), "+f"(d[2][1]), "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), "+f"(d[3][6]), "+f"(d[3][7]), + "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), + "+f"(d[5][0]), "+f"(d[5][1]), "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), "+f"(d[6][6]), "+f"(d[6][7]), + "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]), + "+f"(d[8][0]), "+f"(d[8][1]), "+f"(d[8][2]), "+f"(d[8][3]), "+f"(d[8][4]), "+f"(d[8][5]), "+f"(d[8][6]), "+f"(d[8][7]), + "+f"(d[9][0]), "+f"(d[9][1]), "+f"(d[9][2]), "+f"(d[9][3]), "+f"(d[9][4]), "+f"(d[9][5]), "+f"(d[9][6]), "+f"(d[9][7]), + "+f"(d[10][0]), "+f"(d[10][1]), "+f"(d[10][2]), "+f"(d[10][3]), "+f"(d[10][4]), "+f"(d[10][5]), "+f"(d[10][6]), "+f"(d[10][7]), + "+f"(d[11][0]), "+f"(d[11][1]), "+f"(d[11][2]), "+f"(d[11][3]), "+f"(d[11][4]), "+f"(d[11][5]), "+f"(d[11][6]), "+f"(d[11][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + + +template +__device__ __forceinline__ void wgmma128(float d[8][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " %66, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]), "+f"(d[4][0]), "+f"(d[4][1]), "+f"(d[4][2]), "+f"(d[4][3]), + "+f"(d[4][4]), "+f"(d[4][5]), "+f"(d[4][6]), "+f"(d[4][7]), "+f"(d[5][0]), "+f"(d[5][1]), + "+f"(d[5][2]), "+f"(d[5][3]), "+f"(d[5][4]), "+f"(d[5][5]), "+f"(d[5][6]), "+f"(d[5][7]), + "+f"(d[6][0]), "+f"(d[6][1]), "+f"(d[6][2]), "+f"(d[6][3]), "+f"(d[6][4]), "+f"(d[6][5]), + "+f"(d[6][6]), "+f"(d[6][7]), "+f"(d[7][0]), "+f"(d[7][1]), "+f"(d[7][2]), "+f"(d[7][3]), + "+f"(d[7][4]), "+f"(d[7][5]), "+f"(d[7][6]), "+f"(d[7][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma64(float d[4][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " %34, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]), + "+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]), + "+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]), + "+f"(d[3][6]), "+f"(d[3][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma32(float d[2][8], bf16* sA, bf16* sB) { + uint64_t desc_a = make_smem_desc(&sA[0]); + uint64_t desc_b = make_smem_desc(&sB[0]); + asm volatile( + "{\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}, " + " %16," + " %17," + " %18, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]), + "+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]), + "+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]) + : "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)), + "n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB))); +} + +template +__device__ __forceinline__ void wgmma(float d[WGMMA_N/16][8], bf16* sA, bf16* sB) { + static_assert(WGMMA_N == 32 || WGMMA_N == 64 || WGMMA_N == 128 || WGMMA_N == 192 || WGMMA_N == 208 || WGMMA_N == 256); + if constexpr (WGMMA_N == 256) + wgmma256(d, sA, sB); + if constexpr (WGMMA_N == 192) + wgmma192(d, sA, sB); + if constexpr (WGMMA_N == 128) + wgmma128(d, sA, sB); + if constexpr (WGMMA_N == 64) + wgmma64(d, sA, sB); + if constexpr (WGMMA_N == 32) + wgmma32(d, sA, sB); +} + +template +__device__ void warpgroup_reg_alloc() { + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +template +__device__ void warpgroup_reg_dealloc() { + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount)); +} + +__device__ static __forceinline__ void init_barrier(uint64_t* bar, int thread_count, int transaction_count) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.init.shared::cta.b64 [%0], %1;\n" + :: "r"(bar_ptr), "r"(thread_count+transaction_count) + ); +} + +__device__ static __forceinline__ void expect_bytes(uint64_t* bar, uint32_t bytes) { + uint32_t bar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;\n" + :: "r"(bar_ptr), "r"(bytes)); +} + +__device__ static inline void load_async(bf16 *dst, void const* const src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5}], [%2];" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64) + : "memory" + ); +} + +__device__ static __forceinline__ void wait(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static __forceinline__ void arrive(uint64_t* bar, uint32_t count=1) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n" + : + : "r"(mbar_ptr), "r"(count) + : "memory" + ); +} + +__device__ static __forceinline__ void wait_cluster(uint64_t* bar, int kPhaseBit) { + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile ( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(mbar_ptr), + "r"(kPhaseBit) + ); +} + +__device__ static inline void load_async_multicast(bf16 *dst, void const* const src_tma_map, uint64_t* bar, int global_col_idx, int global_row_idx, uint16_t cluster_mask) { + uint64_t tma_ptr = reinterpret_cast(src_tma_map); + uint32_t mbar_ptr = static_cast(__cvta_generic_to_shared(bar)); + uint32_t dst_ptr = static_cast(__cvta_generic_to_shared(dst)); + + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5}], [%2], %6;" + : + : "r"(dst_ptr), "l"(tma_ptr), "r"(mbar_ptr), + "n"(0), "r"(global_row_idx), "r"(global_col_idx/64), "h"(cluster_mask) + : "memory" + ); +} + +__device__ void arrive_cluster(uint64_t* bar, uint32_t cta_id, uint32_t count=1) { + uint32_t smem_addr = static_cast(__cvta_generic_to_shared(bar)); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + : + : "r"(smem_addr), "r"(cta_id), "r"(count)); +} + +template +struct Schedule; + +constexpr int SPACE_LEN = 128; +int *_dspace; + +template +struct Schedule<0, NUM_SM, BM, BN, TM, TN> { + int st, en; + int blocks_n; + + __device__ __forceinline__ Schedule(int M, int N, int block, int *space) { + int total_blocks = M*N/(BM*BN); + int blocks_per_sm = total_blocks / NUM_SM; + int extra_blocks = total_blocks % NUM_SM; + if (block < extra_blocks) { + st = block*(blocks_per_sm + 1); + en = st + blocks_per_sm + 1; + } else { + st = extra_blocks + block*blocks_per_sm; + en = st + blocks_per_sm; + } + blocks_n = CEIL_DIV(N, BN); + } + + __device__ __forceinline__ bool next(int &block_m, int& block_n) { + if (en == st) { + return false; + } + block_m = st / blocks_n; + block_n = st % blocks_n; + st += 1; + return true; + } +}; + +template +struct Schedule<1, NUM_SM, BM, BN, TM, TN> { + int block; + int it; + int total_blocks_m, total_blocks_n; + + __device__ __forceinline__ Schedule(int M, int N, int _block, int *space) { + block = _block; + it = 0; + total_blocks_m = CEIL_DIV(M, BM); + total_blocks_n = CEIL_DIV(N, BN); + assert(CEIL_DIV(M, BM)%TM == 0 && total_blocks_n%TN == 0); + } + + __device__ __forceinline__ bool next(int &block_m, int& block_n) { + int num = it*NUM_SM + block; + if (num >= total_blocks_m*total_blocks_n) {return false;} + + int cur_tile = num / (TM*TN); + int cur_tile_pos = num % (TM*TN); + block_m = TM*(cur_tile / (total_blocks_n/TN)); + block_n = TN*(cur_tile % (total_blocks_n/TN)); + block_m += cur_tile_pos / TN; + block_n += cur_tile_pos % TN; + ++it; + return true; + } +}; + +template +struct Schedule<2, NUM_SM, BM, BN, TM, TN> { + int it; + int *space; + + __device__ __forceinline__ Schedule(int M, int N, int block, int *_space) { + it = 0; + space = _space; + } + + __device__ __forceinline__ bool next(int &block_m, int& block_n) { + if (it >= SPACE_LEN) { + return false; + } + int now = space[it]; + if (now == -1) { + return false; + } + block_m = now >> 16; + block_n = (now & ((1<<16)-1)); + ++it; + return true; + } +}; + +template +struct SMem { + alignas(128) bf16 A[BM*BK*QSIZE]; + alignas(128) bf16 B[BK*BN*QSIZE]; + alignas(8) uint64_t full[QSIZE], empty[QSIZE], tc[2]; + int space[2][SPACE_LEN]; +}; + +template +__global__ __launch_bounds__(NUM_THREADS) void __cluster_dims__(CLUSTER_M * CLUSTER_N, 1, 1) matmulKernel10(int M, int N, int K, bf16* C, const __grid_constant__ CUtensorMap tensorMapA, const __grid_constant__ CUtensorMap tensorMapB, int* dspace) { + constexpr int WGMMA_M = 64, WGMMA_K = 16, WGMMA_N=BN; + constexpr int num_consumers = (NUM_THREADS / 128) - 1; + constexpr int CLUSTERS = CLUSTER_M * CLUSTER_N; + assert((M / BM) % CLUSTER_M == 0); + assert((N / BN) % CLUSTER_N == 0); + + extern __shared__ __align__(128) uint8_t smem[]; + SMem &s = *reinterpret_cast*>(smem); + bf16 *sA = s.A; + bf16 *sB = s.B; + uint64_t* full = s.full; + uint64_t* empty = s.empty; + int* space[2]; + space[0] = s.space[0], space[1] = s.space[1]; + + uint32_t rank; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(rank) :); + if (threadIdx.x < SPACE_LEN) { + space[0][threadIdx.x] = dspace[2*rank*SPACE_LEN+threadIdx.x]; + space[1][threadIdx.x] = dspace[(2*rank+1)*SPACE_LEN+threadIdx.x]; + } + + const int num_blocks_k = K / BK; + int wg_idx = threadIdx.x / 128; + int tid = threadIdx.x % 128; + + if (threadIdx.x == 0) { + for (int qidx = 0; qidx < QSIZE; ++qidx) { + init_barrier(&full[qidx], 0, 1); + init_barrier(&empty[qidx], 0, CLUSTERS); + arrive(&empty[qidx], 2); + } + for (int c = 0; c < num_consumers; ++c) { + init_barrier(&s.tc[c], 0, 1); + } + arrive(&s.tc[0]); + } + asm volatile("barrier.cluster.arrive;\n" : :); + asm volatile("barrier.cluster.wait;\n" : :); + + asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(rank) :); + uint32_t rank_m = rank / CLUSTER_N; + uint32_t rank_n = rank % CLUSTER_N; + + // Producer + if (wg_idx == 0) { + constexpr int num_regs = (num_consumers <= 2 ? 24 : 32); + warpgroup_reg_dealloc(); + if (tid == 0) { + int p = 0; + int qidx = 0; + uint32_t col_mask = 0; + for (int i = 0; i < CLUSTER_M; ++i) { + col_mask |= (1 << (i * CLUSTER_N)); + } + Schedule<2, NUM_SM/CLUSTERS, BM*CLUSTER_M, BN*CLUSTER_N, 16/CLUSTER_M, 8/CLUSTER_N> schedule[2] = { + Schedule<2, NUM_SM/CLUSTERS, BM*CLUSTER_M, BN*CLUSTER_N, 16/CLUSTER_M, 8/CLUSTER_N>(M, N, rank*2, space[0]), + Schedule<2, NUM_SM/CLUSTERS, BM*CLUSTER_M, BN*CLUSTER_N, 16/CLUSTER_M, 8/CLUSTER_N>(M, N, rank*2+1, space[1]) + }; + bool can = true; + int num_block_m, num_block_n; + while (can) { + can = false; + // Schedule blocks for warpgroups one after another + #pragma unroll + for (int widx = 0; widx < num_consumers; ++widx) if (schedule[widx].next(num_block_m, num_block_n)) { + can = true; + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + for (int block_k_iter = 0; block_k_iter < num_blocks_k; ++block_k_iter, ++qidx) { + if (qidx == QSIZE) { qidx = 0; p ^= 1; } + // printf("Try filling qidx %d\n", qidx); + wait(&empty[qidx], p); + // printf("Filling qidx %d\n", qidx); + expect_bytes(&full[qidx], (BK*BM + BK*BN)*sizeof(bf16)); + if constexpr (CLUSTER_M > 1) { + if (rank_m == 0) { + load_async_multicast(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN, col_mask << rank_n); + } + } else { + load_async(&sB[qidx*BK*BN], &tensorMapB, &full[qidx], block_k_iter*BK, num_block_n*BN); + } + + if constexpr (CLUSTER_N > 1) { + uint32_t mask = ((1 << CLUSTER_N) - 1) << (rank_m * CLUSTER_N); + if (rank_n == 0) { + load_async_multicast(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM, mask); + } + } else { + load_async(&sA[qidx*BK*BM], &tensorMapA, &full[qidx], block_k_iter*BK, num_block_m*BM); + } + } + + } + } + } + } else { + constexpr int num_regs = (num_consumers == 1 ? 256 : (num_consumers == 2 ? 240 : 160)); + warpgroup_reg_alloc(); + float d[BM/WGMMA_M][WGMMA_N/16][8]; + --wg_idx; + int p = 0, tcp = 0; + int qidx = 0; + int lane = tid % 32, warp = tid / 32, row = warp*16 + lane / 4; + int num_block_m, num_block_n; + Schedule<2, NUM_SM/CLUSTERS, BM*CLUSTER_M, BN*CLUSTER_N, 16/CLUSTER_M, 8/CLUSTER_N> schedule(M, N, rank*2+(wg_idx-1), space[wg_idx]); + while (schedule.next(num_block_m, num_block_n)) { + num_block_n = num_block_n * CLUSTER_N + rank_n; + num_block_m = num_block_m * CLUSTER_M + rank_m; + if (wg_idx == 1) { + qidx += num_blocks_k; + p ^= (qidx / QSIZE) % 2; + qidx %= QSIZE; + } + // if (tid == 0) printf("Trying MMA %d\n", wg_idx); + wait(&s.tc[wg_idx], tcp); + // if (tid == 0) printf("Starting MMA %d\n", wg_idx); + { + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < BM/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*m_it*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + { + wgmma(d[m_it], &wgmma_sA[0], &wgmma_sB[0]); + #pragma unroll + for (int k_it = 1; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + #pragma unroll + for (int bk = 64; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + ++qidx; if (qidx == QSIZE) { qidx = 0; p ^= 1; } + } + for (int block_k_iter = 1; block_k_iter < num_blocks_k; ++block_k_iter) { + wait(&full[qidx], p); + warpgroup_arrive(); + #pragma unroll + for (int m_it = 0; m_it < BM/WGMMA_M; ++m_it) { + bf16 *wgmma_sA = sA + qidx*BK*BM + 64*m_it*WGMMA_M; + bf16 *wgmma_sB = sB + qidx*BK*BN; + #pragma unroll + for (int bk = 0; bk < BK; bk += 64) { + #pragma unroll + for (int k_it = 0; k_it < 64/WGMMA_K; ++k_it) { + wgmma(d[m_it], &wgmma_sA[k_it*WGMMA_K], &wgmma_sB[k_it*WGMMA_K]); + } + wgmma_sA += 64*BM; + wgmma_sB += 64*BN; + } + } + warpgroup_commit_batch(); + warpgroup_wait<0>(); + if (tid < CLUSTERS) arrive_cluster(&empty[qidx], tid); + ++qidx; if (qidx == QSIZE) { qidx = 0; p ^= 1; } + } + if (tid == 0) arrive(&s.tc[wg_idx ^ 1]); + tcp ^= 1; + if (wg_idx == 0) { + qidx += num_blocks_k; + p ^= (qidx / QSIZE) % 2; + qidx %= QSIZE; + } + // if (tid == 0) printf("Finished MMA %d\n", wg_idx); + + // __nv_bfloat162* block_C = reinterpret_cast<__nv_bfloat162*>(C + num_block_m*BM*N + num_block_n*BN); + bf16 *block_C = C + num_block_n*BN*M + num_block_m*BM; + + #pragma unroll + for (int m_it = 0; m_it < BM/WGMMA_M; ++m_it) { + int yo = m_it*WGMMA_M; + if (row + 8 + yo + num_block_m*BM >= M) continue; + #pragma unroll + for (int w = 0; w < WGMMA_N; w+=16) if (w+num_block_n*BN < N) { + // int col = 8*w + (tid % 4); + // #define IDX(i, j) ((((i) + yo)*N/2 + (j))) + // block_C[IDX(row, col)] = __halves2bfloat162(d[m_it][w][0], d[m_it][w][1]); + // block_C[IDX(row + 8, col)] = __halves2bfloat162(d[m_it][w][2], d[m_it][w][3]); + // block_C[IDX(row, col + 4)] = __halves2bfloat162(d[m_it][w][4], d[m_it][w][5]); + // block_C[IDX(row + 8, col + 4)] = __halves2bfloat162(d[m_it][w][6], d[m_it][w][7]); + // #undef IDX + int col = w + 2*(tid % 4); + #define IDX(i, j) ((j)*M + ((i) + yo)) + #define ST(i, j, v) __stwt(&block_C[IDX(i, j)], v); + + ST(row, col, d[m_it][w/16][0]); + ST(row+8, col, d[m_it][w/16][2]); + ST(row, col+1, d[m_it][w/16][1]); + ST(row+8, col+1, d[m_it][w/16][3]); + ST(row, col+8, d[m_it][w/16][4]); + ST(row+8, col+8, d[m_it][w/16][6]); + ST(row, col+9, d[m_it][w/16][5]); + ST(row+8, col+9, d[m_it][w/16][7]); + #undef IDX + #undef ST + } + } + } + } +} + +// Rotate/flip quadrant appropriately +void rot(int n, int& x, int& y, int rx, int ry) { + if (ry == 0) { + if (rx == 1) { + x = n-1 - x; + y = n-1 - y; + } + // Swap x and y + int t = x; + x = y; + y = t; + } +} + +// Convert distance along curve to (x,y) point +void d2xy(int n, int d, int& x, int& y) { + int rx, ry, s, t = d; + x = y = 0; + for (s = 1; s < n; s *= 2) { + rx = 1 & (t/2); + ry = 1 & (t ^ rx); + rot(s, x, y, rx, ry); + x += s * rx; + y += s * ry; + t /= 4; + } +} + +void createHilbert(int M, int N, int CORES, int *space) { + int dim = (1 << (32 - __builtin_clz(max(M, N) - 1))); + int core = 0, loc = 0; + std::vector v(dim, std::string(dim, '.')); + memset(space, -1, sizeof(int)*CORES*SPACE_LEN); + int total = 0; + for (int i = 0; i < dim*dim; ++i) { + int x, y; + d2xy(dim, i, x, y); + if (x < M && y < N) { + assert(loc < SPACE_LEN); + assert(v[x][y] == '.'); + v[x][y] = '*'; + ++total; + space[core*SPACE_LEN+loc] = ((x << 16) | y); + ++core; + if (core == CORES) {core = 0; loc++;} + } + } + // for (int i = 0; i < CORES; ++i) { + // printf("%d: ", i); + // for (int j = 0; j < SPACE_LEN; ++j) { + // int value = space[i*SPACE_LEN+j]; + // if (value == -1) break; + // printf("[%d,%d] ", value>>16, value&((1<<16)-1)); + // } + // printf("\n"); + // } + assert(total == M*N); + // for (int i = 0; i < dim; ++i) { + // std::cout << v[i] << std::endl; + // } +} + +// Factor: 2263.915771, Times: 4096, Load: 1641.335504, Compute: 1189.903631, Store: 9115.878662 +void runKernel10(int M, int N, int K, bf16 *A, bf16 *B, bf16 *C, int *DB) { + constexpr int BM = 128; + constexpr int BN = 192; + constexpr int BK = 128; + constexpr int NUM_THREADS = 128*3; + constexpr int QSIZE = 2; + constexpr int CLUSTER_M = 2; + constexpr int CLUSTER_N = 1; + constexpr int NUM_SM = 128; + static_assert(NUM_SM % (CLUSTER_M*CLUSTER_N) == 0); + + if (_prev_m != M) { + d_tma_map_A = create_tensor_map(A, M, K); + d_tma_map_B = create_tensor_map(B, N, K); + _prev_m = M; + _prev_n = N; + _prev_k = K; + int *space; + space = (int*)malloc(sizeof(int)*NUM_SM*SPACE_LEN); + createHilbert(CEIL_DIV(M, BM*CLUSTER_M), CEIL_DIV(N, BN*CLUSTER_N), NUM_SM*2/CLUSTER_M/CLUSTER_N, space); + cudaCheck(cudaMalloc((void **)&_dspace, sizeof(int)*NUM_SM*SPACE_LEN)); + cudaCheck(cudaMemcpy(_dspace, space, sizeof(int)*NUM_SM*SPACE_LEN, cudaMemcpyHostToDevice)); + } + // Assert cached values are of same size + assert (M == _prev_m && N == _prev_n && K == _prev_k); + auto* kernel = matmulKernel10; + size_t sMemSize = sizeof(SMem); + cudaCheck(cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize)); + + kernel<<>>(M, N, K, C, d_tma_map_A, d_tma_map_B, _dspace); +} + +} // namespace M10 + +using M10::runKernel10; diff --git a/code/labs/fast_cu/upstream/h100/sum.cu b/code/labs/fast_cu/upstream/h100/sum.cu new file mode 100644 index 000000000..649589df0 --- /dev/null +++ b/code/labs/fast_cu/upstream/h100/sum.cu @@ -0,0 +1,332 @@ +/** +make sum && ./sum + +N = 1<<30, GPU = H100: +<---------------| kernel1 |---------------> +Bandwidth: 3232.76 GB/s +Time taken: 1.32858 ms +Sum: 1609549285 + +<---------------| kernel2 |---------------> +Bandwidth: 3239.69 GB/s +Time taken: 1.32573 ms +Sum: 1609549285 + +<---------------| kernel3 |---------------> +Bandwidth: 3240.11 GB/s +Time taken: 1.32556 ms +Sum: 1609549285 + +<---------------| kernel4 |---------------> +Bandwidth: 3241.47 GB/s +Time taken: 1.32501 ms +Sum: 1609549285 + +<---------------| cub |---------------> +Bandwidth: 3193.36 GB/s +Time taken: 1.34497 ms +Sum: 1609549285 +*/ + +#include +#include +#include + +const int WARP_SIZE = 32; + +int ceilDiv(int N, int D) { + return (N + D - 1) / D; +} + +__device__ void warpReduce1(volatile int *sdata, int tid) { + #pragma unroll + for (int i = WARP_SIZE; i >= 1; i >>= 1) { + sdata[tid] += sdata[tid + i]; + } +} + +// Inspired from https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf +template +__global__ void sumKernel1(int *d_in, int *d_out, int n) { + __shared__ int sdata[BlockSize]; + + unsigned int tid = threadIdx.x; + unsigned int i = blockIdx.x*BlockSize*Batch + tid; + + // Set output to 0 + if (i == 0) *d_out = 0; + + // Read Batch elements(strided) + float value = i < n ? d_in[i] : 0; + #pragma unroll + for (int j = 1; j < Batch; ++j) { + if (i + j*BlockSize < n) value += d_in[i + j*BlockSize]; + } + // Save in shmem and compute sum(before warp level) + sdata[tid] = value; + __syncthreads(); + + #pragma unroll + for (int s = 512; s > WARP_SIZE; s >>= 1) { + if (BlockSize >= s * 2) { + if (tid < s) sdata[tid] += sdata[tid + s]; + __syncthreads(); + } + } + + // Compute warp sum + if (tid < WARP_SIZE) warpReduce1(sdata, tid); + // Final addition to global memory + if (tid == 0) atomicAdd(d_out, sdata[0]); +} + +inline int __device__ warpReduce2(int sum) { + #pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + sum += __shfl_down_sync(0xFFFFFFFF, sum, offset); + } + return sum; +} + +template +__global__ void sumKernel2(int4 *d_in, int *d_out, int n) { + // This kernel assumes 32 warps in block + static_assert(BlockSize == 1024); + __shared__ __align__(16) int sdata[WARP_SIZE]; + + unsigned int tid = threadIdx.x; + unsigned int i = blockIdx.x*BlockSize*Batch + tid; + if (i == 0) { + *d_out = 0; + } + + int4 value = i < n ? d_in[i] : make_int4(0, 0, 0, 0); + int sum = value.x + value.y + value.z + value.w; + #pragma unroll + for (int j = 1; j < Batch; ++j) { + if (i + j*BlockSize < n) { + value = d_in[i + j*BlockSize]; + sum += value.x + value.y + value.z + value.w; + } + } + // Sum within warps + sum = warpReduce2(sum); + + // Store warp sums in sdata + if (tid % WARP_SIZE == 0) { + sdata[tid / WARP_SIZE] = sum; + } + __syncthreads(); + if (tid < WARP_SIZE) { + sum = sdata[tid]; + sum = warpReduce2(sum); + } + if (tid == 0) atomicAdd(d_out, sum); +} + +template +__global__ void sumKernel3(int4 *d_in, int *d_out, int n) { + __shared__ __align__(16) int sdata[1]; + + unsigned int tid = threadIdx.x; + unsigned int i = blockIdx.x*BlockSize*Batch + tid; + + if (i == 0) *d_out = 0; + if (tid == 0) sdata[0] = 0; + + int4 value = i < n ? d_in[i] : make_int4(0, 0, 0, 0); + int sum = value.x + value.y + value.z + value.w; + #pragma unroll + for (int j = 1; j < Batch; ++j) { + if (i + j*BlockSize < n) { + value = d_in[i + j*BlockSize]; + sum += value.x + value.y + value.z + value.w; + } + } + // Sum within warps + sum = warpReduce2(sum); + __syncthreads(); + if (tid % WARP_SIZE == 0) { + atomicAdd(&sdata[0], sum); + } + __syncthreads(); + if (tid == 0) { + atomicAdd(d_out, sdata[0]); + } +} + +inline int __device__ warpReduce3(int sum) { + int sum_warp = 0; + // https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-redux-sync + // Use `redux.sync` to reduce within warp + asm volatile("redux.sync.add.s32 %0, %1, 0xff;" : "=r"(sum_warp) : "r"(sum)); + return sum_warp; +} + +template +__global__ void sumKernel4(int4 *d_in, int *d_out, int n) { + __shared__ __align__(16) int sdata[1]; + + unsigned int tid = threadIdx.x; + unsigned int i = blockIdx.x*BlockSize*Batch + tid; + + if (i == 0) *d_out = 0; + if (tid == 0) sdata[0] = 0; + + int4 value = i < n ? __ldcg(&d_in[i]) : make_int4(0, 0, 0, 0); + int sum = value.x + value.y + value.z + value.w; + #pragma unroll + for (int j = 1; j < Batch; ++j) { + if (i + j*BlockSize < n) { + value = d_in[i + j*BlockSize]; + sum += value.x + value.y + value.z + value.w; + } + } + // Sum within warps + sum = warpReduce3(sum); + __syncthreads(); + if (tid % WARP_SIZE == 0) { + atomicAdd(&sdata[0], sum); + } + __syncthreads(); + if (tid == 0) { + atomicAdd(d_out, sdata[0]); + } +} + +const int NUM_KERNELS = 4; +const int NUM_KERNEL_RUNS = 100; +const int NUM_STD_RUNS = NUM_KERNEL_RUNS; + +void kernelDispatch(int kernelNum, int *d_in, int *d_out, int *h_out, int N) { + assert(N % 4 == 0); + switch (kernelNum) { + case 1: + sumKernel1<512, 20><<>>(d_in, d_out, N); + break; + case 2: + sumKernel2<1024, 16><<>>(reinterpret_cast(d_in), d_out, N/4); + break; + case 3: + sumKernel3<1024, 16><<>>(reinterpret_cast(d_in), d_out, N/4); + break; + case 4: + sumKernel4<1024, 16><<>>(reinterpret_cast(d_in), d_out, N/4); + break; + } +} + +void printDetails(std::string info, float timeMs, int N, int sum) { + std::cout << "<---------------| " << info << " |--------------->" << std::endl; + float bw = N*sizeof(int)/timeMs/1e6; + std::cout << "Bandwidth: " << bw << " GB/s" << std::endl; + std::cout << "Time taken: " << timeMs << " ms" << std::endl; + std::cout << "Sum: " << sum << std::endl << std::endl; +} + +void sumKernelCall(int kernelNum, int *d_in, int *d_out, int *h_out, int N, int times) { + cudaEvent_t start, stop; + cudaEventCreate(&start); + cudaEventCreate(&stop); + cudaEventRecord(start); + + for (int i = 0; i < times; ++i) kernelDispatch(kernelNum, d_in, d_out, h_out, N); + + cudaEventRecord(stop); + cudaEventSynchronize(stop); + + float elapsedTime; + cudaEventElapsedTime(&elapsedTime, start, stop); + + cudaMemcpy(h_out, d_out, sizeof(int), cudaMemcpyDeviceToHost); + cudaDeviceSynchronize(); + printDetails(std::string("kernel") + std::to_string(kernelNum), elapsedTime / times, N, *h_out); + + cudaEventDestroy(start); + cudaEventDestroy(stop); +} + +void sumCubCall(int *d_in, int *d_out, int *h_out, int N, int times) { + void* d_temp = nullptr; + size_t temp_storage = 0; + + // First call to determine temporary storage size + cub::DeviceReduce::Sum(d_temp, temp_storage, d_in, d_out, N); + + // Allocate temporary storage + assert(temp_storage > 0); + cudaMalloc(&d_temp, temp_storage); + + cudaEvent_t start, stop; + cudaEventCreate(&start); + cudaEventCreate(&stop); + + cudaEventRecord(start); + + for (int i = 0; i < times; ++i) cub::DeviceReduce::Sum(d_temp, temp_storage, d_in, d_out, N); + + cudaEventRecord(stop); + cudaEventSynchronize(stop); + + cudaMemcpy(h_out, d_out, sizeof(int), cudaMemcpyDeviceToHost); + + float elapsedTime; + cudaEventElapsedTime(&elapsedTime, start, stop); + + printDetails("cub", elapsedTime / times, N, *h_out); + + cudaFree(d_temp); + cudaEventDestroy(start); + cudaEventDestroy(stop); +} + +void sumCpuCall(int *h_in, int *h_out, int N) { + std::chrono::steady_clock::time_point begin = std::chrono::steady_clock::now(); + int value = 0; + for (int i = 0; i < N; ++i) { + value += h_in[i]; + } + *h_out = value; + float timeTaken = std::chrono::duration_cast(std::chrono::steady_clock::now() - begin).count()/1e6f; + printDetails("cpu", timeTaken, N, value); +} + +__global__ void warmupKernel() { + extern __shared__ int sdata[]; +} + +int main() { + warmupKernel<<<1024, 1024, 1024*sizeof(int)>>>(); + cudaDeviceSynchronize(); + + const int N = 1 << 30; + size_t size = N * sizeof(int); + + // Allocate host memory + int* h_in = new int[N]; + int h_out = 0.0f; + + srand(42); + for (int i = 0; i < N; ++i) { + h_in[i] = rand() % 100; + } + + // Allocate device memory + int* d_in; + int* d_out; + cudaMalloc(&d_in, size); + cudaMalloc(&d_out, sizeof(int)); + + // Copy input data from host to device + cudaMemcpy(d_in, h_in, size, cudaMemcpyHostToDevice); + cudaMemcpy(d_out, &h_out, sizeof(int), cudaMemcpyHostToDevice); + + for (int i = 1; i <= NUM_KERNELS; ++i) sumKernelCall(i, d_in, d_out, &h_out, N, NUM_KERNEL_RUNS); + sumCubCall(d_in, d_out, &h_out, N, NUM_STD_RUNS); + + // Free memory + cudaFree(d_in); + cudaFree(d_out); + delete[] h_in; + return 0; +} \ No newline at end of file diff --git a/code/labs/fast_cu/upstream_manifest.json b/code/labs/fast_cu/upstream_manifest.json new file mode 100644 index 000000000..a1d606b6c --- /dev/null +++ b/code/labs/fast_cu/upstream_manifest.json @@ -0,0 +1,123 @@ +{ + "schema_version": 1, + "repository": "https://github.com/pranjalssh/fast.cu", + "revision": "2dfe5e26aecfd9e5f27bf9d5837deea01acda24b", + "license": "MIT", + "license_file": "LICENSE", + "hardware": { + "h100": { + "gpu_family": "NVIDIA Hopper H100", + "compute_capability": [9, 0], + "architecture": "sm_90a", + "minimum_cuda": "12.0", + "minimum_cuda_basis": "Inferred from the compute_90a build, WGMMA, TMA, and CUtensorMap APIs; upstream does not state an exact minimum.", + "libraries": ["CUDA Runtime", "CUDA Driver", "cuBLAS", "CUB"] + }, + "nvfp4": { + "gpu_family": "NVIDIA Blackwell GB300/B300", + "compute_capability": [10, 3], + "architecture": "sm_103a", + "minimum_cuda": "13.1", + "minimum_cuda_basis": "Declared by the upstream README; the toolkit must support the required architecture-specific tcgen05 forms.", + "libraries": ["CUDA Runtime", "CUDA Driver", "cuBLASLt"] + } + }, + "families": { + "h100_bf16_matmul": { + "hardware": "h100", + "entrypoint": "h100/matmul.cu", + "build_target": "matmul", + "binary": "out/matmul", + "kernel_snapshots": [ + "h100/matmul/matmul_1.cuh", + "h100/matmul/matmul_2.cuh", + "h100/matmul/matmul_3.cuh", + "h100/matmul/matmul_4.cuh", + "h100/matmul/matmul_5.cuh", + "h100/matmul/matmul_6.cuh", + "h100/matmul/matmul_7.cuh", + "h100/matmul/matmul_8.cuh", + "h100/matmul/matmul_9.cuh", + "h100/matmul/matmul_10.cuh", + "h100/matmul/matmul_11.cuh", + "h100/matmul/matmul_12.cuh" + ], + "fixed_shape": [8192, 8192, 8192] + }, + "h100_sum_reduction": { + "hardware": "h100", + "entrypoint": "h100/sum.cu", + "build_target": "sum", + "binary": "out/sum", + "fixed_elements": 1073741824 + }, + "gb300_nvfp4_matmul": { + "hardware": "nvfp4", + "entrypoint": "gb300/nvfp4/main.cu", + "build_targets": [ + "nvfp4", + "nvfp4-r0", + "nvfp4-r1", + "nvfp4-r2", + "nvfp4-r3", + "nvfp4-r4", + "nvfp4-r5", + "nvfp4-r6", + "nvfp4-r7", + "nvfp4-r8", + "nvfp4-r9", + "nvfp4-ladder" + ], + "default_binary": "out/nvfp4", + "kernel_snapshots": [ + "gb300/nvfp4/gemm0.cuh", + "gb300/nvfp4/gemm1.cuh", + "gb300/nvfp4/gemm2.cuh", + "gb300/nvfp4/gemm3.cuh", + "gb300/nvfp4/gemm4.cuh", + "gb300/nvfp4/gemm5.cuh", + "gb300/nvfp4/gemm6.cuh", + "gb300/nvfp4/gemm7.cuh", + "gb300/nvfp4/gemm8.cuh", + "gb300/nvfp4/gemm9.cuh" + ], + "upstream_runner": "gb300/nvfp4/run_bench_modes.sh" + } + }, + "files": { + ".gitignore": "a888b9df89402a815bdb15d30c5280b090be11655ab4cc2299ce748d44cd6ed4", + "LICENSE": "076fb1df0d5125f25033dc51498629925fac3b705e86f8bf9808f5eec9d41785", + "Makefile": "54d8599fd3b11cef68fc15b0c848d3d058d3ec83bf0d1cbd40d72983a1c10f0e", + "README.md": "63e4e285b216a913d071993d95f55f64740d23abaa430efece99b9d05337946d", + "gb300/nvfp4/README.md": "c195cd506d7d1cdf5fda122a0f637cbf00039a8b35097249b588f575087e0449", + "gb300/nvfp4/gemm0.cuh": "dd4475b68d5ac1c358073e5dbc8f2b6f3c8f5b8e000bfbb9c643441c4576c131", + "gb300/nvfp4/gemm1.cuh": "02d22dc7ef843626bebe67e5a99c728cf2c9a5fec643bb5c4f5f0f183db6cb57", + "gb300/nvfp4/gemm2.cuh": "2a4bd066b522ad63056e30452c67846a98db5175bd9b126b80d1f1614ec4394a", + "gb300/nvfp4/gemm3.cuh": "1dca64c9d44ac00c2d00f770f7859f08b52ddcf2923864479087c5f0226c1a9d", + "gb300/nvfp4/gemm4.cuh": "4f6c92015b4ef4ec71f89e974725d86e2f40a44d3ef32b0ea9c9dbd542552389", + "gb300/nvfp4/gemm5.cuh": "f3de02a730648a96cd8e694e6d650d180558706dfd37a96ef860bebf8708f6f9", + "gb300/nvfp4/gemm6.cuh": "2eb4bdb07adac774bffcd55d7a130576f17ac3af2f1a06c8b97e01c7f8562c72", + "gb300/nvfp4/gemm7.cuh": "73ce3cb9687468ba107525200c4a3bc6f8614041a09ab0ba5ba4ab6a8b976cf8", + "gb300/nvfp4/gemm8.cuh": "28ead4f7c56e95026331a28af9b70e71ba0f5b6ab0af45d112b0f68dd723f1d2", + "gb300/nvfp4/gemm9.cuh": "b9ec1b3f07ed934728499421b7a90bb7f31cbc179960fd2707ecefbf2ec8606c", + "gb300/nvfp4/main.cu": "175bdd8b09265f5f012ba3594159ba669364cd20da949b3d40a8135fabcc4d92", + "gb300/nvfp4/run_bench_modes.sh": "474fc1c2f6662539c8d7e2209fe2010b27d4fdd9736e5a8f2bad37696ae537d9", + "h100/logs.txt": "f99566db425eb41e2dad76c926195a9a77963b36c1ecb21671881167ab180fc6", + "h100/matmul.cu": "d6029098a62ea95093e8415464dfe0714a371703f558a5ea2087b3c9f2e8aa36", + "h100/matmul/cublaslt_matmul.cu": "34eec28f850b3512ec64125ce3adf44bd847834be2a0d6aae44340f3f14178e2", + "h100/matmul/matmul_1.cuh": "d5987339865a7a7c982c3ea333c86ecea90929513f33533fe9efe1238e77c9b6", + "h100/matmul/matmul_10.cuh": "ce6965431164d6c7f5f2be20932f75e4547db5ceaa764955da6a392edb0c673f", + "h100/matmul/matmul_11.cuh": "3c1699e96f4664c55cbc74520bdee3f3fe23e1348387b2a910118fbfcedc5872", + "h100/matmul/matmul_12.cuh": "f6c40cf385f7bd4a566c989cb98e04df30680ed6bea434461fce799e987ced4b", + "h100/matmul/matmul_2.cuh": "09dafdef86a55d9b59ef6b61f9794f92bc1d7829df39e39e1709fe98b48a10dd", + "h100/matmul/matmul_3.cuh": "920c06c9aa74228c544b10a508ff72e1e01a51c438384ab7a9cfb3055f48671d", + "h100/matmul/matmul_4.cuh": "1cf4ed06360d8dfbbd2ef62c6918459dad46af9fe86ea7f29a5779ed2cea4137", + "h100/matmul/matmul_5.cuh": "606b3b849e02530884c7fc9d8c784fa846cd5b4505ef0bf39039de4925494933", + "h100/matmul/matmul_6.cuh": "0c916956b957ea6fcbf38f282ce99c29478874d8deae970db7c6ed9cad69dcb3", + "h100/matmul/matmul_7.cuh": "d8b4d046427a1c2462cf074fccd2171c9cc81c0a3b92ddac6b90bf80169f4c82", + "h100/matmul/matmul_8.cuh": "d4353ad5b4608d43e9565d4b71dafc75e18e076d632a81e98e54859f73478f01", + "h100/matmul/matmul_9.cuh": "3185afad674ed624ef97b8994669884ceac176cbf554523af9ef746153577069", + "h100/matmul/pingpong_experimental.cuh": "9b7e44221b7eb37158e60587b6767c8d3baefceaf649325cca2b6771b86f161b", + "h100/sum.cu": "7aaed0a7952f925ba6f72f6b6026851ea478eede64e19ffe5cea27865122f207" + } +} diff --git a/code/labs/fast_cu/validation.json b/code/labs/fast_cu/validation.json new file mode 100644 index 000000000..cc69f97c8 --- /dev/null +++ b/code/labs/fast_cu/validation.json @@ -0,0 +1,79 @@ +{ + "schema_version": 1, + "date": "2026-09-16", + "qualification": "noncanonical_virtualized_host_exploration", + "source_base_commit": "0299307fa", + "upstream_revision": "2dfe5e26aecfd9e5f27bf9d5837deea01acda24b", + "hardware": { + "gpu": "NVIDIA B200", + "compute_capability": "10.0", + "visible_gpus": 1, + "virtualized": true, + "cuda_toolkit": "13.0.88", + "pytorch": "2.9.1+cu130", + "sm_clock_mhz": 1965, + "memory_clock_mhz": 3996 + }, + "measurement": { + "profile": "deep_dive", + "validity_profile": "strict", + "requested_iterations": 20, + "requested_warmup": 5, + "protocol": "repository harness fixed resident buffers; not source blog rotating paired-round protocol" + }, + "native_sources_sha256": { + "h100_reduction_kernels.cu": "a8d786df751482c5a431b22b279548adbed09600621a167ca578d701e44aed48", + "h100_gemm_extension.cu": "e0442d4dfc4910ba981034d1c2544dcdaebabb464d62c115e2f755f4c272174c", + "nvfp4_native.cu": "ac4e8ea997c85ba4d4c01244208a1163d5febda0396d8bd3068e84f6d9c890f9", + "b200_epilogue.cu": "806f6226b44b6f1cd99a8e944650ae58571ee4647fb3b1b1d5163e2bb283d920" + }, + "results": [ + { + "target": "labs/fast_cu:int32_reduction", + "status": "failed_no_speedup", + "baseline_ms": 0.1883358495159113, + "candidate_ms": 0.1801341006920921, + "speedup": 1.0455313502124657, + "minimum_required_speedup": null, + "input_verification": true, + "output_verification": true, + "baseline_profilers": { + "nsys": "succeeded", + "ncu": "succeeded", + "torch": "succeeded" + }, + "candidate_profilers": { + "nsys": "succeeded", + "ncu": "succeeded", + "torch": "succeeded" + }, + "source_result_sha256": "b8cec133bd9b886311f8784d938783b46f91a0a4726c90c61fa4129a38f501f2" + }, + { + "target": "labs/fast_cu:b200_epilogue", + "status": "succeeded", + "baseline_ms": 0.05168658194423338, + "candidate_ms": 0.048537385734480486, + "speedup": 1.0648818670824238, + "minimum_required_speedup": null, + "input_verification": true, + "output_verification": true, + "baseline_profilers": { + "nsys": "succeeded", + "ncu": "succeeded", + "torch": "succeeded" + }, + "candidate_profilers": { + "nsys": "succeeded", + "ncu": "succeeded", + "torch": "succeeded" + }, + "source_result_sha256": "6a63865270a5fdfbb5a652e028bce169a68f565855caf427b2fb181fe01be4ec" + } + ], + "limitations": [ + "No H100 runtime validation", + "No SM103 runtime validation", + "No published speedup or expectation baseline from this single exploratory run" + ] +} diff --git a/code/labs/fast_cu/validation.md b/code/labs/fast_cu/validation.md new file mode 100644 index 000000000..53e2c9a3a --- /dev/null +++ b/code/labs/fast_cu/validation.md @@ -0,0 +1,94 @@ +# Integration validation - 2026-09-16 + +The [machine-readable receipt](validation.json) records the B200 experiment, +native CUDA source hashes, and digests of the original result reports. The +vendored source remains byte-for-byte identical to the revision and 34 file +hashes in [upstream_manifest.json](upstream_manifest.json). + +## B200 results + +Both comparisons ran on one visible NVIDIA B200 (SM100), with PyTorch +2.9.1+cu130 and CUDA toolkit 13.0.88. The repository harness used strict validity, +20 requested iterations, 5 warmups, and the `deep_dive` profile. It locked SM +and memory clocks to 1965 and 3996 MHz and recorded application-clock telemetry. +Nsight Systems, Nsight Compute, and PyTorch profiling succeeded for both arms +of both comparisons. Input equivalence and full-output verification passed. + +| Target | Baseline | Candidate | Ratio | Harness outcome | +| --- | ---: | ---: | ---: | --- | +| `int32_reduction` | 0.188336 ms | 0.180134 ms | 1.0455x | Below the required 1.05x speed threshold | +| `b200_epilogue` | 0.051687 ms | 0.048537 ms | 1.0649x | Passed the speed threshold in this run | + +These are **exploratory, non-canonical measurements on a virtualized host**. +They are one fixed-resident-buffer harness run, not the blog's rotating-input, +alternating-round protocol. No expectation file or published speedup baseline +was created. Repeat interleaved measurements on the intended deployment hardware +before treating either ratio as a durable improvement. + +The reduction is not marked as a passing optimization merely because it was +slightly faster. The 256-bit epilogue result applies to that isolated conversion +and store workload, not to a full B200 NVFP4 GEMM. + +## Correctness and integration checks + +The final focused suite passed **66 tests with 2 expected hardware skips** on +B200, and **63 tests with 5 expected hardware skips** on the local CPU host. + +- Real B200 tests compare CUB and the safe reduction against an independent + integer sum, on a nondefault CUDA stream. +- Real B200 epilogue tests compare every output element with PyTorch's FP16 + conversion, cover nondefault-stream ordering, and verify caller-owned RNG and + setup/verification lifecycle behavior. +- CPU tests cover source integrity, changed/missing/extra source rejection, + build-cache identity, explicit hardware/toolkit contracts, target discovery, + output lifecycle, and selectable NVFP4 rungs. +- All four targets are discoverable through `aisp bench list-targets`. +- Repository-wide benchmark lint checked 944 entrypoints with zero errors and + zero warnings. Focused Ruff, formatting, syntax, YAML, and local-link checks + also passed. +- The H100 adapter compiled and linked for `sm_90a` with CUDA 13.0. No H100 + kernel was executed; numerical and performance qualification still needs H100. +- The SM103 adapter retains its CUDA 13.1+ requirement. Its independent host + oracle, poisoned-output, guard, determinism, and placement checks are part of + the adapter, but SM103 execution remains unqualified without B300/GB300. + +## Failures caught and corrected + +1. The original reduction resets a shared output from inside one CTA while other + CTAs can already update it. The adaptation resets before the grid on the + current stream and uses bounded int32 inputs. +2. The first B200 epilogue conversion swapped adjacent FP16 lanes. The full-output + GPU oracle rejected it; the PTX operands were corrected and the test passed. +3. Publishing an uninitialized epilogue output during setup triggered the + harness's precomputation check. Setup now retains only a private allocation; + the public output is assigned by the timed invocation. +4. Verification payloads retained previous GPU inputs across teardown. Setup and + teardown now release those payloads in every adapter. +5. PyTorch's extension flags disabled the BF16 conversion used by the original + Hopper header. The H100 adapter explicitly restores that conversion when + compiling, without modifying the pinned source. +6. The host's system Python reported conflicting installed-library versions. + The harness correctly rejected that comparison; reported results use an + isolated environment with unambiguous package metadata. + +## Reproduce + +From `code/`, run the tests on the allocated GPU and then the two B200 targets +serially using the [README commands](README.md#run-with-the-repository-harness): + +```bash +python -m pytest tests/test_fast_cu_*.py -q -ra +python -m core.scripts.linting.check_benchmarks labs/fast_cu --fail-on-warnings +``` + +For actual SM103 validation, explicitly enable its GPU test on B300/GB300 with +CUDA 13.1+: + +```bash +AISP_RUN_FAST_CU_NVFP4_GPU_TEST=1 python -m pytest tests/test_fast_cu_nvfp4.py -q -ra +``` + +The native default is r9; construct `FastCuNvfp4KernelBenchmark(rung=N)` to +select another snapshot. Every rung must pass its own target-host checks before +its results can be used. Retained upstream logs and article figures are source +evidence, not measurements from this integration. diff --git a/code/labs/fast_cu/workload_spec.yaml b/code/labs/fast_cu/workload_spec.yaml new file mode 100644 index 000000000..1b0446604 --- /dev/null +++ b/code/labs/fast_cu/workload_spec.yaml @@ -0,0 +1,47 @@ +run_metadata: + benchmark_layer: component + benchmark_grade: realism + use_case: engineering_debug +comparison_axis: + variable_under_test: kernel_implementation_or_store_width + constants_locked: [shape, input_bits, dtype, output_layout, caller_seed, selected_device, stream, cache_protocol] +workload: + h100_bf16_gemm: + control: cublas + candidate: upstream_matmul_12 + shape: [8192, 8192, 8192] + input_dtype: bf16 + accumulation: fp32 + output_dtype: bf16 + int32_reduction: + control: cub + candidate: safe_vectorized_warp_reduction + elements: 268435456 + input_values: [-1, 0, 1] + dtype: int32 + b200_epilogue: + control: two_128bit_stores + candidate: one_256bit_store + elements: 16777216 + input_dtype: fp32 + output_dtype: fp16 + shared_cache_policy: L1_no_allocate_L2_evict_first + nvfp4_gemm: + control: cublaslt + candidate: upstream_r9 + shape: [8192, 8192, 8192] + packed_dtype: e2m1 + scale_dtype: ue4m3 + scale_vector_size: 16 + output_dtype: fp16 +measurement: + timing_scope: steady_state_benchmark_invocation + excluded_setup: [jit_compile, allocations, random_inputs, descriptors, workspace_selection, schedule_preparation] + verification: full_actual_output_after_execution + input_cache_protocol: fixed_resident_buffers + canonical_profile: strict + required_profilers: [ncu, nsys] + expected_speedup: null +source: + manifest: upstream_manifest.json + revision: 2dfe5e26aecfd9e5f27bf9d5837deea01acda24b diff --git a/code/tests/test_fast_cu_b200_epilogue.py b/code/tests/test_fast_cu_b200_epilogue.py new file mode 100644 index 000000000..47698b02f --- /dev/null +++ b/code/tests/test_fast_cu_b200_epilogue.py @@ -0,0 +1,266 @@ +"""CPU contracts plus optional exact-SM100 execution for the B200 epilogue pair.""" + +from __future__ import annotations + +import importlib +import re +from pathlib import Path + +import pytest +import torch + +from core.discovery import discover_benchmarks +from core.harness.validity_checks import check_setup_precomputation +from labs.fast_cu import b200_epilogue +from labs.fast_cu.b200_epilogue import ( + DEFAULT_NUM_ELEMENTS, + ELEMENTS_PER_THREAD, + EXPECTED_COMPUTE_CAPABILITY, + MINIMUM_CUDA, + B200EpilogueBenchmark, + ensure_b200_capability_supported, + ensure_cuda_build_supported, + parse_cuda_build, + require_b200_runtime, + validate_epilogue_tensors, + validate_num_elements, +) + +LAB_DIR = Path(__file__).resolve().parents[1] / "labs" / "fast_cu" +CUDA_SOURCE = LAB_DIR / "b200_epilogue.cu" + + +def test_default_workload_is_representative_and_aligned() -> None: + assert DEFAULT_NUM_ELEMENTS == 1 << 24 + assert ELEMENTS_PER_THREAD == 16 + assert DEFAULT_NUM_ELEMENTS % ELEMENTS_PER_THREAD == 0 + assert MINIMUM_CUDA == (12, 9) + assert EXPECTED_COMPUTE_CAPABILITY == (10, 0) + + +@pytest.mark.parametrize("value", [True, 0, -16, 17, 31, 1.5, "16"]) +def test_num_elements_contract_rejects_invalid_values(value) -> None: + with pytest.raises(ValueError): + validate_num_elements(value) + + +def test_cuda_build_parser_is_explicit() -> None: + assert parse_cuda_build("12.9") == (12, 9) + assert parse_cuda_build("13.0") == (13, 0) + assert parse_cuda_build("13.0.1") == (13, 0) + with pytest.raises(RuntimeError, match="CUDA-enabled"): + parse_cuda_build(None) + with pytest.raises(RuntimeError, match="Cannot parse"): + parse_cuda_build("development") + + +def test_benchmark_construction_does_not_resolve_cuda() -> None: + baseline = B200EpilogueBenchmark(optimized=False, num_elements=32) + optimized = B200EpilogueBenchmark(optimized=True, num_elements=32) + + assert baseline._device is None + assert optimized._device is None + assert baseline.get_config() == optimized.get_config() + assert baseline.get_workload_metadata() == optimized.get_workload_metadata() + assert baseline.get_custom_metrics()["fast_cu.store_instructions_per_thread"] == 2.0 + assert optimized.get_custom_metrics()["fast_cu.store_instructions_per_thread"] == 1.0 + + +def test_real_discovery_finds_the_pair() -> None: + pairs = { + example: (baseline.name, [path.name for path in optimized]) + for baseline, optimized, example in discover_benchmarks(LAB_DIR, warn_missing=False) + if example == "b200_epilogue" + } + assert pairs == { + "b200_epilogue": ( + "baseline_b200_epilogue.py", + ["optimized_b200_epilogue.py"], + ) + } + + +@pytest.mark.parametrize( + ("module_name", "optimized"), + [ + ("baseline_b200_epilogue", False), + ("optimized_b200_epilogue", True), + ], +) +def test_wrapper_factories_construct_without_cuda(module_name: str, optimized: bool) -> None: + module = importlib.import_module(f"labs.fast_cu.{module_name}") + benchmark = module.get_benchmark() + assert isinstance(benchmark, B200EpilogueBenchmark) + assert benchmark.optimized is optimized + assert benchmark._device is None + + +def test_tensor_contract_rejects_dtype_layout_shape_alignment_and_cpu() -> None: + accumulators = torch.empty(16, dtype=torch.float32) + output = torch.empty(16, dtype=torch.float16) + + with pytest.raises(TypeError, match="float32"): + validate_epilogue_tensors(accumulators.to(torch.float64), output) + with pytest.raises(TypeError, match="float16"): + validate_epilogue_tensors(accumulators, output.to(torch.float32)) + noncontiguous = torch.empty((16, 2), dtype=torch.float32)[:, 0] + with pytest.raises(ValueError, match="contiguous"): + validate_epilogue_tensors(noncontiguous, output) + with pytest.raises(ValueError, match="matching shapes"): + validate_epilogue_tensors(accumulators, torch.empty(32, dtype=torch.float16)) + + output_storage = torch.empty(17, dtype=torch.float16) + misaligned_output = output_storage[1:17] + assert misaligned_output.is_contiguous() + assert misaligned_output.data_ptr() % 32 != 0 + with pytest.raises(ValueError, match="32-byte aligned"): + validate_epilogue_tensors(accumulators, misaligned_output) + + with pytest.raises(ValueError, match="CUDA tensors"): + validate_epilogue_tensors(accumulators, output) + + +def test_pure_runtime_contract_requires_cuda_12_9_and_exact_sm100() -> None: + ensure_b200_capability_supported((10, 0)) + ensure_cuda_build_supported((12, 9)) + ensure_cuda_build_supported((13, 0)) + with pytest.raises(RuntimeError, match="exact SM100"): + ensure_b200_capability_supported((10, 3)) + with pytest.raises(RuntimeError, match=r"12\.9\+"): + ensure_cuda_build_supported((12, 8)) + + +def test_runtime_gate_reports_real_host_without_mocking_success() -> None: + device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cuda") + ) + supported = ( + torch.cuda.is_available() + and torch.cuda.get_device_capability(device) == EXPECTED_COMPUTE_CAPABILITY + and parse_cuda_build(torch.version.cuda) >= MINIMUM_CUDA + ) + if supported: + require_b200_runtime(device) + else: + with pytest.raises(RuntimeError, match=r"SKIPPED:"): + require_b200_runtime(device) + + +def test_extension_build_is_pinned_to_sm100_and_cuda_12_9(monkeypatch) -> None: + captured = {} + sentinel = object() + + def fake_load(name, sources, **kwargs): + captured.update(name=name, sources=sources, **kwargs) + return sentinel + + b200_epilogue.load_b200_epilogue_extension.cache_clear() + monkeypatch.setattr(b200_epilogue, "load_cuda_extension", fake_load) + try: + assert b200_epilogue.load_b200_epilogue_extension() is sentinel + finally: + b200_epilogue.load_b200_epilogue_extension.cache_clear() + + assert captured["name"] == "fast_cu_b200_epilogue" + assert captured["sources"] == [CUDA_SOURCE] + assert captured["minimum_cuda"] == (12, 9) + assert "-gencode=arch=compute_100,code=sm_100" in captured["extra_cuda_cflags"] + + +def test_cuda_source_keeps_conversion_policy_and_store_width_isolated() -> None: + source = CUDA_SOURCE.read_text(encoding="utf-8") + + assert source.count("cvt.rn.f16x2.f32") == 1 + assert ': "f"(hi), "f"(lo)' in source + assert source.count("createpolicy.fractional.L2::evict_first") == 1 + assert source.count("st.global.L1::no_allocate.L2::cache_hint.v4.b32") == 1 + assert source.count("st.global.L1::no_allocate.L2::cache_hint.v8.b32") == 1 + assert source.count("pack_f16x2(accumulators[") == 8 + + branch = re.search( + r"if constexpr \(WideStore\) \{(?P.*?)\} else \{(?P.*?)\n \}", + source, + re.DOTALL, + ) + assert branch is not None + assert branch.group("wide").count("store_b256(") == 1 + assert "store_b128(" not in branch.group("wide") + assert branch.group("narrow").count("store_b128(") == 2 + assert "store_b256(" not in branch.group("narrow") + + +def _require_actual_sm100() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + if torch.cuda.get_device_capability() != (10, 0): + pytest.skip("requires exact SM100") + if parse_cuda_build(torch.version.cuda) < MINIMUM_CUDA: + pytest.skip("requires a CUDA 12.9+ PyTorch build") + + +def test_actual_sm100_kernels_match_full_reference_on_current_stream() -> None: + _require_actual_sm100() + extension = b200_epilogue.load_b200_epilogue_extension() + extension.require_exact_sm100() + + count = 1 << 20 + accumulators = torch.randn(count, device="cuda", dtype=torch.float32) + expected = accumulators.to(torch.float16) + baseline = torch.empty_like(expected) + optimized = torch.empty_like(expected) + validate_epilogue_tensors(accumulators, baseline) + validate_epilogue_tensors(accumulators, optimized) + + current = torch.cuda.current_stream() + stream = torch.cuda.Stream() + stream.wait_stream(current) + with torch.cuda.stream(stream): + extension.baseline(accumulators, baseline) + extension.optimized(accumulators, optimized) + stream.synchronize() + + torch.testing.assert_close(baseline, expected, rtol=0, atol=0) + torch.testing.assert_close(optimized, expected, rtol=0, atol=0) + torch.testing.assert_close(optimized, baseline, rtol=0, atol=0) + + +def test_actual_sm100_setup_uses_caller_seed_and_resets_verification_state() -> None: + _require_actual_sm100() + b200_epilogue.load_b200_epilogue_extension() + device_index = torch.cuda.current_device() + device = torch.device("cuda", device_index) + + with torch.random.fork_rng(devices=[device_index]): + snapshots = [] + for seed in (42, 42, 1042): + torch.manual_seed(seed) + benchmark = B200EpilogueBenchmark(optimized=False, num_elements=32) + benchmark.device = device + benchmark._ran = True + benchmark._verification_output = torch.ones(1, device=device) + benchmark._verification_payload = object() + setup_valid, setup_error = check_setup_precomputation( + lambda benchmark=benchmark: {"output": benchmark.output}, + benchmark.setup, + ) + assert setup_valid, setup_error + snapshots.append( + ( + benchmark.accumulators.detach().clone(), + benchmark._caller_seed, + int(torch.initial_seed()), + ) + ) + assert benchmark._ran is False + assert benchmark.output is None + assert benchmark._output_buffer is not None + assert benchmark._verification_output is None + assert benchmark._verification_payload is None + benchmark.teardown() + + assert torch.equal(snapshots[0][0], snapshots[1][0]) + assert not torch.equal(snapshots[0][0], snapshots[2][0]) + assert snapshots[0][1:] == snapshots[1][1:] == (42, 42) + assert snapshots[2][1:] == (1042, 1042) diff --git a/code/tests/test_fast_cu_build.py b/code/tests/test_fast_cu_build.py new file mode 100644 index 000000000..92b7aeb8b --- /dev/null +++ b/code/tests/test_fast_cu_build.py @@ -0,0 +1,59 @@ +"""Build provenance checks that execute without a CUDA toolchain.""" + +from __future__ import annotations + +import pytest + +from labs.fast_cu.build import extension_name, parse_nvcc_release + + +def test_nvcc_release_requires_explicit_toolkit_version(): + assert parse_nvcc_release("Cuda compilation tools, release 13.1, V13.1.80") == (13, 1) + assert parse_nvcc_release("Cuda compilation tools, release 12.8, V12.8.61") == (12, 8) + with pytest.raises(RuntimeError, match="Cannot determine"): + parse_nvcc_release("driver version 590.0") + + +def test_cache_key_tracks_source_headers_and_flags(tmp_path): + source = tmp_path / "extension.cu" + source.write_text("// original\n") + manifest = {"revision": "pinned", "files": {"gemm.cuh": "a" * 64}} + original = extension_name("fast_cu", [source], ["-O3"], ["-lcuda"], manifest) + assert original == extension_name("fast_cu", [source], ["-O3"], ["-lcuda"], manifest) + source.write_text("// changed\n") + assert original != extension_name("fast_cu", [source], ["-O3"], ["-lcuda"], manifest) + source.write_text("// original\n") + assert original != extension_name("fast_cu", [source], ["-O2"], ["-lcuda"], manifest) + assert original != extension_name("fast_cu", [source], ["-O3"], ["-lcublas"], manifest) + manifest["files"]["gemm.cuh"] = "b" * 64 + assert original != extension_name("fast_cu", [source], ["-O3"], ["-lcuda"], manifest) + + +def test_extension_name_is_safe_and_sources_required(tmp_path): + with pytest.raises(ValueError, match="identifier"): + extension_name("bad-name", [], [], [], {}) + with pytest.raises(ValueError, match="source"): + extension_name("fast_cu", [], [], [], {}) + + +def test_cache_key_tracks_in_place_toolkit_upgrade(tmp_path): + source = tmp_path / "extension.cu" + source.write_text("// unchanged source\n") + manifest = {"revision": "pinned", "files": {"gemm.cuh": "a" * 64}} + old = extension_name( + "fast_cu", + [source], + [], + [], + manifest, + toolchain_identity="/cuda/bin/nvcc\nrelease 13.0 V13.0.88", + ) + new = extension_name( + "fast_cu", + [source], + [], + [], + manifest, + toolchain_identity="/cuda/bin/nvcc\nrelease 13.1 V13.1.80", + ) + assert old != new diff --git a/code/tests/test_fast_cu_h100.py b/code/tests/test_fast_cu_h100.py new file mode 100644 index 000000000..4d9b2a1c1 --- /dev/null +++ b/code/tests/test_fast_cu_h100.py @@ -0,0 +1,211 @@ +"""Contracts and real-GPU coverage for the fast.cu Hopper adapters.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +from labs.fast_cu.baseline_h100_bf16_gemm import ( + BaselineH100Bf16GemmBenchmark, +) +from labs.fast_cu.baseline_int32_reduction import ( + BaselineInt32ReductionBenchmark, +) +from labs.fast_cu.h100_common import ( + GEMM_K, + GEMM_M, + GEMM_N, + H100_CAPABILITY, + REDUCTION_CAPABILITIES, + REDUCTION_ELEMENTS, + _gemm_cuda_flags, + _reduction_build_contract, + ensure_h100_capability_supported, + ensure_reduction_capability_supported, + load_h100_gemm_extension, + load_int32_reduction_extension, +) +from labs.fast_cu.optimized_h100_bf16_gemm import ( + OptimizedH100Bf16GemmBenchmark, +) +from labs.fast_cu.optimized_int32_reduction import ( + OptimizedInt32ReductionBenchmark, +) + +LAB_DIR = Path(__file__).resolve().parents[1] / "labs" / "fast_cu" + + +def _require_cuda_capability(supported: set[tuple[int, int]] | frozenset[tuple[int, int]]): + if not torch.cuda.is_available(): + pytest.skip("real CUDA GPU is unavailable") + capability = tuple(torch.cuda.get_device_capability()) + if capability not in supported: + pytest.skip(f"real GPU has unsupported capability {capability}") + return capability + + +def test_h100_gate_is_exact_and_reports_skip(): + ensure_h100_capability_supported((9, 0)) + for capability in ((8, 0), (8, 9), (9, 1), (10, 0), (10, 3), (12, 0)): + with pytest.raises(RuntimeError, match=r"SKIPPED:.*exact SM 9\.0"): + ensure_h100_capability_supported(capability) + + +def test_reduction_gate_and_build_contract_cover_h100_and_b200_only(): + assert {(9, 0), (10, 0)} == REDUCTION_CAPABILITIES + for capability, arch, minimum in ( + ((9, 0), "90a", (12, 0)), + ((10, 0), "100a", (12, 8)), + ): + ensure_reduction_capability_supported(capability) + flags, minimum_cuda = _reduction_build_contract(capability) + assert f"-gencode=arch=compute_{arch},code=sm_{arch}" in flags + assert minimum_cuda == minimum + for capability in ((8, 0), (9, 1), (10, 3), (12, 0)): + with pytest.raises(RuntimeError, match=r"SKIPPED:.*SM 9\.0.*SM 10\.0"): + ensure_reduction_capability_supported(capability) + + +def test_h100_gemm_build_targets_arch_specific_wgmma(): + flags = _gemm_cuda_flags(H100_CAPABILITY) + assert "-gencode=arch=compute_90a,code=sm_90a" in flags + assert "-DNDEBUG" in flags + assert "-U__CUDA_NO_BFLOAT16_CONVERSIONS__" in flags + assert not any("compute_100" in flag for flag in flags) + reduction_flags, _ = _reduction_build_contract(H100_CAPABILITY) + assert "-U__CUDA_NO_BFLOAT16_CONVERSIONS__" not in reduction_flags + + +def test_benchmark_contracts_are_equivalent_and_full_output(): + baseline_gemm = BaselineH100Bf16GemmBenchmark() + optimized_gemm = OptimizedH100Bf16GemmBenchmark() + assert baseline_gemm.get_input_signature() == optimized_gemm.get_input_signature() + gemm_signature = baseline_gemm.get_input_signature() + assert gemm_signature.shapes == { + "matrix_a": (GEMM_M, GEMM_K), + "matrix_b": (GEMM_N, GEMM_K), + "output": (GEMM_M, GEMM_N), + } + assert gemm_signature.dtypes["output"] == "torch.bfloat16" + assert baseline_gemm.get_config().warmup == 5 + assert optimized_gemm.get_config().iterations == 20 + + baseline_reduction = BaselineInt32ReductionBenchmark() + optimized_reduction = OptimizedInt32ReductionBenchmark() + assert baseline_reduction.get_input_signature() == optimized_reduction.get_input_signature() + reduction_signature = baseline_reduction.get_input_signature() + assert reduction_signature.shapes == { + "input": (REDUCTION_ELEMENTS,), + "output": (1,), + } + assert reduction_signature.dtypes == { + "input": "torch.int32", + "output": "torch.int32", + } + assert baseline_reduction.get_config().warmup == 10 + assert optimized_reduction.get_config().iterations == 50 + + +def test_teardown_releases_full_verification_payloads_on_cpu(): + gemm = BaselineH100Bf16GemmBenchmark() + gemm.matrix_a = torch.ones((2, 3), dtype=torch.bfloat16) + gemm.matrix_b = torch.ones((4, 3), dtype=torch.bfloat16) + gemm._physical_output = torch.ones((4, 2), dtype=torch.bfloat16) + gemm._logical_output = gemm._physical_output.transpose(0, 1) + gemm.output = gemm._logical_output + gemm.capture_verification_payload() + assert gemm._verification_payload is not None + + gemm.teardown() + + assert gemm._verification_payload is None + assert gemm.matrix_a is None and gemm.matrix_b is None + assert gemm._physical_output is None and gemm._logical_output is None + with pytest.raises(RuntimeError, match="must be called before verification"): + gemm.get_verify_output() + + reduction = BaselineInt32ReductionBenchmark() + reduction.input = torch.tensor([-1, 0, 1], dtype=torch.int32) + reduction._output_buffer = torch.tensor([0], dtype=torch.int32) + reduction.output = reduction._output_buffer + reduction.capture_verification_payload() + assert reduction._verification_payload is not None + + reduction.teardown() + + assert reduction._verification_payload is None + assert reduction.input is None and reduction._output_buffer is None + with pytest.raises(RuntimeError, match="must be called before verification"): + reduction.get_verify_output() + + +def test_safe_reduction_source_orders_reset_before_launch_on_current_stream(): + source = (LAB_DIR / "h100_reduction_kernels.cu").read_text(encoding="utf-8") + adapter = source[source.index("void run_safe_vectorized_reduction") :] + reset = adapter.index("cudaMemsetAsync") + launch = adapter.index("safe_vectorized_sum_kernel<<<") + assert reset < launch + assert "at::cuda::getCurrentCUDAStream" in adapter[:launch] + kernel = source[ + source.index("__global__ __launch_bounds__") : source.index( + "int64_t cub_temp_storage_bytes" + ) + ] + assert "*output = 0" not in kernel + assert "int4* input" in kernel + assert "atomicAdd(output, sum)" in kernel + + +def test_gemm_source_uses_pinned_latest_kernel_and_current_stream(): + source = (LAB_DIR / "h100_gemm_extension.cu").read_text(encoding="utf-8") + assert '#include "upstream/h100/matmul/matmul_12.cuh"' in source + assert "M12::matmulKernel12" in source + launch = source[source.index("void run_upstream_gemm") :] + assert "at::cuda::getCurrentCUDAStream" in launch + assert "<<>>" in launch + assert "physical output must have transposed [N,M] layout" in source + + +def test_real_supported_gpu_int32_reduction_matches_full_reference(): + _require_cuda_capability(REDUCTION_CAPABILITIES) + extension = load_int32_reduction_extension() + count = 1 << 20 + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + values = (torch.arange(count, device="cuda", dtype=torch.int32) % 3) - 1 + cub_output = torch.empty((1,), device="cuda", dtype=torch.int32) + optimized_output = torch.empty_like(cub_output) + storage_bytes = extension.cub_temp_storage_bytes(values, cub_output) + temp_storage = torch.empty((storage_bytes,), device="cuda", dtype=torch.uint8) + extension.run_cub_reduction(values, cub_output, temp_storage) + extension.run_safe_vectorized_reduction(values, optimized_output) + reference = values.sum(dtype=torch.int64).to(torch.int32).reshape(1) + stream.synchronize() + assert torch.equal(cub_output, reference) + assert torch.equal(optimized_output, reference) + + +def test_real_h100_gemm_matches_cublas_full_logical_output(): + _require_cuda_capability({H100_CAPABILITY}) + extension = load_h100_gemm_extension() + # Two M-cluster tiles avoid the pinned host schedule's one-tile clz edge. + m, n, k = 512, 256, 64 + a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + b = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + cublas_physical = torch.empty((n, m), device="cuda", dtype=torch.bfloat16) + upstream_physical = torch.empty_like(cublas_physical) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + extension.prepare_cublas_gemm(a, b, cublas_physical) + extension.prepare_upstream_gemm(a, b, upstream_physical) + extension.run_cublas_gemm(a, b, cublas_physical) + extension.run_upstream_gemm(a, b, upstream_physical) + fp32_reference = a.float() @ b.float().transpose(0, 1) + stream.synchronize() + cublas_logical = cublas_physical.transpose(0, 1) + upstream_logical = upstream_physical.transpose(0, 1) + assert torch.allclose(upstream_logical, cublas_logical, rtol=1e-2, atol=1e-1) + assert torch.allclose(cublas_logical.float(), fp32_reference, rtol=1e-2, atol=1e-1) + assert torch.allclose(upstream_logical.float(), fp32_reference, rtol=1e-2, atol=1e-1) diff --git a/code/tests/test_fast_cu_nvfp4.py b/code/tests/test_fast_cu_nvfp4.py new file mode 100644 index 000000000..955118bda --- /dev/null +++ b/code/tests/test_fast_cu_nvfp4.py @@ -0,0 +1,235 @@ +"""Static, CPU, and opt-in GB300 coverage for the native fast.cu NVFP4 pair.""" + +from __future__ import annotations + +import os +from dataclasses import FrozenInstanceError +from pathlib import Path + +import pytest +import torch + +from core.harness.benchmark_harness import BaseBenchmark +from labs.fast_cu.baseline_nvfp4_gemm import FastCuNvfp4CublasLtBenchmark +from labs.fast_cu.nvfp4 import ( + COMPUTE_CAPABILITY, + CUDA_VERSION, + DEFAULT_K, + DEFAULT_M, + DEFAULT_N, + DEFAULT_RUNG, + Nvfp4Workload, + require_nvfp4_runtime, + validate_rung, + workload_bytes, + workload_flops, +) +from labs.fast_cu.optimized_nvfp4_gemm import FastCuNvfp4KernelBenchmark, get_benchmark +from tests.protection_test_utils import preserve_rng_state + +LAB_DIR = Path(__file__).resolve().parents[1] / "labs" / "fast_cu" +NATIVE_SOURCE = LAB_DIR / "nvfp4_native.cu" +PYTHON_SOURCE = LAB_DIR / "nvfp4.py" + + +def test_default_workload_and_rung_are_explicit_and_frozen() -> None: + workload = Nvfp4Workload() + + assert workload == Nvfp4Workload(DEFAULT_M, DEFAULT_N, DEFAULT_K) + assert (DEFAULT_M, DEFAULT_N, DEFAULT_K) == (8192, 8192, 8192) + assert DEFAULT_RUNG == 9 + assert CUDA_VERSION == (13, 1) + assert COMPUTE_CAPABILITY == (10, 3) + workload.validate() + assert workload_bytes(workload) == 209_715_200 + assert workload_flops(workload) == 1_099_511_627_776 + with pytest.raises(FrozenInstanceError): + workload.m = 4096 # type: ignore[misc] + + +@pytest.mark.parametrize( + ("workload", "message"), + [ + (Nvfp4Workload(0, 128, 128), "m must be"), + (Nvfp4Workload(128, -1, 128), "n must be"), + (Nvfp4Workload(128, 128, True), "k must be"), + (Nvfp4Workload(128, 128, 48), "divisible by 32"), + ], +) +def test_workload_rejects_invalid_or_cublaslt_unsupported_shapes( + workload: Nvfp4Workload, + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + workload.validate() + + +@pytest.mark.parametrize("rung", range(10)) +def test_every_upstream_optimization_rung_is_selectable(rung: int) -> None: + assert validate_rung(rung) == rung + benchmark = FastCuNvfp4KernelBenchmark(Nvfp4Workload(256, 256, 256), rung=rung) + assert benchmark.rung == rung + + +@pytest.mark.parametrize("rung", [-1, 10, True, 1.5, "9"]) +def test_invalid_rung_is_rejected(rung: object) -> None: + with pytest.raises(ValueError, match=r"integer in \[0, 9\]"): + validate_rung(rung) # type: ignore[arg-type] + + +def test_benchmark_factories_are_import_safe_and_harness_native() -> None: + baseline = FastCuNvfp4CublasLtBenchmark(Nvfp4Workload(256, 256, 256)) + optimized = get_benchmark() + + assert isinstance(baseline, BaseBenchmark) + assert isinstance(optimized, FastCuNvfp4KernelBenchmark) + assert baseline.optimized is False + assert optimized.optimized is True + assert optimized.rung == 9 + assert baseline.output is None and optimized.output is None + assert baseline._output_buffer is None and optimized._output_buffer is None + metadata = optimized.get_workload_metadata() + assert metadata is not None + assert metadata.bytes_per_iteration == float(workload_bytes(optimized.workload)) + assert metadata.custom_units_per_iteration == float(workload_flops(optimized.workload)) + assert metadata.custom_unit_name == "nvfp4_flops" + config = optimized.get_config() + assert config.iterations == 25 + assert config.warmup == 5 + assert config.single_gpu is True + + +def test_verification_capture_rejects_pre_execution_state() -> None: + benchmark = FastCuNvfp4CublasLtBenchmark(Nvfp4Workload(256, 256, 256)) + + with pytest.raises(RuntimeError, match=r"benchmark_fn\(\) must run"): + benchmark.capture_verification_payload() + + +def test_teardown_releases_the_full_verification_payload() -> None: + benchmark = FastCuNvfp4CublasLtBenchmark(Nvfp4Workload(256, 256, 256)) + benchmark._native_context = object() + benchmark._inputs = {"packed": torch.ones(4, dtype=torch.uint8)} + benchmark.output = torch.ones((2, 2), dtype=torch.float16) + benchmark._setup_gate_passed = True + benchmark.capture_verification_payload() + assert benchmark._verification_payload is not None + + benchmark.teardown() + + assert benchmark._verification_payload is None + assert benchmark._inputs == {} + assert benchmark.output is None + + +def test_native_source_preserves_the_equivalent_workload_contract() -> None: + source = NATIVE_SOURCE.read_text(encoding="utf-8") + python_source = PYTHON_SOURCE.read_text(encoding="utf-8") + + assert '#include "main.cu"' in source + assert "host::make_fixture(M, N, K, seed" in source + assert "fixture.inputs.A == fixture.A_logical" in source + assert "CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3" in source + assert "CUDA_R_4F_E2M1" in (LAB_DIR / "upstream" / "gb300" / "nvfp4" / "main.cu").read_text( + encoding="utf-8" + ) + assert "output=self.output" in python_source + assert "output.sum" not in python_source + assert "checksum" not in python_source.lower() + + +def test_native_source_uses_current_stream_and_rejects_r9_degradation() -> None: + source = NATIVE_SOURCE.read_text(encoding="utf-8") + + assert "at::cuda::getCurrentCUDAStream(device)" in source + assert "cublasLtMatmul(" in source and "stream)," in source + assert "<<>>" in source + assert "strict_setup_r9_schedule" in source + assert "sched::build_schedule" in source + assert "table.empty()" in source + assert "raster fallback is forbidden" in source + assert "l2a_placement_errors" in source + assert "bench::small_correctness_gates()" in source + assert "small host-reference, guard, or determinism gate failed" in source + assert "bench::setup_schedule" not in source + + +def test_build_contract_is_exact_sm103a_cuda131_and_setup_only() -> None: + source = PYTHON_SOURCE.read_text(encoding="utf-8") + + assert "verify_upstream()" in source + assert "require_nvfp4_runtime()" in source + assert "require_nvfp4_toolkit()" in source + assert '"-gencode=arch=compute_103a,code=sm_103a"' in source + assert "minimum_cuda=CUDA_VERSION" in source + benchmark_body = source.split("def benchmark_fn(self)", 1)[1].split( + "def capture_verification_payload", 1 + )[0] + assert "load_nvfp4_extension" not in benchmark_body + assert "subprocess" not in benchmark_body + + +def test_runtime_gate_reports_the_real_host_without_mocking_success() -> None: + supported = ( + torch.cuda.is_available() + and torch.version.cuda is not None + and tuple(map(int, torch.version.cuda.split(".")[:2])) >= (13, 1) + and torch.cuda.get_device_capability(torch.cuda.current_device()) == (10, 3) + ) + if supported: + assert require_nvfp4_runtime() == torch.cuda.current_device() + else: + with pytest.raises(RuntimeError, match=r"SKIPPED:.*(?:CUDA|SM103)"): + require_nvfp4_runtime() + + +@pytest.mark.skipif( + os.environ.get("AISP_RUN_FAST_CU_NVFP4_GPU_TEST") != "1", + reason="set AISP_RUN_FAST_CU_NVFP4_GPU_TEST=1 on CUDA 13.1 SM103 hardware", +) +def test_real_sm103_native_pair_matches_full_output_on_current_stream() -> None: + """Compile and run both native arms; no mocked GPU/runtime success is allowed.""" + require_nvfp4_runtime() + workload = Nvfp4Workload(4096, 4096, 1024) + seed = 20260916 + + with preserve_rng_state(): + torch.manual_seed(seed) + baseline = FastCuNvfp4CublasLtBenchmark(workload) + optimized = FastCuNvfp4KernelBenchmark(workload, rung=9) + stream = torch.cuda.Stream() + try: + with torch.cuda.stream(stream): + baseline.setup() + optimized.setup() + assert baseline.output is None and optimized.output is None + assert baseline._setup_gate_passed and optimized._setup_gate_passed + with pytest.raises(RuntimeError, match=r"benchmark_fn\(\) must run"): + baseline.capture_verification_payload() + with pytest.raises(RuntimeError, match=r"benchmark_fn\(\) must run"): + optimized.capture_verification_payload() + assert all( + torch.equal(baseline._inputs[name], optimized._inputs[name]) + for name in baseline._inputs + ) + with torch.cuda.stream(stream): + baseline.benchmark_fn() + optimized.benchmark_fn() + stream.synchronize() + baseline.capture_verification_payload() + optimized.capture_verification_payload() + baseline_output = baseline.get_verify_output() + optimized_output = optimized.get_verify_output() + + assert baseline_output.numel() == workload.m * workload.n + assert torch.isfinite(baseline_output).all() + assert torch.isfinite(optimized_output).all() + torch.testing.assert_close( + optimized_output, + baseline_output, + rtol=2e-3, + atol=0.5, + ) + finally: + optimized.teardown() + baseline.teardown() diff --git a/code/tests/test_fast_cu_source.py b/code/tests/test_fast_cu_source.py new file mode 100644 index 000000000..bac7e97ac --- /dev/null +++ b/code/tests/test_fast_cu_source.py @@ -0,0 +1,73 @@ +"""Real file-system checks for the pinned fast.cu source supply chain.""" + +from __future__ import annotations + +import json +import shutil + +import pytest + +from labs.fast_cu.source import ( + MANIFEST_PATH, + UPSTREAM_DIR, + UPSTREAM_REVISION, + verify_upstream, +) + + +def test_complete_pinned_source_tree(): + manifest = verify_upstream() + assert manifest["revision"] == UPSTREAM_REVISION + assert len(manifest["files"]) == 34 + assert "MIT License" in (UPSTREAM_DIR / "LICENSE").read_text() + for rung in range(10): + assert f"gb300/nvfp4/gemm{rung}.cuh" in manifest["files"] + for rung in range(1, 13): + assert f"h100/matmul/matmul_{rung}.cuh" in manifest["files"] + + +@pytest.fixture +def copied_source(tmp_path): + root = tmp_path / "upstream" + shutil.copytree(UPSTREAM_DIR, root) + manifest_path = tmp_path / "manifest.json" + shutil.copyfile(MANIFEST_PATH, manifest_path) + return root, manifest_path + + +def test_modified_source_rejected(copied_source): + root, manifest_path = copied_source + (root / "h100" / "sum.cu").write_text("// modified\n") + with pytest.raises(ValueError, match="integrity mismatch: h100/sum.cu"): + verify_upstream(root, manifest_path) + + +def test_missing_license_rejected(copied_source): + root, manifest_path = copied_source + (root / "LICENSE").rename(root / "LICENSE.moved") + with pytest.raises(ValueError, match="file set changed"): + verify_upstream(root, manifest_path) + + +def test_extra_header_rejected(copied_source): + root, manifest_path = copied_source + (root / "injected.cuh").write_text("// unpinned include\n") + with pytest.raises(ValueError, match="extra=.*injected.cuh"): + verify_upstream(root, manifest_path) + + +def test_symlink_rejected(copied_source): + root, manifest_path = copied_source + (root / "linked.cuh").symlink_to(root / "h100" / "sum.cu") + with pytest.raises(ValueError, match="Symlink"): + verify_upstream(root, manifest_path) + + +@pytest.mark.parametrize("field,value", [("revision", "0" * 40), ("schema_version", 99)]) +def test_changed_manifest_identity_rejected(copied_source, field, value): + root, manifest_path = copied_source + manifest = json.loads(manifest_path.read_text()) + manifest[field] = value + manifest_path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="manifest schema|source revision"): + verify_upstream(root, manifest_path) From a103b5c842736e95601dc10cf2620b5388386ac1 Mon Sep 17 00:00:00 2001 From: Chris Fregly Date: Wed, 16 Sep 2026 17:23:10 -0700 Subject: [PATCH 2/2] Validate all SM103 rungs and preserve generated chapter mappings --- code/core/scripts/refresh_readmes.py | 25 ++++++++++++++ code/labs/fast_cu/nvfp4_native.cu | 9 ++--- code/labs/fast_cu/validation.json | 50 ++++++++++++++++++++++++---- code/labs/fast_cu/validation.md | 20 +++++++++-- code/tests/test_fast_cu_nvfp4.py | 5 ++- 5 files changed, 95 insertions(+), 14 deletions(-) diff --git a/code/core/scripts/refresh_readmes.py b/code/core/scripts/refresh_readmes.py index afa35c7dc..eafebc5bc 100644 --- a/code/core/scripts/refresh_readmes.py +++ b/code/core/scripts/refresh_readmes.py @@ -48,6 +48,7 @@ class Entry: validation_heading: str = "Validation Checklist" extra_sections: Sequence[str] = field(default_factory=tuple) notes: Sequence[str] = field(default_factory=tuple) + trailing_sections: Sequence[str] = field(default_factory=tuple) def _format_markdown(entry: Entry) -> str: @@ -97,6 +98,12 @@ def _format_markdown(entry: Entry) -> str: lines.append("## Notes") for note in entry.notes: lines.append(f"- {note}") + if entry.trailing_sections: + lines.append("") + for idx, section in enumerate(entry.trailing_sections): + if idx: + lines.append("") + lines.append(section.strip()) lines.append("") return "\n".join(lines) @@ -404,6 +411,7 @@ def chapter_entry( "when you need repeatable artifact capture." ), validation_heading: str = "Validation Checklist", + trailing_sections: Sequence[str] = (), ) -> Entry: """Create a chapter README entry.""" return Entry( @@ -418,6 +426,7 @@ def chapter_entry( run_intro=run_intro, validation_heading=validation_heading, notes=notes, + trailing_sections=trailing_sections, ) @@ -660,6 +669,17 @@ def lab_entry( """ ).strip() +FAST_CU_RELATED_SECTION = dedent( + """\ + ## Related fast.cu experiments + + The [fast.cu lab](../labs/fast_cu/README.md) preserves the complete H100 and GB300 + source progressions and adds native harness pairs, including B200 adaptations. + See its [chapter mapping and optimization guide](../labs/fast_cu/optimization_guide.md) + for the overlap with this chapter and the SM100 versus SM103 hardware boundaries. + """ +) + ENTRIES["README.md"] = Entry( title="AI Systems Performance Engineering", summary=dedent( @@ -965,6 +985,7 @@ def lab_entry( | `labs/cutlass_profiler_kernel_selector/` | CUTLASS profiler-based kernel selection | ch06, ch09 | | `labs/decode_optimization/` | Decoder hot-path optimization | ch18, ch19 | | `labs/dynamic_router/` | Dynamic prefill/decode routing | ch17, ch19 | + | `labs/fast_cu/` | Pinned fast.cu H100 GEMM, B200 reduction and wide-store epilogue adaptations, and the complete GB300 NVFP4 ladder; see the [chapter mapping](fast_cu/optimization_guide.md) | ch09, ch10, ch11, ch14 | | `labs/flashattention4/` | FlashAttention-4 pipeline co-design, including a [Colfax diary guide](flashattention4/colfax_optimization_diaries.md) to decode S/P ping-pong and backward hdim-64 TMEM de-aliasing | ch10, ch11, ch18 | | `labs/flashattention_gluon/` | FlashAttention experimentation | ch18 | | `labs/flashinfer_attention/` | FlashInfer block-sparse attention lab | ch16 | @@ -1972,6 +1993,7 @@ def lab_entry( "`requirements.txt` includes the shared `requirements_latest.txt` pins for PyTorch 2.9.1 and its compatible Triton release.", "`optimized_cublaslt_gemm_fp4` is intentionally capability-gated: if cuBLASLt cannot provide the native block-scaled NVFP4 heuristic, the benchmark reports a clean skip instead of silently falling back to a different FP4 mode.", ], + trailing_sections=[FAST_CU_RELATED_SECTION], ) ENTRIES["ch10"] = chapter_entry( @@ -2077,6 +2099,7 @@ def lab_entry( "`requirements_cufile.txt` holds the optional `cufile` wheel; install it only on hosts with GPUDirect Storage enabled.", "The CUTLASS-style warp-specialization pair provides a reference implementation aligned with `sm100_mma_array_warpspecialized` for performance comparison.", ], + trailing_sections=[FAST_CU_RELATED_SECTION], ) ENTRIES["ch11"] = chapter_entry( @@ -2183,6 +2206,7 @@ def lab_entry( "`warp_specialized_triton.py` provides a Triton analogue for the CUDA concurrency demos so you can compare compiler-generated schedules.", "`kv_prefetch_pipeline_enhanced_demo.cu` builds on the DSMEM kernels bundled in this directory so you can study the entire pipeline locally.", ], + trailing_sections=[FAST_CU_RELATED_SECTION], ) ENTRIES["ch12"] = chapter_entry( @@ -2554,6 +2578,7 @@ def lab_entry( "For repo-native supporting examples that fill the training hot-path gaps without changing this chapter's primary compile narrative, see `labs/training_hotpath`.", "`cublas_vs_cutlass` is a supplementary comparison pair. Chapter-native performance claims stay anchored on `model_compile_reduced_precision`, `regional_triton`, and `triton_persistent`.", ], + trailing_sections=[FAST_CU_RELATED_SECTION], ) ENTRIES["ch15"] = chapter_entry( diff --git a/code/labs/fast_cu/nvfp4_native.cu b/code/labs/fast_cu/nvfp4_native.cu index 4046d2a8c..24d06d1b1 100644 --- a/code/labs/fast_cu/nvfp4_native.cu +++ b/code/labs/fast_cu/nvfp4_native.cu @@ -1,7 +1,7 @@ #include -#include #include +#include #include #include @@ -82,7 +82,7 @@ cudaStream_t current_stream(int device) { device, " but current device is ", current_device); - return at::cuda::getCurrentCUDAStream(device); + return c10::cuda::getCurrentCUDAStream(device); } void require_exact_runtime(cudaDeviceProp* properties, int* device) { @@ -322,8 +322,9 @@ public: TORCH_CHECK(fast_prepared_, "fast.cu kernel was not prepared for this context"); c10::cuda::CUDAGuard guard(device_); #if NVFP4_HAS_SCHEDULE + const std::array context_shape{M_, N_, K_}; TORCH_CHECK( - scheduled_device == device_ && scheduled_shape == std::array{M_, N_, K_}, + scheduled_device == device_ && scheduled_shape == context_shape, "fast.cu r9 schedule no longer matches this context"); #endif cudaStream_t stream = current_stream(device_); @@ -555,7 +556,7 @@ private: } // namespace PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { - py::class_(module, "Nvfp4Context") + py::class_(module, "Nvfp4Context", py::module_local()) .def(py::init()) .def("launch_cublaslt", &Nvfp4Context::launch_cublaslt) .def("launch_fast", &Nvfp4Context::launch_fast) diff --git a/code/labs/fast_cu/validation.json b/code/labs/fast_cu/validation.json index cc69f97c8..60ed46e4f 100644 --- a/code/labs/fast_cu/validation.json +++ b/code/labs/fast_cu/validation.json @@ -24,7 +24,7 @@ "native_sources_sha256": { "h100_reduction_kernels.cu": "a8d786df751482c5a431b22b279548adbed09600621a167ca578d701e44aed48", "h100_gemm_extension.cu": "e0442d4dfc4910ba981034d1c2544dcdaebabb464d62c115e2f755f4c272174c", - "nvfp4_native.cu": "ac4e8ea997c85ba4d4c01244208a1163d5febda0396d8bd3068e84f6d9c890f9", + "nvfp4_native.cu": "2ae5c0b8bb5b8a6e60adae4e28ab5541e072c269144aea62d3eb89d405f066be", "b200_epilogue.cu": "806f6226b44b6f1cd99a8e944650ae58571ee4647fb3b1b1d5163e2bb283d920" }, "results": [ @@ -34,7 +34,7 @@ "baseline_ms": 0.1883358495159113, "candidate_ms": 0.1801341006920921, "speedup": 1.0455313502124657, - "minimum_required_speedup": null, + "minimum_required_speedup": 1.05, "input_verification": true, "output_verification": true, "baseline_profilers": { @@ -47,7 +47,8 @@ "ncu": "succeeded", "torch": "succeeded" }, - "source_result_sha256": "b8cec133bd9b886311f8784d938783b46f91a0a4726c90c61fa4129a38f501f2" + "source_result_sha256": "b8cec133bd9b886311f8784d938783b46f91a0a4726c90c61fa4129a38f501f2", + "speedup_threshold_source": "repository harness default; raw expectation override was null" }, { "target": "labs/fast_cu:b200_epilogue", @@ -55,7 +56,7 @@ "baseline_ms": 0.05168658194423338, "candidate_ms": 0.048537385734480486, "speedup": 1.0648818670824238, - "minimum_required_speedup": null, + "minimum_required_speedup": 1.05, "input_verification": true, "output_verification": true, "baseline_profilers": { @@ -68,12 +69,49 @@ "ncu": "succeeded", "torch": "succeeded" }, - "source_result_sha256": "6a63865270a5fdfbb5a652e028bce169a68f565855caf427b2fb181fe01be4ec" + "source_result_sha256": "6a63865270a5fdfbb5a652e028bce169a68f565855caf427b2fb181fe01be4ec", + "speedup_threshold_source": "repository harness default; raw expectation override was null" } ], "limitations": [ "No H100 runtime validation", "No SM103 runtime validation", "No published speedup or expectation baseline from this single exploratory run" - ] + ], + "compile_only_checks": { + "nvfp4": { + "status": "passed", + "rungs": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9 + ], + "cuda_compiler": "13.1.80", + "cuda_architecture": "compute_103a,sm_103a", + "pytorch": "2.9.1+cu130", + "compiled_linked_and_imported": true, + "modules_coexisting_in_one_process": 10, + "gpu_visible": false, + "kernels_executed": false, + "qualification": "build_only_not_sm103_runtime_validation", + "source_log_sha256": "91deed619c03dbee5414c0d7116e156632cf471522536b26bd928d4aea4643e8", + "cublas_package": "nvidia-cublas 13.0.0.19", + "runtime_headers_and_cudart": "CUDA 13.1.80", + "aggregate_compile_link_import_seconds": 277.093 + }, + "h100_gemm": { + "status": "passed", + "cuda_compiler": "13.0.88", + "cuda_architecture": "compute_90a,sm_90a", + "compiled_and_linked": true, + "kernels_executed": false + } + } } diff --git a/code/labs/fast_cu/validation.md b/code/labs/fast_cu/validation.md index 53e2c9a3a..bc38babb6 100644 --- a/code/labs/fast_cu/validation.md +++ b/code/labs/fast_cu/validation.md @@ -43,14 +43,19 @@ B200, and **63 tests with 5 expected hardware skips** on the local CPU host. build-cache identity, explicit hardware/toolkit contracts, target discovery, output lifecycle, and selectable NVFP4 rungs. - All four targets are discoverable through `aisp bench list-targets`. +- All 12 README-generator tests pass, and the lab index plus four chapter + READMEs match their generator definitions. - Repository-wide benchmark lint checked 944 entrypoints with zero errors and zero warnings. Focused Ruff, formatting, syntax, YAML, and local-link checks also passed. - The H100 adapter compiled and linked for `sm_90a` with CUDA 13.0. No H100 kernel was executed; numerical and performance qualification still needs H100. -- The SM103 adapter retains its CUDA 13.1+ requirement. Its independent host - oracle, poisoned-output, guard, determinism, and placement checks are part of - the adapter, but SM103 execution remains unqualified without B300/GB300. +- Every SM103 rung, r0-r9, compiled, linked, and imported together in one Python + process with CUDA compiler 13.1.80 and `compute_103a,sm_103a`. This isolated + compile-only check used PyTorch 2.9.1+cu130, cuBLAS 13.0.0.19 headers/library, + and CUDA 13.1.80 runtime headers/library, with no GPU visible and no kernel launches. The production adapter still requires CUDA runtime and toolkit 13.1+ + on exact SM103 hardware. Its independent host oracle, poisoned-output, guard, + determinism, and placement checks remain unqualified without B300/GB300. ## Failures caught and corrected @@ -70,6 +75,15 @@ B200, and **63 tests with 5 expected hardware skips** on the local CPU host. 6. The host's system Python reported conflicting installed-library versions. The harness correctly rejected that comparison; reported results use an isolated environment with unambiguous package metadata. +7. The final documentation audit caught links added only to generated READMEs. + Their definitions now live in the README generator, so regeneration preserves + the chapter mappings. +8. Compiling with CUDA 13.1 exposed an overly broad PyTorch include and a macro + argument containing a template comma. The adapter now includes only the CUDA + stream API it uses and gives the schedule shape a named local variable. +9. Importing multiple compiled NVFP4 rungs exposed a Pybind global-type collision. + Each module now registers its context type locally. The all-rung import check + passed, and the opt-in SM103 test also imports r0 after setting up r9. ## Reproduce diff --git a/code/tests/test_fast_cu_nvfp4.py b/code/tests/test_fast_cu_nvfp4.py index 955118bda..1ad1c6b7b 100644 --- a/code/tests/test_fast_cu_nvfp4.py +++ b/code/tests/test_fast_cu_nvfp4.py @@ -19,6 +19,7 @@ DEFAULT_N, DEFAULT_RUNG, Nvfp4Workload, + load_nvfp4_extension, require_nvfp4_runtime, validate_rung, workload_bytes, @@ -141,7 +142,7 @@ def test_native_source_preserves_the_equivalent_workload_contract() -> None: def test_native_source_uses_current_stream_and_rejects_r9_degradation() -> None: source = NATIVE_SOURCE.read_text(encoding="utf-8") - assert "at::cuda::getCurrentCUDAStream(device)" in source + assert "c10::cuda::getCurrentCUDAStream(device)" in source assert "cublasLtMatmul(" in source and "stream)," in source assert "<<>>" in source assert "strict_setup_r9_schedule" in source @@ -202,6 +203,8 @@ def test_real_sm103_native_pair_matches_full_output_on_current_stream() -> None: with torch.cuda.stream(stream): baseline.setup() optimized.setup() + r0_module = load_nvfp4_extension(0) + assert hasattr(r0_module, "Nvfp4Context") assert baseline.output is None and optimized.output is None assert baseline._setup_gate_passed and optimized._setup_gate_passed with pytest.raises(RuntimeError, match=r"benchmark_fn\(\) must run"):