diff --git a/.gitignore b/.gitignore index 4ad6f92ff..e0e116ee3 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,10 @@ build/ +build_new/ +build_cast/ +build_adam/ .cache/ .vscode/ +nsys/out/ *.log *.report.rank* diff --git a/CMakeLists.txt b/CMakeLists.txt index 709bc30f2..59819e709 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2,6 +2,7 @@ cmake_minimum_required(VERSION 3.28) option(USE_CUDA "Support NVIDIA CUDA" OFF) option(PROFILE_MODE "ENABLE PROFILE MODE" OFF) +option(NVTX_MODE "Emit NVTX ranges for nsys timeline analysis (requires USE_CUDA; adds no CUDA synchronization)" OFF) option(USE_OMP "Use OpenMP as backend for Eigen" ON) option(USE_NCCL "Build project for distributed running on CUDA using NCCL" ON) option(BUILD_TEST "Build InfiniTrain tests" OFF) @@ -65,6 +66,15 @@ if(PROFILE_MODE) add_compile_definitions(PROFILE_MODE=1) endif() +if(NVTX_MODE) + if(NOT USE_CUDA) + message(FATAL_ERROR "NVTX_MODE=ON requires USE_CUDA=ON: ships with the CUDA toolkit.") + endif() + # Unlike PROFILE_MODE this does not serialize the pipeline, so the two can be + # enabled independently; NVTX_MODE alone is the one to use for timing work. + add_compile_definitions(NVTX_MODE=1) +endif() + # ------------------------------------------------------------------------------ # Sources # ------------------------------------------------------------------------------ @@ -104,10 +114,27 @@ endif() if(USE_CUDA) add_compile_definitions(USE_CUDA=1) + + # Must be set before enable_language(CUDA): CMake's built-in default (52) is + # rejected by nvcc 13.x, whose minimum supported architecture is compute_75. + # 75=Turing, 80=Ampere, 90=Hopper, 120=Blackwell GeForce (RTX 50 series). + # Pass -DCMAKE_CUDA_ARCHITECTURES=120 for a faster build targeting one GPU. + if(NOT CMAKE_CUDA_ARCHITECTURES) + set(CMAKE_CUDA_ARCHITECTURES "75;80;90;120") + endif() + enable_language(CUDA) find_package(CUDAToolkit REQUIRED) include_directories(${CUDAToolkit_INCLUDE_DIRS}) + if(NVTX_MODE) + # nvtx3 is header-only and needs no library (libnvToolsExt.so was dropped in + # CUDA 13); the injection library is resolved by dlopen at runtime. Fail at + # configure time rather than halfway through compiling 24 CUDA sources. + find_path(NVTX3_INCLUDE_DIR nvtx3/nvToolsExt.h HINTS ${CUDAToolkit_INCLUDE_DIRS} REQUIRED) + message(STATUS "NVTX_MODE enabled, nvtx3 headers found in: ${NVTX3_INCLUDE_DIR}") + endif() + # CUDA compilation options set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-extended-lambda --expt-relaxed-constexpr") @@ -115,13 +142,14 @@ if(USE_CUDA) file(GLOB_RECURSE CUDA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/*.cu) add_library(infini_train_cuda_kernels STATIC ${CUDA_KERNELS}) - set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90") + set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}") target_link_libraries(infini_train_cuda_kernels PUBLIC glog CUDA::cudart CUDA::cublas + CUDA::cublasLt CUDA::cuda_driver ) @@ -162,6 +190,7 @@ if(USE_CUDA) PUBLIC CUDA::cudart CUDA::cublas + CUDA::cublasLt CUDA::cuda_driver ) diff --git a/README.md b/README.md index abd8070b2..72c42ea91 100644 --- a/README.md +++ b/README.md @@ -50,7 +50,7 @@ Build Options: > Both options are optional and can be disabled for CPU-only builds. -## ✨ InfiniTrain Overview +## ✨ InfiniTrain Overview ### ✔ Support Matrix @@ -96,160 +96,166 @@ For example, the `llama3` example produces a binary named `llama3`. To view available runtime options: -```bash -./build/llama3 --help +```bash +./build/llama3 --help +``` + +### Getting Started + +#### Prepare Datasets and Weights + +Run the asset preparation script from the repository root. Prepared files are +written to `data/` by default. + +```bash +# MNIST dataset +./scripts/assets/prepare-infinitrain-assets.sh mnist + +# GPT-2 124M weights, tokenizer, and tokenized TinyShakespeare data +./scripts/assets/prepare-infinitrain-assets.sh gpt2 + +# LLaMA 3.2 1B weights and tokenized TinyShakespeare data +HF_TOKEN=hf_xxx ./scripts/assets/prepare-infinitrain-assets.sh llama3 + +# Same flow through ModelScope +MODEL_SOURCE=modelscope MODEL_REPO_ID=LLM-Research/Meta-Llama-3.2-1B \ + ./scripts/assets/prepare-infinitrain-assets.sh llama3 +``` + +Preparing LLaMA requires access to the gated +`meta-llama/Llama-3.2-1B` repository. Accept its license on Hugging Face and +provide `HF_TOKEN`, or authenticate with `hf auth login`, before running the +command. If the Hugging Face download is blocked, set `MODEL_SOURCE=modelscope` +and optionally override `MODEL_REPO_ID` to the mirror you have access to. +The complete LLaMA preparation requires approximately 8.5 GB of free disk +space, including the downloaded checkpoint and converted FP32 weights. + +Use `DATA_DIR` to write the assets elsewhere, or prepare all supported assets +in one invocation: + +```bash +DATA_DIR=/path/to/data \ +HF_TOKEN=hf_xxx \ +./scripts/assets/prepare-infinitrain-assets.sh all +``` + +#### Model Examples + +The generated files can be passed directly to the corresponding executables: + +##### MNIST + +```bash +./build/mnist \ + --device cpu \ + --dataset data/mnist +``` + +##### GPT-2 124M + +```bash +./build/gpt2 \ + --device cuda \ + --input_bin data/gpt2/tiny_shakespeare_train.bin \ + --input_val_bin data/gpt2/tiny_shakespeare_val.bin \ + --tokenizer_bin data/gpt2/gpt2_tokenizer.bin \ + --llmc_filepath data/gpt2/gpt2_124M.bin \ + --num_iteration 10 +``` + +##### LLaMA 3.2 1B + +```bash +./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --input_val_bin data/llama3/tiny_shakespeare_val.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 10 +``` + +### Launch Modes + +GPT-2 and LLaMA training support both thread-based and process-based launches. +The examples below use LLaMA, but the same launch modes also apply to GPT-2. + +#### Direct Launch + +Running a model executable directly uses one process and one device by default. +Set `--nthread_per_process` to use multiple execution threads and devices in the +same process: + +```bash +./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --nthread_per_process 8 \ + --num_iteration 10 +``` + +#### Single-node Multi-process Launch + +Use `infini_run` to start multiple training processes on one node. Each process +uses one execution thread by default: + +```bash +./build/infini_run \ + --nnodes=1 \ + --nproc_per_node=8 \ + ./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 10 +``` + +#### Multi-node Multi-process Launch + +Run the following command on every node with the same rendezvous settings and +a distinct `node_rank`: + +```bash +./build/infini_run \ + --nnodes=2 \ + --nproc_per_node=4 \ + --node_rank=[rank_id] \ + --rdzv_endpoint=[master_addr]:29500 \ + --rdzv_id=[job_id] \ + ./build/llama3 \ + --device cuda \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 10 \ + --tensor_parallel 2 \ + --pipeline_parallel 2 \ + --sequence_parallel ``` -### Getting Started - -#### Prepare Datasets and Weights - -Run the asset preparation script from the repository root. Prepared files are -written to `data/` by default. - -```bash -# MNIST dataset -./scripts/assets/prepare-infinitrain-assets.sh mnist - -# GPT-2 124M weights, tokenizer, and tokenized TinyShakespeare data -./scripts/assets/prepare-infinitrain-assets.sh gpt2 - -# LLaMA 3.2 1B weights and tokenized TinyShakespeare data -HF_TOKEN=hf_xxx ./scripts/assets/prepare-infinitrain-assets.sh llama3 -``` - -Preparing LLaMA requires access to the gated -`meta-llama/Llama-3.2-1B` repository. Accept its license on Hugging Face and -provide `HF_TOKEN`, or authenticate with `hf auth login`, before running the -command. The complete LLaMA preparation requires approximately 8.5 GB of free -disk space, including the downloaded checkpoint and converted FP32 weights. - -Use `DATA_DIR` to write the assets elsewhere, or prepare all supported assets -in one invocation: - -```bash -DATA_DIR=/path/to/data \ -HF_TOKEN=hf_xxx \ -./scripts/assets/prepare-infinitrain-assets.sh all -``` - -#### Model Examples - -The generated files can be passed directly to the corresponding executables: - -##### MNIST - -```bash -./build/mnist \ - --device cpu \ - --dataset data/mnist -``` - -##### GPT-2 124M - -```bash -./build/gpt2 \ - --device cuda \ - --input_bin data/gpt2/tiny_shakespeare_train.bin \ - --input_val_bin data/gpt2/tiny_shakespeare_val.bin \ - --tokenizer_bin data/gpt2/gpt2_tokenizer.bin \ - --llmc_filepath data/gpt2/gpt2_124M.bin \ - --num_iteration 10 -``` - -##### LLaMA 3.2 1B - -```bash -./build/llama3 \ - --device cuda \ - --input_bin data/llama3/tiny_shakespeare_train.bin \ - --input_val_bin data/llama3/tiny_shakespeare_val.bin \ - --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ - --num_iteration 10 -``` - -### Launch Modes - -GPT-2 and LLaMA training support both thread-based and process-based launches. -The examples below use LLaMA, but the same launch modes also apply to GPT-2. - -#### Direct Launch - -Running a model executable directly uses one process and one device by default. -Set `--nthread_per_process` to use multiple execution threads and devices in the -same process: - -```bash -./build/llama3 \ - --device cuda \ - --input_bin data/llama3/tiny_shakespeare_train.bin \ - --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ - --nthread_per_process 8 \ - --num_iteration 10 -``` - -#### Single-node Multi-process Launch - -Use `infini_run` to start multiple training processes on one node. Each process -uses one execution thread by default: - -```bash -./build/infini_run \ - --nnodes=1 \ - --nproc_per_node=8 \ - ./build/llama3 \ - --device cuda \ - --input_bin data/llama3/tiny_shakespeare_train.bin \ - --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ - --num_iteration 10 -``` - -#### Multi-node Multi-process Launch - -Run the following command on every node with the same rendezvous settings and -a distinct `node_rank`: - -```bash -./build/infini_run \ - --nnodes=2 \ - --nproc_per_node=4 \ - --node_rank=[rank_id] \ - --rdzv_endpoint=[master_addr]:29500 \ - --rdzv_id=[job_id] \ - ./build/llama3 \ - --device cuda \ - --input_bin data/llama3/tiny_shakespeare_train.bin \ - --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ - --num_iteration 10 \ - --tensor_parallel 2 \ - --pipeline_parallel 2 \ - --sequence_parallel -``` - -`--nproc_per_node` and `--nthread_per_process` can be combined. The total -training world size is: - -```text -world_size = nnodes × nproc_per_node × nthread_per_process -``` +`--nproc_per_node` and `--nthread_per_process` can be combined. The total +training world size is: + +```text +world_size = nnodes × nproc_per_node × nthread_per_process +``` ### Parallelism Strategies -#### Distributed Data Parallelism (DDP) - -For a direct launch with TP and PP disabled, the following starts eight -data-parallel workers in one process: - -```bash ---nthread_per_process 8 # 8-way DDP when TP=1 and PP=1 -``` - -For all launch modes, the data-parallel size is derived from the total world -size after accounting for tensor and pipeline parallelism: - -```text -data_parallel_size = world_size / (tensor_parallel × pipeline_parallel) -``` +#### Distributed Data Parallelism (DDP) + +For a direct launch with TP and PP disabled, the following starts eight +data-parallel workers in one process: + +```bash +--nthread_per_process 8 # 8-way DDP when TP=1 and PP=1 +``` + +For all launch modes, the data-parallel size is derived from the total world +size after accounting for tensor and pipeline parallelism: + +```text +data_parallel_size = world_size / (tensor_parallel × pipeline_parallel) +``` #### Tensor Parallelism (TP) @@ -269,6 +275,182 @@ data_parallel_size = world_size / (tensor_parallel × pipeline_parallel) Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scale training across devices and nodes. +## 🔍 Profiling with nsys + +InfiniTrain emits NVTX ranges for every training step and every major phase +(`Forward`, `Backward`, `Optimizer`, `LossReadback`, ...) when built with +`-DNVTX_MODE=ON`. Combine this with `nsys` to locate GPU bottlenecks without +adding any CUDA synchronization. + +### Build with NVTX + +```bash +cmake .. -DUSE_CUDA=ON -DUSE_NCCL=ON -DNVTX_MODE=ON +make -j llama3 +``` + +### Profile LLaMA 3 for 5 iterations + +```bash +nsys profile \ + --trace=cuda,nvtx,osrt,cudnn,cublas \ + --sample=none \ + --cpuctxsw=none \ + --force-overwrite=true \ + --output=nsys/out/llama3_5iter_nvtx \ + --export=sqlite \ + ./build/llama3 \ + --device cuda \ + --dtype bfloat16 \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --input_val_bin data/llama3/tiny_shakespeare_val.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 5 \ + --batch_size 4 \ + --sequence_length 64 \ + --total_batch_size 256 +``` + +Flags worth noting: + +- `--trace=cuda,nvtx,osrt,cudnn,cublas` collects CUDA runtime/driver, NVTX + ranges, OS runtime, cuDNN and cuBLAS events. Add `nccl` for multi-GPU runs. +- `--sample=none --cpuctxsw=none` disables CPU sampling and context-switch + tracing to keep the capture small and the overhead low. +- `--export=sqlite` produces a `.sqlite` file alongside the `.nsys-rep` so the + results can be queried with SQL or Python. +- `--force-overwrite=true` lets repeated runs reuse the same output name. + +### Generate summary reports + +```bash +nsys stats \ + --force-export=true \ + --report cuda_gpu_kern_sum \ + --report cuda_gpu_sum \ + --report cuda_api_sum \ + --report nvtx_sum \ + --report nvtx_pushpop_sum \ + --report cuda_gpu_mem_time_sum \ + --report cuda_gpu_mem_size_sum \ + --format table \ + --output nsys/out/llama3_5iter_stats \ + nsys/out/llama3_5iter_nvtx.nsys-rep +``` + +This writes one `.txt` per report next to the capture. Open the `.nsys-rep` in +the Nsight Systems GUI for the full timeline. + +### Per-step Python breakdown + +`nsys/out/analyze_steps.py` queries the exported sqlite and prints per-step +wall time, GPU busy ratio, the hot kernels inside the first steady-state step, +and the memcpy breakdown: + +```bash +python3 nsys/out/analyze_steps.py +``` + +### Reference result (A100-SXM4-40GB, LLaMA 3.2 1B, BF16 autocast, 5 iters) + +The run above uses `--dtype bfloat16`, which turns on **autocast** mixed +precision: `Matmul`/`Linear` run on BF16 Tensor Cores while the master weights +and the Adam optimizer stay in FP32 (see `infini_train/include/autocast.h`). + +| step | wall (ms) | Σ kernel (ms) | kernels | GPU util | +| ------ | --------- | ------------- | ------- | -------- | +| Step_0 | 273.031 | 54.854 | 3483 | 21.14% | +| Step_1 | 118.333 | 83.362 | 3569 | 73.13% | +| Step_2 | 114.822 | 76.690 | 3574 | 69.33% | +| Step_3 | 112.973 | 73.618 | 3574 | 67.52% | +| Step_4 | 112.404 | 73.628 | 3574 | 68.05% | + +Step_0 is warm-up (CUDA context init, cuBLAS BF16 kernel selection, first-time +cast-buffer growth); it is host-bound, so GPU util is only ~21%. From Step_1 +onward the GPU is busy ~68–73% of the wall time — noticeably lower than FP32's +~93%, because BF16 compute is much lighter and the step becomes +**launch/cast-bound** on a single CUDA stream (no kernel overlap). + +Top steady-state (Step_1) kernels: + +| rank | kernel | calls | Σ (ms) | share | +| ---- | ------------------------------------------------- | ----- | --------- | --------- | +| 1 | `AdamAccumulateGradKernel` | 110 | 31.7 | 38.1% | +| 2 | `CastKernel<__nv_bfloat16, float>` (f32→bf16) | 356 | 9.5 | 11.4% | +| 3 | `ampere_s16816gemm_bf16_128x128_..._nt` | 49 | 4.6 | 5.5% | +| 4 | `ampere_bf16_s16816gemm_bf16_128x256_..._f2f_tn` | 49 | 4.5 | 5.4% | +| 5 | `BinaryBackwardKernel` (Mul) | 194 | 3.7 | 4.4% | +| 6 | `FillKernel` | 678 | 3.2 | 3.8% | +| 7 | `BinaryForwardKernel` (Mul) | 194 | 2.9 | 3.5% | +| 8 | `TransposeForwardKernel` | 128 | 2.5 | 3.0% | +| – | **all GEMM kernels combined (BF16 Tensor Core)** | 339 | **16.8** | **20.0%** | +| – | **all Cast kernels combined (autocast)** | 661 | **10.7** | **12.8%** | + +Sub-phase breakdown of Step_1 (118.3 ms wall; Σ kernel = GPU time executing +inside each NVTX window): + +| phase | wall (ms) | Σ kernel (ms) | util | +| ------------------------- | --------- | ------------- | ----- | +| `Forward` | 61.9 | 53.1 | 85.7% | +| └ `CrossEntropyForward` | 6.5 | 5.4 | 82.6% | +| `Backward` | 53.6 | 28.2 | 52.7% | +| `LossReadback` | 0.8 | 0.0 | – | +| `Optimizer` | 1.1 | 1.3 | – | + +> The `Forward` window shows 53.1 ms of GPU execution but only 23.4 ms was +> actually launched in-phase — the ~30 ms gap is Step_0's Adam optimizer +> draining on the GPU at the start of Step_1. `nsys/out/llama3_kernel_timeline.md` +> uses host-launch attribution (via `correlationId`) to correct for this and +> gives the per-phase totals: Forward 1343 kernels / 23.4 ms, Backward 2116 / +> 28.2 ms, Optimizer 115 / 32.3 ms. + +FP32 → BF16 comparison (same 5-iteration workload, Step_1 steady state): + +| metric | FP32 | BF16 (autocast) | +| ---------------------------- | ------------- | --------------- | +| Step_1 wall | 181.8 ms | 118.3 ms (1.54×)| +| steady-state throughput | ~1 424 tok/s | ~2 260 tok/s | +| GPU util (steady) | ~93% | ~68–73% | +| all GEMM (339 calls) | 113.9 ms (68.5%) | 16.8 ms (20.0%) — **6.8× faster** | +| fp32 `sgemm` calls | 339 | ≈0 (all Tensor Core) | +| autocast Cast | — | 661 calls / 10.7 ms | +| Adam optimizer | 31.9 ms (19.2%) | 32.3 ms (38.5%) | +| `CrossEntropyForward` phase | 32.5 ms | 6.5 ms | +| `LossReadback` phase | 29.3 ms | 0.8 ms | + +Key observations: + +- **BF16 Tensor Core GEMM (6.8× faster).** All 339 GEMMs run on BF16 + `s16816`/`s161616` tensor cores (fp32 `sgemm` ≈ 0), cutting GEMM time from + 113.9 ms to 16.8 ms. The `lm_head` GEMM (vocab = 128 256) alone drops from + ~7.3 ms to ~0.70 ms — which is why `CrossEntropyForward` shrinks 32.5 → 6.5 ms + and `LossReadback` (the sync that waits on the backward tail) 29.3 → 0.8 ms. +- **Autocast Cast is the main new cost (12.8%).** 661 `CastKernel` launches / + 10.7 ms per step. The FP32→BF16 **weight** casts dominate: each MLP weight + (2048×8192) cast costs ~100–102 μs, *more* than the BF16 GEMM it feeds + (~73–89 μs). Caching the BF16 weight copy across fwd+bwd within a step would + recover most of this. +- **Adam optimizer is now the #1 cost.** `AdamAccumulateGradKernel` + (110 calls, 31.7 ms, 38.1%) is untouched by autocast (master weights / + optimizer stay FP32), so with GEMM 6.8× faster it becomes the largest single + contributor. +- **Launch/cast-bound.** GPU util fell 93% → ~68–73%, and `cudaLaunchKernel` + rose to ~3 526 per step (from 2 929) due to the extra casts (~27.9 ms + host-side). CUDA Graphs, or fusing Cast into the GEMM prologue, are the + natural next step. +- **`FillKernel`** is still launched ~678–710 times per step (~3.2 ms); + merging these into fewer larger fills would cut launch overhead. +- **Model load** (before Step_0) still transfers 5.99 GB HtoD (~0.82 s): the + checkpoint is FP32 and master weights stay FP32, so BF16 does not change this + one-shot cost. + +Steady-state throughput: **~2 260 tok/s** (256 tokens / ~113 ms) on a single +A100-SXM4-40GB at BF16 autocast — about **1.6×** the FP32 ~1 424 tok/s. + +For the full per-kernel timeline (methodology, layer-0 launch sequence, +kernel→architecture mapping, and the FP32→BF16 breakdown) see +[`nsys/out/llama3_kernel_timeline.md`](nsys/out/llama3_kernel_timeline.md). + ## 🗺 Roadmap - **2025/03/10** — InfiniTrain **v0.1.0** @@ -316,4 +498,4 @@ Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scal optimizations. Integrated a CTest + GTest based testing infrastructure to strengthen the - framework's automated test workflow. + framework's automated test workflow. diff --git a/docs/deepgemm_sm100_bf16_gemm_design.md b/docs/deepgemm_sm100_bf16_gemm_design.md new file mode 100644 index 000000000..1230c8c05 --- /dev/null +++ b/docs/deepgemm_sm100_bf16_gemm_design.md @@ -0,0 +1,1709 @@ +# DeepGEMM `sm100_bf16_gemm` Kernel Detailed Design + +面向 Blackwell（SM100，B200/GB200）的 BF16 GEMM 内核详细设计。本文覆盖：warp specialization 与 pipeline 组织、多级 tiling、SMEM/TMEM 排布、Tensor Core（`tcgen05.mma`）调用方式、生产者-消费者 mbarrier 同步协议、TMA 指令的发射与完成语义，以及若干容易被忽略但对正确性/性能关键的设计点。 + +## 0. Code Index + +| 层次 | 文件 | 职责 | +| --- | --- | --- | +| Device 主体 | [sm100_bf16_gemm.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh) | kernel 本体：SMEM/TMEM 布局、三个 warp 角色、流水推进 | +| UMMA 描述符 | [mma/sm100.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/mma/sm100.cuh) | `SmemDescriptor` 构造、SBO/LBO 推导、K 方向描述符推进 | +| tcgen05 PTX | [ptx/tcgen05.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh) | `tcgen05.mma.*` 内联汇编、`tcgen05.fence` | +| TMA load | [common/tma_copy.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/common/tma_copy.cuh) | swizzle-atom 循环、1SM/2SM/multicast 分支 | +| TMA PTX | [ptx/tma.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/ptx/tma.cuh) | `cp.async.bulk*`、`mbarrier.*`、tensormap 改写 | +| 调度器 | [scheduler/gemm.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/scheduler/gemm.cuh) | persistent 块分配、L2 swizzle、grouped/batched 索引 | +| Epilogue | [epilogue/sm100_store_cd.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd.cuh)
[epilogue/sm100_store_cd_swap_ab.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh) | TMEM→RF→SMEM→GMEM，swizzle 计算与 STSM 转置 | +| 通用工具 | [common/utils.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/common/utils.cuh) | `PatternVisitor`、编译期循环展开、TMEM 列对齐 | +| Host JIT | [impls/sm100_bf16_gemm.hpp](../third_party/DeepGEMM/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp) | 模板实参拼装、TMA descriptor 构造、launch | +| Host 启发式 | [heuristics/sm100.hpp](../third_party/DeepGEMM/csrc/jit_kernels/heuristics/sm100.hpp)
[heuristics/config.hpp](../third_party/DeepGEMM/csrc/jit_kernels/heuristics/config.hpp)
[heuristics/utils.hpp](../third_party/DeepGEMM/csrc/jit_kernels/heuristics/utils.hpp) | BLOCK_M/N/K、cluster、swizzle、stage 数、线程数推导 | + +> **注**：本仓库的 `third_party/DeepGEMM/third-party/cutlass` 子模块为空目录，CUTLASS 头文件未随仓库落地。文中涉及 `cute::UMMA::SmemDescriptor` 位域、`cutlass::arch::umma_arrive`、`cute::SM100_TMA_2SM_LOAD_2D` 等具体 PTX 文本时，均标注为「依据 CUTLASS 约定/由本仓库调用方式反推」，不做逐字断言。 + +--- + +## 1. 设计总览 + +### 1.1 一句话概括 + +一个 **persistent + warp-specialized + 全异步** 的三段流水线内核。数据通路为： + +``` +GMEM ──(cp.async.bulk.tensor / TMA)──► SMEM(A,B ring) + │ + ▼ + tcgen05.mma.cta_group::{1,2}.kind::f16 + │ + ▼ + TMEM(累加器双缓冲) + │ + ▼ tcgen05.ld (32dp32b*) + RF + │ + ▼ st.shared.v4 / stmatrix.trans + SMEM(C/D ring) + │ + ▼ cp.async.bulk.tensor store / reduce.add + GMEM(D) +``` + +从 GMEM 到 GMEM **没有任何一处让数据「停在通用寄存器里等一个 `__syncthreads()`」**：所有跨阶段依赖都由 `mbarrier` 的 parity 相位 + `tcgen05.commit` 表达，每个阶段都能独立超前推进。 + +### 1.2 框图（单个 CTA 内） + +``` + ┌──────────────────────────── 256 threads / 8 warps ────────────────────────────┐ + │ │ + warp0 (1 lane) │ warp1 (1 lane, leader CTA) warp2 warp4..7 (128 thr) │ + ┌──────────────┐ │ ┌────────────────────────┐ ┌──────────────┐ ┌───────────────────┐ │ + │ TMA LOAD │ │ │ MMA ISSUE │ │ TMEM ALLOC │ │ EPILOGUE │ │ + │ persistent │ │ │ persistent │ │ (prologue │ │ persistent │ │ + │ k-loop │ │ │ k-loop │ │ only) │ │ store-loop │ │ + └──────┬───────┘ │ └───┬──────────────▲─────┘ └──────────────┘ └────▲──────┬───────┘ │ + │ │ │ │ │ │ │ + wait │ empty[s] │ wait │ full[s] │ tcgen05.commit wait │ │ arrive │ + ▼ │ ▼ │ → empty[s] tmem_full ▼ tmem_empty│ + ┌───────────┐ │ ┌───────────┐ │ → tmem_full[a] (last k) [a] │ ┌───────────┐ │ + │ SMEM ring │─────┼─►│ tcgen05 │──────┼───────────────────────────────────────┼─►│ TMEM ring │ │ + │ A[s],B[s] │ │ │ .mma │ │ │ │ acc[a] │ │ + │ kNumStages│ │ └───────────┘ │ │ │ 2 stage │ │ + └───────────┘ │ │ │ └───────────┘ │ + └─────────────────────┴───────────────────────────────────────┴──────────────────┘ +``` + +warp3 以及 warp1 在 peer CTA 上的副本是**空转**的（见 §3.1）。注意这里的「空转」仅指 SIMT 线程层面：**grid 是 `kNumSMs` 个 persistent CTA，每个 SM 上都有一个自己的 lane 在发 UMMA，没有任何一个 SM 的 tensor core 被浪费**——完整推导见 §3.4。 + +### 1.3 关键设计选择 + +| 设计点 | 取值 | 理由 | +| --- | --- | --- | +| 编译方式 | 每个 (shape, config) 组合 JIT 生成一份全量模板特化的 `.cu` 编到 cubin | BLOCK_*、swizzle、stage 数、甚至 N/K 本身都成为编译期常量 → 内层 K 循环完全展开、描述符偏移常量折叠、tail-K 分支整段消除 | +| Grid | `gridDim.x = num_sms`，`__launch_bounds__(256, 1)` | persistent kernel，1 CTA/SM；配合近满额 SMEM 与独占 TMEM，物理上排除 2 CTA 共 SM | +| Cluster | 1 或 2（`cluster_m×cluster_n ≤ 2`） | 用 `cta_group::2` UMMA 把 M 做到 256，B 沿 N 对半切 → B 的 SMEM 占用与 L2 流量双双减半 | +| MMA 指令 | `tcgen05.mma.cta_group::{1,2}.kind::f16`，`UMMA_K = 16` | BF16 的 UMMA K 原子固定 16；A/B 均来自 SMEM（`_SS` 变体），累加器在 TMEM | +| `UMMA_M` | 恒为 `128 × kNumMulticast` | TMEM 的 datapath（行）固定 128，A/D 布局原子就是 128 行；`BLOCK_M ∈ {32,64}` 时也照发 M=128 的指令（见 §7.6） | +| 累加器缓冲 | TMEM 双缓冲（`kNumEpilogueStages = 2`） | 让第 i+1 块的 MMA 与第 i 块的 epilogue 重叠 | +| C/D 缓冲 | SMEM 双缓冲（`kNumTMAStoreStages = 2`） | 让 STSM 写与 TMA store 读重叠 | +| A/B 环 | 尽可能多的 stage（host 端按 SMEM 预算反解，上限 32） | 掩盖 HBM 延迟；stage ≥ 8 且 NT-Normal 时再做「stage 合并」把 `BLOCK_K` 放大（见 §7.5） | +| PDL | `cudaGridDependencySynchronize()` 放在 prologue **之后** | barrier 初始化、TMEM 分配、TMA descriptor prefetch 与前驱 kernel 的尾巴重叠 | + +--- + +## 2. Host 侧：JIT 特化与配置推导 + +Kernel 的全部行为由 26 个模板实参决定，它们在 [sm100_bf16_gemm.hpp](../third_party/DeepGEMM/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp) 的 `generate_impl()` 里被格式化成一段只包含 `__instantiate_kernel()` 的 `.cu` 源码，再交给 nvcc 编成 cubin。 + +### 2.1 模板实参来源 + +``` +kMajorA / kMajorB ← 由 a/b 的 stride 推断（get_major_type_ab） +SHAPE_M / SHAPE_N / SHAPE_K ← get_compiled_dim(dim, 'm'/'n'/'k', compiled_dims)： + 在 compiled_dims 里 → 填真实值；否则填 0（= 运行期参数） +BLOCK_M / BLOCK_N / BLOCK_K_ ← Layout +kNumGroups ← GemmDesc +kSwizzle{A,B,CD}Mode ← StorageConfig +kNumStages_ ← PipelineConfig.num_stages +kNumNonEpilogueThreads / kNumEpilogueThreads ← LaunchConfig（固定 128 / 128） +kNumMulticast ← Layout.get_cluster_size() +kIsMulticastOnA ← (Layout.cluster_n > 1) +kNumSMs ← LaunchConfig.num_sms（= gridDim.x） +kKAlignment ← heuristics_runtime->get_mk_alignment_for_contiguous_layout() +kSwapAB / kEnsureZeroPadding / kGemmType / kWithAccumulation / cd_dtype_t ← GemmDesc +kTensorCoreUtilControl ← device_runtime->get_tc_util()（默认 100） +``` + +Python 侧 `bf16_gemm_nt` 等 API 的 `compiled_dims` 默认是 `"nk"`（grouped contiguous/masked 也是 `"nk"`，两个 batched einsum 变体是 `"mn"`）。因此**典型场景下 `SHAPE_M == 0`、`SHAPE_N`/`SHAPE_K` 为编译期常量**，kernel 内第 117–119 行的覆写： + +```cpp +shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; +``` + +会把 N/K 变成常量，直接影响第 313 行： + +```cpp +constexpr bool kMayHaveTailKBlock = is_k_grouped_contiguous(kGemmType) + ? (kKAlignment % BLOCK_K != 0) + : (SHAPE_K == 0 or SHAPE_K % BLOCK_K != 0); +``` + +当 K 被编译进来且能被 `BLOCK_K` 整除时，`kMayHaveTailKBlock == false`，整个 tail-K 处理分支（`for_each_static_prefix` + `switch`）**不会生成任何 SASS**。这是 JIT 相比通用库最大的收益之一。 + +### 2.2 Layout 候选枚举 + +`SM100ArchSpec::get_layout_candidates()`： + +- **`block_k` 恒定**：`128 / element_size(BF16=2) = 64`。即 `BLOCK_K_ == 64`，对应 kernel 里的 `DG_STATIC_ASSERT(BLOCK_K_ == 64)`。物理含义：一个 swizzle atom 的 K 方向字节数固定为 128 B。 +- **m-grouped 三种类型强制走 swap-AB**：`block_n = 128`，`block_m = get_mk_alignment_for_contiguous_layout()`（SM100 上按 `expected_m` 在 `[32, 224]` 内以 32 为步长收缩），`cluster_m = 1`，`cluster_n = 2`（当 `ceil(n/128)` 与 `num_sms` 均为偶数）。 +- **其余类型全枚举** `swap_ab × block_m × block_n × cluster_m × cluster_n`，逐条过滤： + - `swap_ab == 1 && cluster_m > 1` → 跳过；`swap_ab == 0 && cluster_n > 1` → 跳过（只支持 layout A/D 方向的 cluster）。 + - `cluster_m * cluster_n > 2`、`num_sms % cluster_size != 0` → 跳过。 + - 非 swap 时 `block_m ∈ {32, 64, 128}`，按 `desc.m ≤ 32 / ≤ 64 / else` 三选一（注释：*smaller block M can avoid TMA L2 OOB bound*）。 + - 非 swap 时 `block_n` 上界：`desc.k <= 256 ? 128 : 256`（注释：*For small K, fewer store blocks improve store/compute overlap*），步长 `lcm(32, block_n_multiple_of)`，另加一个 16 的候选。 + - MN-major 时要求 `(block_m/cluster_n) % 64 == 0`（swizzle 对齐）；K-major 时只要求 `% 8 == 0`。 + - `ceil_div(desc.m, block_m) % cluster_m != 0` 或 `ceil_div(desc.n, block_n) % cluster_n != 0` → 跳过（保证 cluster 不跨边界）。 + - `swap_ab && block_n != 128` → 跳过（`LAYOUT_AD_M` 必须是 128）。 + - **TMEM 容量**：`2 * umma_n + tmem_sf_cols > 512` → 跳过（BF16 无 SF，`tmem_sf_cols = 0`，故 `umma_n ≤ 256`）。 + - 当 A 或 B 至少有一个是 K-major 时，要求算出的 `swizzle_a_mode == swizzle_b_mode == 128`（注释：*32B swizzle yields poor performance*）。 + +打分 `compare()` 的优先级序列：**单 wave 最优 → cluster 大者优 → wave 数少者优 → 末 wave 利用率高者优 → `block_m + block_n` 小者优（= stage 更多） → `block_m × block_n` 小者优**。 + +### 2.3 StorageConfig + +```cpp +load_block_m = block_m / cluster_n; +load_block_n = block_n / cluster_m; +store_block_m = swap_ab ? 16 /* umma_step_n */ : min(128 /* layout_ad_m */, block_m); +store_block_n = block_n; + +swizzle_mode_a = get_swizzle_mode(major_a == K ? block_k : load_block_m, sizeof(a_dtype)); +swizzle_mode_b = get_swizzle_mode(major_b == K ? block_k : load_block_n, sizeof(b_dtype)); +swizzle_mode_cd = get_swizzle_mode(store_block_n, sizeof(cd_dtype)); +``` + +`get_swizzle_mode()` 从 `{128, 64, 32, 16}` 里挑第一个能整除 `block_size * elem_size` 的值。BF16 + `block_k = 64` → 128 B，恒为 `swizzle = 128`。 + +注意 host 的 `store_block_n = block_n` 与 kernel 的 `STORE_BLOCK_N` 并不相等：kernel 里非 swap 分支是 `kSwizzleCDMode / sizeof(cd_dtype_t)`（即 TMA box 的内维被压到一个 swizzle atom 宽）。host 的 `store_block_n` 只用于算 swizzle mode，`make_tma_2d_desc()` 里又会被 `smem_inner_dim = swizzle_mode / elem_size` 覆盖掉。两者最终一致，但读代码时容易误判。 + +### 2.4 PipelineConfig（SMEM 预算 → stage 数） + +```cpp +constexpr int smem_capacity = 232448; // 227 KB +int smem_cd = swap_ab ? store_block_m * store_block_n * elemsize * 2 + : store_block_m * swizzle_cd_mode * 2; // × 2 = 双缓冲 +int smem_barriers = 32 * 8 * 3 + 2 * 8 * 2 + 8; // = 808 B，按 kNumMaxStages=32 预留 +int smem_tmem_ptr = 4; +int smem_a_per_stage = load_block_m * block_k * elemsize_a; +int smem_b_per_stage = load_block_n * block_k * elemsize_b; + +num_stages = min((smem_capacity - (smem_cd + smem_barriers + smem_tmem_ptr)) + / (smem_a_per_stage + smem_b_per_stage), 32); +smem_size = smem_extra + num_stages * smem_per_stage; +``` + +`smem_barriers` 的三项与 device 侧布局精确对应（见 §5.2）：`32*8*3` 是每 stage 三组 barrier（full / empty / **with-SF full**），`2*8*2` 是 `tmem_full[2] + tmem_empty[2]`，`+8` 是 `tensor_core_full_barrier`。BF16 没有 scale factor，第三组永不使用，但仍按 1D1D（FP8/FP4）kernel 的约定占位——所以 host 的预留量对 `kNumStages ≤ 32` 恰好是**紧上界**。 + +### 2.5 LaunchConfig + +```cpp +return { desc.num_sms, layout.get_cluster_size(), 256, 32, 128, 128, 128 }; +// num_sms num_sms_per_cluster num_threads tma math non_epi epi +``` + +只有 `num_threads = 256`、`num_non_epilogue_threads = 128`、`num_epilogue_threads = 128` 会进模板实参；`num_tma_threads = 32` / `num_math_threads = 128` 是给 SM90 kernel 用的字段，SM100 路径忽略。 + +### 2.6 TMA descriptor + +三个 descriptor 都以 `__grid_constant__ cute::TmaDescriptor` 按值传参（128 B 常量内存，避免走 GMEM）。构造见 [runtime_utils.hpp](../third_party/DeepGEMM/csrc/jit_kernels/impls/runtime_utils.hpp)： + +| descriptor | gmem (inner, outer) | smem box (inner, outer) | 备注 | +| --- | --- | --- | --- | +| A | K-major: `(k, m*G)`；MN-major: `(m*G, k)` | K-major: `(block_k→64, block_m)`；MN-major: `(block_m→64, block_k)` | `num_groups > 1` 时强制 K-major；box 内维被 swizzle 覆写为 `swizzle/elem = 64` | +| B | K-major: `(k, n)`；MN-major: `(n, k)` | 同上，`block_n` 换 `load_block_n` | `num_groups` 只作用在外维：`gmem_outer_dim * num_groups` | +| C/D | `(n, m*G)` | `(store_block_n→swizzle/elem, store_block_m)` | D 必须 N-major | + +公共属性：`CU_TENSOR_MAP_INTERLEAVE_NONE`、`CU_TENSOR_MAP_L2_PROMOTION_L2_256B`、`CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE`（越界元素零填充——这是 tail-K 与 M/N 非整除时结果仍正确的硬件保证）、swizzle 由 `mode_into_tensor_map_swizzle()` 映射到 `CU_TENSOR_MAP_SWIZZLE_{NONE,32B,64B,128B}`。 + +Batched（`sm100_bf16_bhr_hdr_bhd` / `bhd_hdr_bhr`）走 `make_tma_3d_desc()`，第三维是 head，device 侧 `kIsBatchedMM` 打开 `SM90_TMA_LOAD_3D` / `SM100_TMA_2SM_LOAD_3D` 分支。 + +### 2.7 Launch 属性 + +`construct_launch_config()`（[handle.hpp](../third_party/DeepGEMM/csrc/jit/handle.hpp)）： + +1. `cuFuncSetAttribute(CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size)` —— 动态 SMEM 远超 48 KB 静态上限，必须显式抬。 +2. `cluster_dim > 1` → `CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION = {cluster_dim, 1, 1}`。 +3. `enable_pdl` → `CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION = 1`。默认 `DeviceRuntime::enable_pdl = false`，需 `deep_gemm.set_pdl(True)` 打开。 +4. `cuLaunchKernelEx` 发射。 + +编译 flag：`--gpu-architecture=sm_100f`（nvcc ≥ 12.9 用 arch-family 后缀 `f`，否则 `100a`）、`-O3 --expt-relaxed-constexpr --expt-extended-lambda`，产物直接是 `-cubin`。kernel 内 `#if __CUDA_ARCH__ >= 1000` 之外的编译分支只留一个 `DG_DEVICE_ASSERT(false and "This kernel only support sm_100f")`。 + +--- + +## 3. 线程组织与 warp 角色 + +### 3.1 角色表（256 threads = 8 warps） + +| warp | 条件 | 角色 | 实际活跃线程 | +| --- | --- | --- | --- | +| 0 | `warp_idx == 0 && elect_one_sync()` | **TMA load**：persistent 遍历所有块，每块跑完整 K 循环，发 A/B 的 TMA | **1 lane**（TMA 是单线程指令） | +| 1 | `warp_idx == 1 && elect_one_sync()` | barrier 初始化（prologue） | 1 lane | +| 1 | `warp_idx == 1 && is_leader_cta` | **MMA issue**：persistent 遍历，每块跑完整 K 循环，发 `tcgen05.mma` 与 `tcgen05.commit` | 发射 UMMA 时 1 lane，`commit` 时整 warp | +| 2 | `warp_idx == 2` | **TMEM 分配**（`tcgen05.alloc`，prologue 一次性） | 整 warp（`.sync.aligned` 要求） | +| 3 | — | **空转** | 0 | +| 4 … 4+`kNumUMMAStoreThreads/32`-1 | `warp_idx >= kNumNonEpilogueThreads/32 && < (kNumNonEpilogueThreads+kNumUMMAStoreThreads)/32` | **Epilogue**：TMEM→RF→SMEM→TMA store | `kNumUMMAStoreThreads` 个线程 | +| 其余 | — | 空转 | 0 | + +`kNumUMMAStoreThreads` 的取值决定了 epilogue 用几个 warp： + +| 场景 | `STORE_BLOCK_M` | `kNumUMMAStoreThreads` | 参与 warp | +| --- | --- | --- | --- | +| 非 swap，`BLOCK_M = 128` | 128 | 128 | w4–w7（满 warpgroup） | +| 非 swap，`BLOCK_M = 64` | 64 | 64 | w4–w5 | +| 非 swap，`BLOCK_M = 32` | 32 | 32 | w4 | +| swap-AB | 16 | `kNumEpilogueThreads = 128` | w4–w7 | + +非 swap 时 `kNumUMMAStoreThreads = STORE_BLOCK_M` 的原因：**TMEM 的一个 datapath（行）由一个线程负责**，读 `STORE_BLOCK_M` 行就需要 `STORE_BLOCK_M` 个线程。swap-AB 时 `STORE_BLOCK_N == 128`（静态断言强制），要覆盖全部 128 个 TMEM 行，因此必须是完整 warpgroup。 + +Peer CTA（`block_rank_in_cluster() != 0`）上的 warp 1 什么都不做——`tcgen05.mma.cta_group::2` 只由 CTA-pair 的 leader 发射，硬件自己从两个 CTA 的 SMEM 取操作数、往两个 CTA 的 TMEM 写结果。但这**不**意味着 peer CTA 的 tensor core 闲置，见 §3.4.4。 + +### 3.2 为什么必须 1 CTA/SM + +三条互相加强的约束： + +1. `__launch_bounds__(256, 1)` 显式声明每 SM 最多 1 个 block。 +2. SMEM 占用接近 227 KB 上限（§2.4 就是按「装满」反解 stage 数的）。 +3. Epilogue 开头的 `DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(tmem_ptr_in_smem) == 0)`（第 405 行）——断言 `tcgen05.alloc` 返回的基址列号为 0，即本 CTA 拿到了 TMEM 的第 0 列起的全部 `kNumTmemCols`。源码注释写明：*we also forbid two CTAs to share the same SM and its tensor memory*。 + +第 3 条同时是一个**功能简化**：因为基址恒为 0，代码里所有 TMEM 地址都可以直接写成 `accum_stage_idx * UMMA_N + 列内偏移`，不需要把分配返回的 base 加进去；行号（datapath）也只由 warp/lane 隐式决定，注释指出*hardware will ignore the warp index bits, i.e., no need for `tmem_ptr |= (epilogue_warp_idx * 32) << 16`*。 + +### 3.3 `elect_one_sync()` 的使用纪律 + +`cute::elect_one_sync()` 底层是 `elect.sync`，**要求整 warp 收敛执行**。本 kernel 所有调用都写成 `if (warp_idx == K and cute::elect_one_sync())` 或在已经 warp-uniform 的分支内部，`warp_idx` 对同一 warp 的 32 lane 恒等，因此短路求值不会造成 `elect.sync` 的部分参与。同理，第 324 行发射 UMMA 的 `if (cute::elect_one_sync())` 位于 `full_barriers[...]->wait()` 与 `__shfl_sync(0xffffffff, ...)` 之后——`__shfl_sync` 用满 mask，也隐含要求收敛。 + +### 3.4 为什么「1 个 lane 发射 UMMA」不等于算力浪费 + +§3.1 的角色表里，MMA issue 一栏写的是「发射 UMMA 时 1 lane」。这个描述的**作用域是单个 CTA 内部**，很容易被误读成「整个 GPU 只有一个线程在驱动 tensor core，其他 SM 的算力全浪费了」。本节把这件事彻底说清。 + +#### 3.4.1 层次：grid 上是 `kNumSMs` 个发射者，不是 1 个 + +``` +gridDim.x = kNumSMs（B200 = 148） ← LaunchConfig 第一项，persistent kernel +__launch_bounds__(256, 1) ← 每 SM 恰好 1 个 CTA（§3.2） + └─► 每个 CTA 都有自己的 warp 1 + └─► 每个 warp 1 都有 1 个 elected lane 在发 tcgen05.mma + ⇒ 整个 GPU 同时有 148 个 lane 在并行发射 MMA +``` + +调度器的 `next_block_idx = (++current_iter) * kNumSMs + blockIdx.x` 保证这 148 个 CTA 各自负责互不重叠的块序列（§4.1），每个 SM 的 tensor core 都有活干。**不存在「其他 SM 被浪费」的情形**。 + +#### 3.4.2 单 CTA 内为什么 1 个 lane 就够：SM90 → SM100 的架构跃迁 + +这不是 DeepGEMM 的取巧，而是 Blackwell tensor core 编程模型的根本变化： + +| | SM90 `wgmma.mma_async` | SM100 `tcgen05.mma` | +| --- | --- | --- | +| 发射宽度 | 整个 warpgroup（128 线程）必须协同 | **单线程** | +| A 操作数 | 可来自寄存器（每线程持有分片） | 只来自 SMEM descriptor | +| B 操作数 | SMEM descriptor | SMEM descriptor | +| D 累加器 | **分散在 128 个线程的寄存器里** | **TMEM**（独立于寄存器文件的 256 KB 存储） | +| 为何需要那么多线程 | 累加器每个元素都得有个线程「拿着」 | 没人需要「拿着」任何东西 | + +SM90 的 128 线程是被**累加器的存储位置**逼出来的，与算力无关。SM100 把累加器搬进 TMEM 之后这个约束消失，UMMA 发射退化成「一个线程往硬件异步队列里投递一条描述符」——它是一条**控制指令**，不是一条需要 32/128 个线程各自贡献数据分片的 SIMT 指令。 + +本 kernel 用的 `_SS` 后缀（`SM100_MMA_F16BF16_SS`）正是这个意思：A 和 B **都**来自 SMEM。三个操作数（A、B、D）没有一个住在寄存器里，自然没有任何理由让多个线程参与发射。 + +#### 3.4.3 量化：发射端有 ≥ 40 倍余裕 + +kernel 自己在 tensor core 利用率控制里给出了执行周期公式（第 382 行）： + +```cpp +constexpr static uint64_t kNumUMMACycles = (2ull * UMMA_M * UMMA_N * BLOCK_K) / 8192ull; +``` + +取典型配置 `UMMA_M = 128`、`UMMA_N = 256`、`BLOCK_K = 64`： + +| 量 | 值 | 来源 | +| --- | --- | --- | +| 单条 UMMA 覆盖的 FLOP | `2 × 128 × 256 × 16` = **1,048,576** | `UMMA_K = 16`（BF16 固定） | +| 一个 stage 发射的 UMMA 条数 | `BLOCK_K / UMMA_K` = **4** | `issue_full_k_block`（第 341 行） | +| 一个 stage 的总 FLOP | `2 × 128 × 256 × 64` = **4,194,304** | — | +| 执行所需 tensor-core cycle | 4,194,304 / 8192 = **512** | 上式 | +| 发射所需 lane-cycle | 4 条 UMMA + 8 条描述符推进 IADD ≈ **12** | `issue_umma` 内两次 `advance_umma_desc_lo` | +| **发射 : 执行** | ≈ **1 : 43**（纯 UMMA 算则 1 : 128） | — | + +`8192` 不是随意取的魔数，它就是 **SM100 每 SM 每 cycle 的 BF16 FLOP 峰值**：B200 约 2.25 PFLOPS ÷ 148 SM ÷ ~1.83 GHz ≈ 8310，取整到 8192。因此 `kNumUMMACycles` 的物理含义是「这个 stage 的 UMMA 在满速 tensor core 上要跑多少 cycle」。 + +**40 倍以上的比值意味着发射端从来不是瓶颈。** 一个 elected lane 每 512 个 cycle 只需忙约 12 个 cycle，其余时间都卡在 `full_barriers[...]->wait()` 上睡觉。 + +最有力的反证是 `kTensorCoreUtilControl` 这个旋钮的存在（§9.7）：作者需要**主动插入 `clock64()` 自旋**才能把 tensor core 拖慢、降低功耗掉频的可能性（第 372–387 行）。如果发射能力不足，这个功能没有任何意义。 + +#### 3.4.4 2-CTA 模式：控制流集中，执行分布 + +`cta_group::2` 时，peer CTA 的 warp 1 完全不执行 MMA 分支，但**它的 tensor core 一点没闲着**：leader CTA 的 1 个 lane 发出的那一条指令，会同时驱动 CTA-pair 两个 SM 的 tensor core——从两个 CTA 的 SMEM 各取一半操作数，往两个 CTA 的 TMEM 各写一半结果，把 `UMMA_M` 拼到 256（§7.2）。 + +这是 SM100 UMMA 最反直觉的地方：**指令流宽度与算力宽度彻底解耦**。1 个 lane 的指令流对应 2 个 SM 的全部 tensor core。 + +#### 3.4.5 真正会让 tensor core 空闲的因素 + +既然发射不是瓶颈，tensor core 停转只可能来自三条依赖边。kernel 的全部复杂度都花在让它们永不成为关键路径上： + +| 等待点 | 物理含义 | 掩盖手段 | +| --- | --- | --- | +| `full_barriers[stage_idx]->wait(phase)` | SMEM 里还没有 A/B（TMA 未返回） | `kNumStages` 深的 A/B 环（最多 32，§2.4） | +| `tmem_empty_barriers[a]->wait(...)` | TMEM 双缓冲都被占（epilogue 未搬完） | `kNumEpilogueStages = 2` + TMEM 早释放（§11.4） | +| 输出侧 TMA store 带宽 | D 写不回 GMEM | `kNumTMAStoreStages = 2` + `wait_group.read`（§11.5） | + +换言之：优化方向是「加深环、提前释放、提高 L2 命中」，而不是「多找几个线程来发射 MMA」。后者在 SM100 上已经没有任何收益空间。 + +#### 3.4.6 观测陷阱：不要用 warp 活跃度评估本 kernel + +256 个线程里稳态活跃的只有约 34 个（TMA 1 lane + MMA 1 lane + epilogue 32~128），warp 3 恒空转（详见 §13.3 第 8 条）。这说的是 **SIMT 通路的活跃度低**，不是 tensor core 闲置。用 nsys / ncu 观察时： + +- ❌ `sm__warps_active`、issue-slot 利用率、`smsp__inst_executed`——会得出「这个 kernel 效率极低」的**完全错误**结论。 +- ✅ `sm__pipe_tensor_cycles_active`（tensor core 管线活跃 cycle）、TMA 的 L2 吞吐、`sm__mio_inst_issued`（TMEM 读写）——才是本 kernel 的真实健康指标。 + +Warp specialization 的设计目标本来就是把通用 SIMT 通路腾空，让数据搬运全部交给 TMA / tensor core / TMEM 这些专用引擎。用 SIMT 时代的指标去衡量它，等于用「多少工人在挥铲子」去评价一台挖掘机。 + +--- + +## 4. Persistent 调度器 + +### 4.1 复制式状态机，不是共享工作队列 + +第 180 行在**角色分派之前**构造 `scheduler`： + +```cpp +auto scheduler = sched::Scheduler( + shape_m, shape_n, shape_k, grouped_layout); +``` + +它是个**寄存器里的值对象**，每个线程各持一份私有副本。三个角色 warp 各自独立调用 `get_next_block()`，靠 `next_block_idx = (++current_iter) * kNumSMs + blockIdx.x` 这一条纯算术式子得到**完全相同**的块序列： + +``` +iter 0 → blockIdx.x +iter 1 → blockIdx.x + kNumSMs +iter 2 → blockIdx.x + 2*kNumSMs +... +``` + +因此： + +- **零原子操作、零全局 ticket 计数器**，调度开销是几条整数指令； +- 三个角色天然锁步，无需为「现在在处理哪一块」建立任何额外通信； +- `scheduler.current_iter` 直接充当**全局逻辑时钟**，MMA warp 与 epilogue warp 各自用它算出累加器缓冲的 stage/phase（§6.3），这就是两条流水之间唯一的「隐式」耦合。 + +代价是负载不能动态窃取：尾波（tail wave）的空闲 CTA 只能干等。启发式打分里的 `last_wave_util` 就是在 host 侧提前把这件事量化。 + +### 4.2 L2 swizzle 分组 + +`get_swizzled_block_idx()` 把线性的 `block_idx` 重映射成 (m, n)，目的是让**同时在飞的 kNumSMs 个 CTA 尽量共享 L2 里的 A/B**。 + +```cpp +kNum1DBlocksPerGroup = get_num_1d_blocks_per_group<...>(); // 编译期，∈ {8, 16} +primary_num_blocks = kIsMulticastOnA ? num_n_blocks : num_m_blocks; +secondary_num_blocks = kIsMulticastOnA ? num_m_blocks : num_n_blocks; +num_blocks_per_group = secondary_num_blocks * kNum1DBlocksPerGroup; +group_idx = block_idx / num_blocks_per_group; +first_block_idx = group_idx * kNum1DBlocksPerGroup; +in_group_idx = block_idx % num_blocks_per_group; +num_blocks_in_group = min(kNum1DBlocksPerGroup, primary_num_blocks - first_block_idx); + +// kIsMulticastOnA == false（在 M 上分组，组内 M 变化最快） +m_block_idx = first_block_idx + in_group_idx % num_blocks_in_group; +n_block_idx = in_group_idx / num_blocks_in_group; +``` + +组大小的选择是**最小化 L2 工作集**： + +```cpp +usage = kIsMulticastOnA ? candidate * BLOCK_N + ceil_div(kNumSMs, candidate) * BLOCK_M // 在 N 上分组 + : candidate * BLOCK_M + ceil_div(kNumSMs, candidate) * BLOCK_N; // 在 M 上分组 +// candidate ∈ {8, 16}，取 usage 最小者 +``` + +`DG_STATIC_ASSERT(kNum1DBlocksPerGroup % kNumMulticast == 0)` 保证一个 cluster 的两个 CTA 不会跨组边界——否则它们拿到的 B tile 就不同了，2-CTA UMMA 会算错。 + +`#if __CUDA_ARCH__ < 1000` 分支里那段「修正不对齐的 TMA multicast」只对 SM90 生效，注释说明 SM100 的 2-CTA 模式**不能动态关闭**，因此 SM100 靠 host 侧的整除性过滤（§2.2）来保证。 + +### 4.3 GemmType 变体 + +| GemmType | `num_blocks` | 额外状态 | 说明 | +| --- | --- | --- | --- | +| `Normal` | `num_m_blocks * num_n_blocks` | — | `get_global_idx` 退化为 `block_idx * block_size` | +| `Batched` | 同上，`× kNumGroups` | `current_group_idx` 作为 batch_idx | 不走 swizzle，按 `kIsMulticastOnA` 决定 m/n 谁变化快；TMA 走 3D | +| `MGroupedContiguous` | 同上 | `grouped_layout[m]` = 每行所属 group | B 的外维加 `group * shape_dim` 偏移 | +| `MGroupedMasked` | 逐 group 累加 | `current_m_cumsum` | 边扫边把 `next_block_idx` 落到对应 group，`num_m_blocks` 每 group 重算 | +| `MGroupedContiguousWithPsumLayout` | 逐 group 累加 | `last_psum_m` / `current_psum_m` / `current_m_block_cumsum` | group 边界按 psum 偏移切分，`m_block_idx += last_psum_m / BLOCK_M` | +| `KGroupedContiguous{,WithPsumLayout}` | 同上 | `current_shape_k` / `current_k_cumsum` / `current_k_start,end` | 每个 group 的 K 长度不同 → `num_total_k_blocks` 逐块变化；要求 A/B 都是 MN-major | + +对本 kernel 最关键的一条约束在第 216 行： + +```cpp +DG_STATIC_ASSERT(kGemmType == Normal or is_k_grouped_contiguous(kGemmType) or kGemmType == Batched or + kMajorA == cute::UMMA::Major::K, "Invalid major"); +``` + +即所有 m-grouped 变体的 A 必须 K-major（注释：*for all m-grouped GEMMs, A must be K-majored*），因为 group 偏移是加在外维上的。 + +### 4.4 跨块连续的流水线状态 + +第 184–191 行： + +```cpp +uint32_t stage_idx = 0, phase = 0, tensor_core_phase = 0; +auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + stage_idx = (stage_idx + 1) % kNumStages; + phase ^= stage_idx == 0; // 只在回绕到 stage 0 时翻转相位 +}; +``` + +`stage_idx`/`phase` 声明在**块循环之外**，`tma_stage_idx` 也在 epilogue 里注释为 *Share store pipeline between blocks*。这意味着 A/B 环与 C/D 环**跨输出块连续运转**：TMA warp 可以在 MMA warp 还在算第 i 块最后一个 k_block 时，就开始往刚被释放的 stage 里灌第 i+1 块的 k=0 数据。块边界上没有任何「排空-重启」的开销，这是 persistent kernel 相较 grid-per-tile 实现的核心优势。 + +`advance_pipeline` 同时被 TMA warp（第 203 行）和 MMA warp（第 314 行）作为 `for` 的递增表达式使用，两边独立维护但推进规则一致，因此 `stage_idx`/`phase` 序列天然对齐。 + +--- + +## 5. 共享内存布局 + +### 5.1 线性布局 + +```cpp +extern __shared__ __align__(1024) uint8_t smem_buffer[]; // 1024 B 对齐，服务于 swizzle-128B +``` + +`utils::PatternVisitor` 是个零开销的「下标 → 指针」闭包包装器（`operator[](i)` 直接调 lambda），用它替代指针数组，避免在 SMEM/寄存器里存 stage 指针表： + +| 区段 | 起址 | 大小 | 访问器 | +| --- | --- | --- | --- | +| C/D store ring | `smem_buffer + 0` | `SMEM_CD_SIZE = STORE_BLOCK_M * STORE_BLOCK_N * sizeof(cd) * kNumTMAStoreStages` | `smem_cd[i]` | +| A ring | `smem_buffer + SMEM_CD_SIZE` | `kNumStages * SMEM_A_SIZE_PER_STAGE`，`SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * 2` | `smem_a[i]` | +| B ring | `+ kNumStages * SMEM_A_SIZE_PER_STAGE` | `kNumStages * SMEM_B_SIZE_PER_STAGE`，`SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * 2` | `smem_b[i]` | +| Barriers | `+ kNumStages * SMEM_B_SIZE_PER_STAGE` | 见 §5.2 | `full_barriers[i]` 等 | +| TMEM 基址 | barriers 之后 | 4 B | `tmem_ptr_in_smem` | + +三个 `DG_STATIC_ASSERT(... % 1024 == 0)` 保证每个区段起点都 1024 B 对齐——这是 swizzle-128B 的硬件要求（一个 swizzle atom 是 8 行 × 128 B = 1 KB，若基址不按 1 KB 对齐，TMA 写入的 swizzle 图案与 UMMA 描述符解读的图案会错位）。 + +### 5.2 Barrier 区（含一段历史包袱） + +以 `Barrier`（= `cutlass::arch::ClusterTransactionBarrier`，8 B）为单位，`barrier_start_ptr` 起： + +| 索引区间 | 名称 | 个数 | `init()` 计数 | +| --- | --- | --- | --- | +| `[0, S)` | `full_barriers` | `kNumStages` | `kNumMulticast` | +| `[S, 2S)` | `empty_barriers` | `kNumStages` | `1` | +| `[2S, 2S+2)` | `tmem_full_barriers` | 2 | `1` | +| `[2S+2, 2S+4)` | `tmem_empty_barriers` | 2 | `kNumMulticast * kNumUMMAStoreThreads` | +| `[2S+4, 3S+4)` | **（空洞，未使用）** | `S` | — | +| `3S+4` | `tensor_core_full_barrier` | 1 | `1`（仅当 `kTensorCoreUtilControl < 100`） | +| `3S+5`（字节偏移 `+4`） | `tmem_ptr_in_smem` | 4 B | — | + +（`S = kNumStages`） + +那个 `S` 大小的空洞来自 `tensor_core_full_barrier = barrier_start_ptr + kNumStages * 3 + kNumEpilogueStages * 2`：索引算术按「每 stage **三**组 barrier」排布，第三组是 FP8/FP4 1D1D kernel 的 *with-SF full barriers*。BF16 没有 scale factor，于是这 `kNumStages * 8` 字节被跳过但**不回收**。host 侧 `smem_barriers = 32*8*3 + 2*8*2 + 8` 与之精确呼应，所以整体仍是紧的。 + +初始化由 warp 1 的单个 lane 完成，随后： + +```cpp +cutlass::arch::fence_barrier_init(); // fence.mbarrier_init.release.cluster +``` + +注释写明目的：*Make initialized barrier visible in async proxy*。mbarrier 会被 TMA/UMMA 这些**异步代理**访问，普通的 `__syncthreads()` 不足以建立 generic proxy → async proxy 的可见性，必须用这条 cluster 作用域的 release fence。之后第 172 行做 `cluster_sync_with_relaxed_arrive()`（2-CTA）或 `__syncthreads()`（1-CTA），确保 peer CTA 也能看到 leader 的 barrier 状态、以及 warp 2 写入的 `tmem_ptr_in_smem`。 + +### 5.3 A/B stage 内部排布 + +**K-major（最常见）**：TMA box 是 `(inner = BLOCK_K = 64 elem = 128 B, outer = LOAD_BLOCK_M)`，swizzle atom = 8 行 × 128 B。一个 stage 就是 `LOAD_BLOCK_M/8` 个 atom 沿 M 方向线性堆叠： + +``` +smem_a[s] (LOAD_BLOCK_M=128, BLOCK_K=64, bf16 → 16 KB) +┌──────────────────────── atom 0 : rows 0.. 7 ────────────────────────┐ ← SBO = 1024 B +│ row r: 128 B = 8 个 16-B bank group，物理位置 g' = g ^ (r % 8) │ +├──────────────────────── atom 1 : rows 8..15 ────────────────────────┤ +│ ... │ +├─────────────────────── atom 15 : rows 120..127 ──────────────────────┤ +└───────────────────────────────────────────────────────────────────────┘ +``` + +`^ (r % 8)` 的 bank-group 置换就是 `CU_TENSOR_MAP_SWIZZLE_128B` / `cute::UMMA::LayoutType::SWIZZLE_128B` 定义的图案，由 TMA 硬件在写入时施加、由 UMMA 硬件在读取时反解，软件两侧都不需要参与。它的作用是把「同一列的 8 个元素」打散到 8 个不同的 bank group，消除 tensor core 按列取数时的 SMEM bank conflict。 + +**MN-major**：TMA box 变成 `(inner = LOAD_BLOCK_MN, outer = BLOCK_K)`，`BLOCK_INNER_ATOM = swizzle/elem = 64`，于是 `tma::copy` 内部会循环 `LOAD_BLOCK_MN / 64` 次，每次的目的地址是 `smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM`。即 SMEM 布局是「**K 外、MN-atom 内**」：先把第 0 个 64 宽的 MN atom 的全部 `BLOCK_K` 行放完，再放第 1 个 atom。UMMA 描述符的 `stride_k` 也随之从 K-major 的 `1` 变成 `get_inner_block_atom_size<...>()`（§9.4）。 + +### 5.4 C/D stage 内部排布 + +非 swap：一个 stage 是 `STORE_BLOCK_M` 行 × `kSwizzleCDMode`（128 B），即 `STORE_BLOCK_M` 个 128-B 行；行内同样是 8 个 16-B bank group 的 `^ (r % 8)` 置换。每个 stage 只覆盖 `STORE_BLOCK_N = 128/sizeof(cd)` 个 N 元素（bf16 → 64，fp32 → 32），正好是一个 swizzle atom 宽，所以一个 stage 对应**一条** TMA store。 + +swap-AB：一个 stage 是 `STORE_BLOCK_M(16) × STORE_BLOCK_N(128)`，被切成 `STORE_BLOCK_N / STORE_BLOCK_N_ATOM = 2` 个 atom，每 atom `16 × 128 B`；4 个 warp 两两负责一个 atom（§11.2）。 + +### 5.5 UMMA 越界读的静态防护 + +第 93–94 行： + +```cpp +static constexpr uint32_t UMMA_A_SIZE_PER_STAGE = + math::constexpr_align(LOAD_BLOCK_M, LAYOUT_AD_M) * BLOCK_K * sizeof(nv_bfloat16); +DG_STATIC_ASSERT(UMMA_A_SIZE_PER_STAGE <= SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE * kNumStages, + "Memory out of bound for UMMA"); +``` + +因为 `UMMA_M` 恒为 128（或 256），当 `BLOCK_M = 32/64` 时 `LOAD_BLOCK_M < 128`，UMMA 仍会**按 128 行去读 A**，越过的部分落在后续 A stage、乃至 B ring 上。那些行算出来的 D 在 epilogue 里根本不会被读（只读 `STORE_BLOCK_M = BLOCK_M` 行），所以读到垃圾无害——但**必须落在本 CTA 已分配的动态 SMEM 内**，否则是非法访问。化简后（两边同除 `BLOCK_K * 2`）该断言等价于： + +``` +128 ≤ LOAD_BLOCK_M + LOAD_BLOCK_N * kNumStages +``` + +### 5.6 Worked Example：8192 × 8192 × 8192，BF16→BF16，B200（148 SM） + +启发式选出的配置：`swap_ab=0, BLOCK_M=128, BLOCK_N=256, BLOCK_K=64, cluster=(2,1)`，`swizzle A/B/CD = 128`，`num_stages = 6`。 + +推导链： + +``` +kNumMulticast = 2 (cluster_m=2, cluster_n=1) +kIsMulticastOnA = false (cluster_n == 1) +LOAD_BLOCK_M = 128 / 1 = 128 +LOAD_BLOCK_N = 256 / 2 = 128 +UMMA_M = 128 * 2 = 256 +UMMA_N = BLOCK_N = 256 UMMA_K = 16 +STORE_BLOCK_M = min(128,128) = 128 +STORE_BLOCK_N = 128 / 2 = 64 kNumUMMAStoreThreads = 128 +kDoMergeStages = false (num_stages = 6 < 8) +kMayHaveTailKBlock= false (SHAPE_K = 8192 编译期常量，8192 % 64 == 0) +kNumAccumTmemCols = 2 * 256 = 512 → kNumTmemCols = 512 （恰好占满 TMEM） +kNum1DBlocksPerGroup: candidate 8 → 8*128 + 19*256 = 5888 + candidate 16 → 16*128 + 10*256 = 4608 ⇒ 取 16 +num_blocks = 64 * 32 = 2048, num_waves = ceil(2048/148) = 14, last_wave_util = 124 +``` + +SMEM 字节表（单 CTA）： + +| 偏移 | 大小 | 内容 | +| --- | --- | --- | +| 0 | 32 768 | C/D ring：2 × (128 行 × 128 B) | +| 32 768 | 98 304 | A ring：6 × (128 × 64 × 2 B = 16 KB) | +| 131 072 | 98 304 | B ring：6 × (128 × 64 × 2 B = 16 KB) | +| 229 376 | 184 | barriers：`3*6 + 5 = 23` 个 Barrier | +| 229 560 | 4 | `tmem_ptr_in_smem` | +| **device 合计** | **229 564** | | +| **host 申请** | **230 188** | `33 580 (extra) + 6 × 32 768` | + +每个 k_block：单 CTA 载入 32 KB（A 16 KB + B 16 KB），CTA-pair 计算 `256 × 256 × 64 × 2 = 8.39 MFLOP`。若无 2-CTA 切分，每 CTA 需要自己存整块 `256 × 64` 的 B（32 KB），stage 数会从 6 掉到 4——这就是 `cta_group::2` 的直接收益。 + +--- + +## 6. Tensor Memory 布局 + +### 6.1 TMEM 与地址格式 + +SM100 每个 SM 有 256 KB 的 Tensor Memory，组织为 **128 行（datapath / lane）× 512 列 × 32 bit**。TMEM 地址是 32 位： + +``` +bit 31 16 15 0 +┌──────────────┬─────────────┐ +│ row (lane) │ column │ +└──────────────┴─────────────┘ +``` + +代码里两处直接操作这个格式： + +- 累加器基址 `accum_stage_idx * UMMA_N` —— 纯列偏移，行号由硬件按 warp/lane 隐式定位（§3.2 的注释）。 +- swap-AB epilogue 的 `cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr | 0x00100000, v4..v7)` —— `0x00100000` 是 `1 << 20`，即行域 `+16`，把第二次加载定位到第 16–31 行（第一次是 0–15 行，`16dp` = 16 datapath）。 + +### 6.2 分配 / 释放协议 + +```cpp +using Allocator = cute::conditional_t; + +// prologue，warp 2 全体 32 lane（tcgen05.alloc 是 .sync.aligned，需整 warp） +Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem); + +// epilogue，warp 0 全体 +Allocator().free(0, kNumTmemCols); +``` + +- 分配结果写进 SMEM 的 `tmem_ptr_in_smem`，再由 `__syncthreads()` / `cluster_sync` 广播给所有角色。 +- `Allocator2Sm` 用于 `cta_group::2`：一次分配覆盖 CTA-pair 的两份 TMEM。 +- 分配**之前**必须先 `cluster_sync_with_relaxed_arrive()`（第 102 行），注释：*Synchronize the cluster before 2-CTA TMEM allocation*。2-CTA 的 alloc 是 pair 级操作，两个 CTA 必须同时在场。 +- 释放传 `0` 作为基址——正是 §3.2 那条断言保证的前提。 +- `kNumTmemCols = utils::get_num_aligned_tmem_cols()` 向上取整到 `{32, 64, 128, 256, 512}`，因为 `tcgen05.alloc` 只接受 2 的幂且 ≥ 32。 + +### 6.3 累加器双缓冲 + +```cpp +constexpr uint32_t kNumAccumTmemCols = kNumEpilogueStages * UMMA_N; // 2 * UMMA_N +accum_stage_idx = scheduler.current_iter % kNumEpilogueStages; // 0,1,0,1,... +accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; // 0,0,1,1,0,0,... +``` + +| 缓冲 | 列区间 | 内容 | +| --- | --- | --- | +| `acc[0]` | `[0, UMMA_N)` | 偶数 iter 的输出块 | +| `acc[1]` | `[UMMA_N, 2*UMMA_N)` | 奇数 iter 的输出块 | + +MMA warp 与 epilogue warp **各自独立**用 `current_iter` 算出同一对 `(accum_stage_idx, accum_phase_idx)`，不需要任何显式传递。这正是 §4.1「复制式调度器 + `current_iter` 当逻辑时钟」的直接收益。 + +握手：MMA warp 在某块最后一个 k_block 上 `tcgen05.commit` → `tmem_full_barriers[accum_stage_idx]`；epilogue 读到寄存器后立刻 `arrive(0u)` → `tmem_empty_barriers[accum_stage_idx]`（§11.4 的「早释放」）。MMA warp 在开始下一块前 `tmem_empty_barriers[accum_stage_idx]->wait(accum_phase_idx ^ 1)`。 + +`UMMA_N = 256` 时 `kNumAccumTmemCols = 512`，双缓冲恰好占满 TMEM，这也是 host 侧 `2 * umma_n > 512 → 跳过候选` 的由来。 + +### 6.4 UMMA 的累加目标 + +`tcgen05.mma` 的第 0 个操作数 `[tmem_c]` 就是 `accum_stage_idx * UMMA_N`： + +```cpp +mma_t::fma(a_desc, b_desc, accum_stage_idx * UMMA_N, + kUMMAKIdx > 0 or k_block_idx > 0, runtime_instr_desc); +// ↑ scale_c 谓词 +``` + +`scale_c` 在 PTX 里被翻成 `setp.ne.b32 p, %4, 0;` 后作为 `tcgen05.mma` 的尾随谓词：`p = false` 时 `D = A*B`（丢弃 TMEM 原值），`p = true` 时 `D = A*B + D`。所以一个输出块的**第一条** UMMA（`k_block_idx == 0 && kUMMAKIdx == 0`）传 0 完成清零，其余全部传 1 累加——省掉了单独 memset TMEM 的一步，也让上一轮遗留的脏数据自动失效。 + +--- + +## 7. 矩阵 Tiling 层次 + +### 7.1 六级 tiling + +以 §5.6 的例子（`BLOCK_M=128, BLOCK_N=256, BLOCK_K=64, cluster=2`）为例，从全局到指令共六级： + +| 级别 | 尺度 | 承载者 | 说明 | +| --- | --- | --- | --- | +| L0 全局 | `M × N × K` | grid | persistent，`gridDim.x = kNumSMs` | +| L1 wave | `kNumSMs` 个输出块 | grid 的一轮 | `iter` 递增一次；L2 swizzle 在此层重排 | +| L2 cluster tile | `(2*BLOCK_M) × BLOCK_N = 256 × 256` | CTA-pair | 一条 `cta_group::2` UMMA 覆盖的范围 | +| L3 CTA tile | `BLOCK_M × BLOCK_N = 128 × 256` | 单 CTA | 落在本 CTA 的 TMEM（128 行 × 256 列 × 2 缓冲） | +| L4 k_block / stage | `BLOCK_M × BLOCK_K`（A）+ `LOAD_BLOCK_N × BLOCK_K`（B） | SMEM ring 的一格 | 流水的调度单位，`kNumStages` 格在飞 | +| L5 UMMA atom | `UMMA_M × UMMA_N × UMMA_K = 256 × 256 × 16` | 一条指令 | `BLOCK_K / UMMA_K = 4` 条 UMMA 消费一个 stage | +| L6 swizzle atom | `8 行 × 128 B` | SMEM 物理布局 | TMA 写入与 UMMA 读出共用的最小图案单位 | + +K 方向的总迭代数：`num_total_k_blocks = ceil_div(scheduler.current_shape_k, BLOCK_K)`。注意用的是 `scheduler.current_shape_k` 而非 `shape_k`——k-grouped 变体里每个 group 的 K 长度不同，且这个值是**运行期**的，所以 K 循环不能整体展开，只能展开内层的 `BLOCK_K/UMMA_K` 次 UMMA。 + +### 7.2 2-CTA（`cta_group::2`）如何切分 A / B / D + +`kNumMulticast == 2` 时，一对 CTA（cluster rank 0 = leader，rank 1 = peer）协作完成一个 `UMMA_M = 256` 的 MMA。切分方式取决于 `kIsMulticastOnA`： + +**情形 A：`kIsMulticastOnA == false`（非 swap-AB，`cluster_m = 2`）** + +``` + N ────────── BLOCK_N = 256 ──────────► + M ┌────────────────────────────────────┐ + CTA0 │ D[0:128, 0:256] → CTA0 TMEM │ A: CTA0 存 rows 0..127 + ▼ ├────────────────────────────────────┤ (LOAD_BLOCK_M = 128) + CTA1 │ D[128:256, 0:256] → CTA1 TMEM │ B: CTA0 存 cols 0..127 + └────────────────────────────────────┘ CTA1 存 cols 128..255 + (LOAD_BLOCK_N = 128) +``` + +- 调度器把**相邻的 `m_block_idx`** 分给 cluster 内的两个 CTA（`kIsMulticastOnA=false` → 组内 M 变化最快）。 +- A 沿 M 对半切，各 CTA 存自己那 128 行；B 沿 N 对半切，各 CTA 存 128 列。 +- 硬件跨 CTA-pair 读取 B，两个 SM 的 tensor core 合起来算出 `256 × 256`，各自把属于自己的 128 行写进本地 TMEM。 +- Epilogue 时每个 CTA 用**自己的** `m_block_idx` 算 `base_m_idx`，独立存自己那 `128 × 256`。 + +**情形 B：`kIsMulticastOnA == true`（swap-AB，`cluster_n = 2`）** + +角色互换：cluster 内两个 CTA 拿**相邻的 `n_block_idx`**（组内 N 变化最快），`LOAD_BLOCK_M = BLOCK_M / 2`、`LOAD_BLOCK_N = BLOCK_N = 128`，`UMMA_M = 256` 对应 GEMM 的 N 方向。 + +两种情形下 `full_barriers[i]->init(kNumMulticast)` 与 `arrive_and_expect_tx(bytes * kNumMulticast)` 的语义都是「pair 内两个 CTA 各自的 TMA 都到齐」。 + +### 7.3 索引计算：`get_global_idx` 的双模板开关 + +```cpp +uint32_t m_idx = scheduler.get_global_idx<(kGemmType == GemmType::MGroupedMasked), IndexType::MN> + (shape_m, BLOCK_M, m_block_idx); +uint32_t n_idx = scheduler.get_global_idx<(kMajorB == cute::UMMA::Major::K), IndexType::MN> + (shape_n, BLOCK_N, n_block_idx, m_block_idx); +uint32_t k_a_idx = scheduler.get_global_idx<(kMajorA == cute::UMMA::Major::MN), IndexType::K> + (shape_k, BLOCK_K, k_block_idx, m_block_idx); +uint32_t k_b_idx = scheduler.get_global_idx<(kMajorB == cute::UMMA::Major::MN), IndexType::K> + (shape_k, BLOCK_K, k_block_idx, m_block_idx); +``` + +第一个模板参数是 `kWithGroupOffset`，第二个是索引语义（`MN` / `K` / `SF_K`）。要点： + +- `n_idx` 的 `kWithGroupOffset` 是 `kMajorB == K`。因为 B 的 group 维度**总是拼在外维**（见 `make_tma_b_desc` 的 `gmem_outer_dim * num_groups`）：K-major 时外维是 N，所以要加 `group * shape_n`；MN-major 时外维是 K，group 偏移由 `IndexType::K` 那条分支处理。 +- `k_a_idx` / `k_b_idx` 的 `kWithGroupOffset` 是 `major == MN`：MN-major 时 K 是外维，k-grouped 的偏移 `current_k_cumsum` / `current_k_start` 加在 K 上；K-major 时 K 是内维，`k_idx = k_block_idx * BLOCK_K` 直接用。 +- 源码注释：*`k_idx` is actually the k index default for K-major, while `k_b_idx` may be MN-major*。 +- 第 218 行的 `uint32_t k_idx = k_block_idx * BLOCK_K;` 实际上在后续代码里**没有被使用**（`tma::copy` 收到的是 `k_a_idx` / `k_b_idx`），是残留变量。 + +随后叠加 2-CTA 偏移（第 225–228 行）： + +```cpp +if constexpr (kNumMulticast > 1) { + m_idx += kIsMulticastOnA ? (block_rank_in_cluster() * load_block_m) : 0; + n_idx += kIsMulticastOnA ? 0 : (block_rank_in_cluster() * LOAD_BLOCK_N); +} +``` + +注意 M 方向用的是**运行期**的 `load_block_m`（swap-AB 时来自 `get_aligned_effective_m_in_block(m_block_idx) / kNumMulticast`，以适配 psum layout 的尾块），N 方向用编译期的 `LOAD_BLOCK_N`。 + +### 7.4 swap-AB：把「参差不齐的维度」放到 UMMA 的 N 上 + +这是本 kernel 最值得注意的一个架构决策。 + +`UMMA_M` 被 TMEM 的 datapath 数量钉死在 `128 × kNumMulticast`，**不能运行期改**；而 `UMMA_N` 只是 instruction descriptor 里的一个 5-bit 字段（`n_dim_ = umma_n >> 3`），**可以逐块改写**，粒度 8（非 swap）/ 16（swap，见 `get_aligned_effective_m_in_block` 里的 `UMMA_STEP_N = 16`）。 + +MoE / m-grouped 场景里每个 group 的有效 token 数（M）是动态且零碎的。若按常规方向做，M 必须向上取整到 `BLOCK_M`，padding 部分的 MMA 全部白算。swap-AB 之后： + +```cpp +// 操作数对调 +mma_t::fma(b_desc, a_desc, accum_stage_idx * UMMA_N, ...); + +// 逐块动态改 UMMA_N = 有效 M +if constexpr (kSwapAB) { + uint32_t umma_n = scheduler.get_aligned_effective_m_in_block(m_block_idx); + mma::sm100::update_instr_desc_with_umma_n(instr_desc, umma_n); // desc.n_dim_ = umma_n >> 3; +} +``` + +于是无效 token 直接**不发射 MMA**。代价是 TMEM 里的累加器变成了 `D^T`（行 = n，列 = m），epilogue 必须转置，这就是 `sm100_store_cd_swap_ab.cuh` 用 `stmatrix...trans` 的原因。 + +### 7.5 Stage 合并：用更大的 `BLOCK_K` 摊薄 `umma_arrive` + +第 43–50 行： + +```cpp +constexpr bool kDoMergeStages = + kNumStages_ >= 8 and kGemmType == GemmType::Normal and + kMajorA == cute::UMMA::Major::K and kMajorB == cute::UMMA::Major::K; +constexpr uint32_t kNumMinStages = 8; +constexpr uint32_t kNumStagesPerMerge = kDoMergeStages ? kNumStages_ / kNumMinStages : 1; +constexpr uint32_t BLOCK_K = BLOCK_K_ * kNumStagesPerMerge; // 64 → 128/192/... +constexpr uint32_t kNumStages = kNumStages_ / kNumStagesPerMerge; +``` + +注释说明动机：*this is for reducing the `umma_arrive()` overhead*。每个 k_block 结束时都要做一次 `tcgen05.commit`（+ 可能的 `tmem_full` commit）和一次 `full_barriers[stage]->wait()`，这些是 MMA warp 的**串行开销**。把 2 个 64-宽的 stage 合成 1 个 128-宽的 stage 后，同步次数减半、每次 UMMA 连发数从 4 增到 8，而**总的 SMEM 占用和流水深度（字节数）不变**。 + +合并只在 `kNumStages_ ≥ 8` 时触发，并保证合并后仍至少有 `kNumMinStages = 8` 个 stage（否则流水深度不足以掩盖延迟）。 + +合并后 SMEM 布局的关键点：一个 stage 内不再是「128 行 × 128 列」的单一 atom，而是 `kNumStagesPerMerge` 个 **K-atom（64 宽）沿「MN 外、K-atom 内」**排列。这一点在三处保持一致： + +1. **TMA**：`tma::copy` 内部 `BLOCK_INNER_ATOM = 128/2 = 64`，循环 2 次，第 i 次写到 `smem + i * LOAD_BLOCK_M * 64`。 +2. **UMMA 描述符**：构造时用 `BLOCK_ATOM_K = BLOCK_K / kNumStagesPerMerge = 64`（**不是** `BLOCK_K`），保证 `DG_STATIC_ASSERT(kSwizzleMode == BLOCK_ATOM_K * sizeof(dtype))` 即 `128 == 64*2` 仍成立；推进时 `kAtomKIdx = kUMMAKIdx * UMMA_K / BLOCK_ATOM_K`，偏移 `kAtomKIdx * LOAD_BLOCK_M * BLOCK_ATOM_K`。 +3. **stage 步长**：`a_desc_lo` 用 `SMEM_A_SIZE_PER_STAGE`（按合并后的 `BLOCK_K` 算）作为 lane 间的步长。 + +举例：`kNumStages_ = 18` → `kNumStagesPerMerge = 2`、`BLOCK_K = 128`、`kNumStages = 9`；每个 k_block 发 `128/16 = 8` 条 UMMA，`kAtomKIdx ∈ {0,0,0,0,1,1,1,1}`，`kInnerKIdx ∈ {0,16,32,48,0,16,32,48}`。 + +### 7.6 `BLOCK_M < 128` 时的算力浪费（有意为之） + +`UMMA_M = LAYOUT_AD_M * kNumMulticast` 恒为 128 或 256，**与 `BLOCK_M` 无关**。当启发式因 `desc.m ≤ 32` / `≤ 64` 选出 `BLOCK_M = 32` / `64` 时： + +- UMMA 仍按 M=128 发射，读 128 行 A（其中 96/64 行是 SMEM 越界垃圾，§5.5），往 TMEM 写 128 行 D； +- epilogue 只读 `STORE_BLOCK_M = BLOCK_M` 行，`kNumUMMAStoreThreads = BLOCK_M` 个线程； +- 结果正确的部分只有前 `BLOCK_M` 行。 + +第 274–277 行的指令形状断言里虽然列了 `UMMA_M == 64` 的合法分支，但 `UMMA_M` 的推导式永远产生不出 64，那段是通用的形状合法性检查（注释也说明 *CUTLASS does not have such checks except the MMA traits, but we are not using these traits*）。 + +这是一个**明确的取舍**：小 M 场景本来就是访存/延迟受限（选小 `BLOCK_M` 的目的正是注释里的 *avoid TMA L2 OOB bound*，即不去 GMEM 白读 128 行），MMA 吞吐富余，用一条统一代码路径换掉「M=64 UMMA + 另一套描述符/断言」的复杂度是划算的。 + +### 7.7 Tail-K + +```cpp +constexpr bool kMayHaveTailKBlock = ...; +for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + ... + if constexpr (kMayHaveTailKBlock) { + auto issue_tail_k_block = [&](const uint32_t& remaining_k) { + const auto num_valid_umma_k = math::ceil_div(remaining_k, UMMA_K); + utils::for_each_static_prefix(std::make_integer_sequence(), + num_valid_umma_k, issue_umma); + }; + const auto is_last_k_block = k_block_idx == num_total_k_blocks - 1; + if (is_last_k_block) { + const auto remaining_k = scheduler.current_shape_k - k_block_idx * BLOCK_K; + if (remaining_k < BLOCK_K) issue_tail_k_block(remaining_k); + else issue_full_k_block(); + } else { + issue_full_k_block(); + } + } else { + issue_full_k_block(); + } +} +``` + +三层设计： + +1. **编译期消除**：`if constexpr (kMayHaveTailKBlock)` —— K 被编译进来且整除时整段不生成代码（§2.1）。 +2. **运行期只在最后一个 k_block 判断**：`is_last_k_block` 之外的迭代走 `issue_full_k_block()`，把动态分支的代价压到 1/`num_total_k_blocks`。 +3. **前缀展开**：`for_each_static_prefix` 在 `BLOCK_K/UMMA_K ≤ 4` 时用 `switch(num_valid)` 跳到对应的编译期前缀（注释：*Prefix expansion uses switch only for small cases to avoid long SASS*），`> 4` 时退化成运行期谓词的 fold expression。 + +`ceil_div(remaining_k, UMMA_K)` 向上取整意味着最后一条 UMMA 可能读到最多 15 个 padding 元素——它们由 TMA 的 OOB 零填充保证为 0（`CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE`），对结果无影响。而 k-grouped 路径有 `DG_STATIC_ASSERT(kKAlignment % UMMA_K == 0)`，`remaining_k` 必是 16 的倍数，`ceil_div` 不会真的向上取。 + +--- + +## 8. TMA 加载路径 + +### 8.1 调用形态 + +TMA warp（warp 0 的单个 lane）在每个 k_block 上按 `kMajorA/kMajorB` 四选一发射： + +```cpp +if constexpr (kMajorA == cute::UMMA::Major::K) + tma::copy( + &tensor_map_a, full_barriers[stage_idx], smem_a[stage_idx], k_a_idx, m_idx, kNumMulticast, batch_idx); +if constexpr (kMajorA == cute::UMMA::Major::MN) + tma::copy( + &tensor_map_a, full_barriers[stage_idx], smem_a[stage_idx], m_idx, k_a_idx, kNumMulticast, batch_idx); +// B 同理 +``` + +模板参数序是 ``，函数参数序是 `(desc, barrier, smem_dst, inner_idx, outer_idx, num_multicast, batch_idx)`。**inner 恒为 SMEM 里连续的那一维**，所以 K-major 时 `(inner, outer) = (k, mn)`，MN-major 时 `(mn, k)`，两个 `if constexpr` 分支只是把实参顺序换了一下。 + +### 8.2 swizzle-atom 循环 + +```cpp +constexpr uint32_t BLOCK_INNER_ATOM = get_inner_block_atom_size(); +// = kSwizzleMode == 0 ? BLOCK_INNER : kSwizzleMode / sizeof(dtype_t) + +#pragma unroll +for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) + SM90_TMA_LOAD_2D::copy(desc_ptr, reinterpret_cast(barrier_ptr), + EVICT_NORMAL, + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, + inner_idx + i * BLOCK_INNER_ATOM, outer_idx); +``` + +TMA box 的内维被 host 压到 `swizzle/elem = 64` 个元素（128 B），所以一个逻辑上 `BLOCK_INNER` 宽的块要拆成 `BLOCK_INNER / 64` 条 TMA： + +| 场景 | `BLOCK_INNER` | atom | TMA 条数 | SMEM 目的地址步进 | +| --- | --- | --- | --- | --- | +| K-major A，未合并（`BLOCK_K=64`） | 64 | 64 | 1 | — | +| K-major A，合并后（`BLOCK_K=128`） | 128 | 64 | 2 | `LOAD_BLOCK_M * 64` | +| MN-major A，`LOAD_BLOCK_M=128` | 128 | 64 | 2 | `BLOCK_K * 64` | + +目的地址步进 `BLOCK_OUTER * BLOCK_INNER_ATOM` 正是 §5.3 描述的「atom 沿外维堆叠」布局，与 `make_umma_desc` 的 SBO/LBO 推导严格对偶。 + +### 8.3 三种 TMA 变体 + +```cpp +if (num_tma_multicast == 1) { + cute::SM90_TMA_LOAD_2D::copy(...); // cp.async.bulk.tensor.2d...（单 CTA） +} else { + #if __CUDA_ARCH__ >= 1000 + cute::SM100_TMA_2SM_LOAD_2D::copy(...); // 带 .cta_group::2 + #elif __CUDA_ARCH__ >= 900 + if (cute::block_rank_in_cluster() == 0) + cute::SM90_TMA_LOAD_MULTICAST_2D::copy(..., (1 << num_tma_multicast) - 1, ...); + #endif +} +``` + +- **1-CTA**：`SM90_TMA_LOAD_2D`，Hopper 就有的 `cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint`。DeepGEMM 在 SM100 上复用它。 +- **SM100 2-CTA**：`SM100_TMA_2SM_LOAD_2D`，即带 `.cta_group::2` 修饰的 bulk tensor copy。源码注释点明关键语义：*2-CTA function will send signals to the leader CTA only* —— 两个 CTA **各自发射**自己那一份 load，但 tx-count 只累加到 CTA-pair leader 的 mbarrier 上。这是 §10.3 里 `full_barriers` 计数推导的基础。 +- **SM90 multicast**：只由 rank 0 发射一条带 CTA mask 的 multicast load，一份 GMEM 读同时写进 cluster 内所有 CTA 的 SMEM。SM100 路径不用它（2-CTA UMMA 不能动态关闭，见 §4.2）。 + +Cache hint 统一是 `EVICT_NORMAL`，并且开头有一条静态断言确保 SM90/SM100 两套枚举值一致。对比 [ptx/tma.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/ptx/tma.cuh) 里手写的 `tma_load_1d` 用的是 `EVICT_FIRST`（注释：*normally, the loaded part will be evicted soon*），而 store 用 `EVICT_NORMAL`（*the stored part will be used soon*）——本 kernel 的 A/B 走 cute 封装故为 NORMAL。 + +### 8.4 Descriptor prefetch + +```cpp +if (warp_idx == 0) { + cute::prefetch_tma_descriptor(&tensor_map_a); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_cd); +} +``` + +在 kernel 最开头、**任何同步之前**由 warp 0 全体执行（`prefetch.tensormap` 不是单线程指令）。descriptor 在常量内存里，首次 TMA 访问会有冷启动延迟，提前 prefetch 可以把它藏进 barrier 初始化与 TMEM 分配的时间里。 + +### 8.5 `expect_tx` 字节数 + +```cpp +constexpr uint32_t kNumArrivalBytes = SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE; +if (is_leader_cta) { + full_barriers[stage_idx]->arrive_and_expect_tx(kNumArrivalBytes * kNumMulticast); +} else { + full_barriers[stage_idx]->arrive(0u); // 远程 arrive 到 cluster rank 0 +} +``` + +- `SMEM_*_SIZE_PER_STAGE` 用的是**合并后**的 `BLOCK_K`，与 TMA 实际搬运的字节数一致。 +- `× kNumMulticast`：pair 内两个 CTA 各搬一份，tx 都记在 leader 的 barrier 上。 +- peer CTA 只做一次「无 tx 的 arrive」，凑够 `init(kNumMulticast)` 的到达计数。 +- **顺序**：TMA 先发射（第 233–244 行），`arrive_and_expect_tx` 后执行（第 248–252 行）。这是合法的——mbarrier 的 tx-count 是「期望值累加」，只要在该相位完成前把期望值补上即可；反过来（先 expect 后发 TMA）同样合法但会让 warp 更早阻塞在计数上。 + +### 8.6 完整的 TMA warp 循环 + +```cpp +while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + const auto load_block_m = kSwapAB ? scheduler.get_aligned_effective_m_in_block(m_block_idx) / kNumMulticast + : LOAD_BLOCK_M; + const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); // ① 等消费者释放 + /* ② 算 m_idx / n_idx / k_a_idx / k_b_idx / batch_idx，叠加 2-CTA 偏移 */ + /* ③ 发 A、B 的 TMA（各 1~2 条指令） */ + /* ④ arrive_and_expect_tx（leader）/ arrive(0u)（peer） */ + } +} +``` + +整个 TMA warp 就是这四步的无限重复，没有任何计算。它的推进速度只受 `empty_barriers` 的释放节奏限制，因此可以超前 MMA 多达 `kNumStages` 个 k_block。 + +--- + +## 9. UMMA 发射路径 + +### 9.1 Instruction descriptor + +```cpp +auto instr_desc = kSwapAB + ? cute::UMMA::make_instr_desc() + : cute::UMMA::make_instr_desc(); +... +const auto runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); +``` + +- 编码了 A/B 数据类型（BF16）、累加类型（FP32）、MMA 形状（`UMMA_M × UMMA_N`）、A/B 的 major 模式。swap-AB 时把两个 major 也对调。 +- `make_runtime_instr_desc` 把这个 32-bit 描述符左移 32 位打包成 `uint64_t`，与 PTX 里的取用方式对应：`"r"(static_cast(desc >> 32))` —— 只取高 32 位作为 `tcgen05.mma` 的第 4 个操作数。 +- swap-AB 时 `update_instr_desc_with_umma_n(instr_desc, umma_n)` 改的是 `desc.n_dim_ = umma_n >> 3`，而 `runtime_instr_desc` 在 k_block 循环**内部**每轮重新计算（第 321 行），所以逐块的 `umma_n` 变化能立刻生效。 + +### 9.2 SMEM descriptor（`SmemDescriptor`） + +`make_smem_desc()` 填充 7 个字段： + +| 字段 | 值 | 含义 | +| --- | --- | --- | +| `version_` | `1` | SM100 版本标记 | +| `lbo_mode_` | `0` | legacy 模式 | +| `layout_type_` | `to_umma_layout_type<...>()` | `SWIZZLE_NONE / 32B / 64B / 128B / 128B_BASE32B` | +| `start_address_` | `cast_smem_ptr_to_uint(p) >> 4` | SMEM 地址，**16 B 为单位** | +| `base_offset_` | `0` | — | +| `stride_byte_offset_` (SBO) | `stride_byte_offset >> 4` | atom 间在某一维上的字节步长 | +| `leading_byte_offset_` (LBO) | `leading_byte_offset >> 4` | atom 间在另一维上的字节步长 | + +`to_umma_layout_type()` 有一个特例：`dtype == float && major == MN`，或显式 `kUseBase32`，返回 `SWIZZLE_128B_BASE32B`；对应 `get_atom_base()` 返回 32 而非 16，进而 `num_non_contiguous = 128 / 32 = 4`（常规是 `128 / 16 = 8`）。BF16 路径永远走常规分支，`num_non_contiguous = 8`。 + +**K-major** 的 SBO/LBO 推导： + +```cpp +DG_STATIC_ASSERT(kSwizzleMode * kPackFactor == BLOCK_K * sizeof(dtype_t)); // 128 == 64 * 2 +const uint32_t stride_byte_offset = num_non_contiguous * BLOCK_K * sizeof(dtype_t) / kPackFactor; // 8*64*2 = 1024 +const uint32_t leading_byte_offset = 0; +``` + +注释解释了为什么 LBO 是 0：*on K, there is only 1 atom as asserted previously*——那条静态断言保证每个 block 在 K 方向恰好只有一个 swizzle atom，于是「K 方向 atom 间步长」无意义。SBO = 一个 atom 的字节大小 = `8 行 × 128 B = 1024 B`，即 MN 方向相邻 atom 的步长。 + +**MN-major**： + +```cpp +constexpr uint32_t BLOCK_MN_ATOM = tma::get_inner_block_atom_size(); // 64 +DG_DEVICE_ASSERT(mn_idx % BLOCK_MN_ATOM == 0); // 不允许 atom 内的 MN 偏移 +uint32_t stride_byte_offset = num_non_contiguous * BLOCK_MN_ATOM * sizeof(dtype_t); // 8*64*2 = 1024 +uint32_t leading_byte_offset = BLOCK_K * BLOCK_MN_ATOM * sizeof(dtype_t); +if constexpr (kSwizzleMode == 16) math::swap(stride_byte_offset, leading_byte_offset); +``` + +注释给出了 SBO/LBO 的语义约定：swizzle 时 `{SBO, LBO}` 是 atom 在 `{K, MN}` 上的步长；非 swizzle（`kSwizzleMode == 16`，*means non-swizzling but interleaving*）时是 `{MN, K}`，所以要 swap。 + +`kPackFactor` 只对 packed FP4 是 2（*Packed FP4 stores two logical elements per byte in SMEM*），BF16 恒为 1，并有 `DG_STATIC_ASSERT(kPackFactor == 1 or sizeof(dtype_t) == 1)` 兜底。 + +### 9.3 用 warp 的 32 个 lane 存 per-stage 描述符 + +```cpp +DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); +constexpr uint32_t BLOCK_ATOM_K = BLOCK_K / kNumStagesPerMerge; +auto a_desc = mma::sm100::make_umma_desc(smem_a[0], 0, 0); +auto b_desc = mma::sm100::make_umma_desc(smem_b[0], 0, 0); +uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * SMEM_A_SIZE_PER_STAGE / 16 : 0u; +uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * SMEM_B_SIZE_PER_STAGE / 16 : 0u; +... +const auto a_desc_base_lo = __shfl_sync(0xffffffff, a_desc_lo, static_cast(stage_idx)); +const auto b_desc_base_lo = __shfl_sync(0xffffffff, b_desc_lo, static_cast(stage_idx)); +``` + +这是一个很巧的优化：`kNumStages` 个 stage 的描述符低 32 位（含 `start_address_`）**分散存放在 MMA warp 的 lane 0 … lane `kNumStages-1` 的寄存器里**，需要哪个 stage 就 `__shfl_sync` 广播出来。 + +- 省掉了 SMEM 里的描述符表（以及随之而来的 `ld.shared` 延迟与 bank 压力）； +- 省掉了每 stage 重算 `make_umma_desc` 的指令； +- `/ 16` 与 `start_address_` 的 16-B 粒度一致，`SMEM_*_SIZE_PER_STAGE` 是 1024 的倍数（§5.1 断言）所以整除无损； +- 代价是 `kNumStages ≤ 32` 这条硬约束（与 host 的 `kNumMaxStages = 32` 对应）。 + +`a_desc` / `b_desc` 这两个 64-bit 结构体被 `issue_umma` lambda 按引用捕获，循环里**只改 `.lo`**——`.hi`（SBO / base_offset / layout_type / version）是 stage 无关的常量。 + +### 9.4 K 内层展开与描述符推进 + +```cpp +auto issue_umma = [&]() { + constexpr uint32_t kAtomKIdx = kUMMAKIdx * UMMA_K / BLOCK_ATOM_K; // 第几个 64-宽 K atom + constexpr uint32_t kInnerKIdx = kUMMAKIdx * UMMA_K % BLOCK_ATOM_K; // atom 内 K 偏移 + a_desc.lo = advance_umma_desc_lo( + a_desc_base_lo, kAtomKIdx * LOAD_BLOCK_M * BLOCK_ATOM_K, kInnerKIdx); + b_desc.lo = advance_umma_desc_lo( + b_desc_base_lo, kAtomKIdx * LOAD_BLOCK_N * BLOCK_ATOM_K, kInnerKIdx); + kSwapAB ? mma_t::fma(b_desc, a_desc, accum_stage_idx * UMMA_N, kUMMAKIdx > 0 or k_block_idx > 0, runtime_instr_desc) + : mma_t::fma(a_desc, b_desc, accum_stage_idx * UMMA_N, kUMMAKIdx > 0 or k_block_idx > 0, runtime_instr_desc); +}; +auto issue_full_k_block = [&]() { + utils::for_each_static_until( + std::make_integer_sequence(), issue_umma); +}; +``` + +`for_each_static_until` 是一个 fold expression：`((kIdx < kNumValid ? func.template operator()() : void()), ...)`。`issue_umma` 是带显式模板参数列表的 generic lambda（C++20），所以 `kUMMAKIdx` 是**编译期常量**，`kAtomKIdx` / `kInnerKIdx` / 元素偏移全部常量折叠，`BLOCK_K/UMMA_K` 条 UMMA 完全展开且各自的描述符增量是立即数。 + +`advance_umma_desc_lo` 的算式： + +```cpp +return base + (((offset + k_idx * stride_k) * sizeof(dtype_t)) >> 4u); +// stride_k = (major == K) ? 1 : get_inner_block_atom_size() +``` + +- **K-major**：`offset` 是 K-atom 的元素偏移（`kAtomKIdx * LOAD_BLOCK_M * 64`，因为 atom 沿 MN 外维堆叠），`k_idx` 是 atom 内的 K 元素偏移，`stride_k = 1`（K 在 atom 内连续）。合计 `× 2 B >> 4` 换成 16-B 单位。未合并时 `kAtomKIdx` 恒为 0，`kInnerKIdx ∈ {0,16,32,48}` → `.lo` 依次 `+0, +2, +4, +6`（每条 UMMA 消耗 `16 × 2 B = 32 B = 2` 个 16-B 单位）。 +- **MN-major**：`stride_k = BLOCK_MN_ATOM = 64`，因为 MN-major 下 K 是**外维**，K 前进 1 要跨过一整个 64 宽的 MN atom。 + +### 9.5 PTX 指令 + +```cpp +using mma_t = cute::conditional_t; +``` + +[ptx/tcgen05.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh) 里的实现： + +```cpp +struct SM100_MMA_F16BF16_SS { + fma(desc_a, desc_b, tmem_c, scale_c, desc) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p; \n\t" + "}\n" + :: "r"(tmem_c), "l"(desc_a), "l"(desc_b), + "r"(static_cast(desc >> 32)), "r"(scale_c)); + } +}; +// 2x1SM 版本仅把 cta_group::1 换成 cta_group::2 +``` + +操作数语义： + +| 位置 | 约束 | 含义 | +| --- | --- | --- | +| `[%0]` | `"r"` | TMEM 累加器地址（列偏移 `accum_stage_idx * UMMA_N`） | +| `%1` | `"l"` | A 的 SMEM 描述符（64-bit） | +| `%2` | `"l"` | B 的 SMEM 描述符（64-bit） | +| `%3` | `"r"` | instruction descriptor（`desc >> 32`） | +| `p` | 由 `%4` 生成 | scale-D 谓词：0 → 覆写，1 → 累加 | + +后缀 `_SS` = 两个操作数都来自 Shared memory（对比 `_RS` 变体的 A 来自寄存器）。BF16 走 `kind::f16`；同文件还有 `mxf8f6f4` / `f8f6f4` / `mxf4`（带 `.block_scale` 与额外的 `[tmem_sfa], [tmem_sfb]` 操作数）以及 `tcgen05.mma.ws`（weight-stationary）变体，本 kernel 都不用。 + +`asm volatile` 且无输出操作数——防止编译器把「重复」的 MMA 指令合并或重排。 + +### 9.6 MMA warp 的完整循环 + +```cpp +while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + accum_stage_idx = current_iter % 2; accum_phase_idx = (current_iter / 2) & 1; + tmem_empty_barriers[accum_stage_idx]->wait(accum_phase_idx ^ 1); // ① 等 epilogue 放掉累加器 + ptx::tcgen05_after_thread_sync(); // ② tcgen05.fence::after_thread_sync + /* 定义 umma_arrive / empty_barrier_arrive lambda；swap-AB 时改 instr_desc 的 n_dim */ + for (k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); // ③ 等 TMA 到货 + ptx::tcgen05_after_thread_sync(); + if (elect_one_sync()) { /* ④ 发 BLOCK_K/UMMA_K 条 UMMA */ } + __syncwarp(); + empty_barrier_arrive(k_block_idx == num_total_k_blocks - 1); // ⑤ tcgen05.commit + /* ⑥ 可选的 tensor core 降速（§9.7） */ + } +} +``` + +```cpp +auto empty_barrier_arrive = [&](const bool& do_tmem_full_arrive) { + umma_arrive(reinterpret_cast(empty_barriers[stage_idx])); + // NOTES: the tensor memory accumulator pipeline has nothing to do with multicasting + if (do_tmem_full_arrive) + umma_arrive(reinterpret_cast(tmem_full_barriers[accum_stage_idx])); + __syncwarp(); +}; +auto umma_arrive = [](const uint64_t* barrier) { + if constexpr (kNumMulticast == 1) cutlass::arch::umma_arrive(barrier); + else cutlass::arch::umma_arrive_multicast_2x1SM(barrier, (1 << kNumMulticast) - 1); +}; +``` + +`tcgen05.commit` 的作用是：让指定的 mbarrier **跟踪此前由本 warp 发起的所有 `tcgen05.mma`**，当它们全部完成时自动产生一次 arrive。这是异步 MMA 与 mbarrier 之间唯一的桥梁——软件完全不需要轮询 MMA 状态。 + +两个 commit 的分工： + +- `empty_barriers[stage_idx]`：**每个** k_block 都 commit，通知 TMA warp「这个 SMEM stage 已被读完，可以覆盖」。 +- `tmem_full_barriers[accum_stage_idx]`：**只在最后一个** k_block commit（`do_tmem_full_arrive = (k_block_idx == num_total_k_blocks - 1)`），通知 epilogue「整块累加完成」。因为 commit 跟踪的是「此前全部 MMA」，最后一次 commit 天然覆盖了整块的所有 UMMA。 + +`umma_arrive_multicast_2x1SM` 的 CTA mask 是 `(1 << kNumMulticast) - 1 = 0b11`，一次 commit 同时向 pair 内两个 CTA 的同名 barrier 各投递一次 arrive——peer CTA 的 TMA warp 也需要知道 leader 的 UMMA 读完了自己那份 SMEM。 + +第 367 行的注释是一个重要的正确性说明：*No explicit `tcgen05.fence::before_thread_sync` is needed, as this is implicitly performed by `tcgen05.commit`*。而 epilogue 侧读 TMEM 后通知 `tmem_empty` 时，**必须**显式调用 `ptx::tcgen05_before_thread_sync()`（因为那里用的是普通 `mbarrier.arrive`，没有 commit 帮忙）。 + +### 9.7 Tensor Core 利用率控制（防掉频） + +```cpp +DG_STATIC_ASSERT(kTensorCoreUtilControl > 0, "Invalid tensor utilization control"); +if constexpr (kTensorCoreUtilControl < 100) { + umma_arrive(reinterpret_cast(tensor_core_full_barrier)); + __syncwarp(); + tensor_core_full_barrier->wait(tensor_core_phase); + tensor_core_phase ^= 1; + + constexpr static uint64_t kNumUMMACycles = (2ull * UMMA_M * UMMA_N * BLOCK_K) / 8192ull; + constexpr static uint64_t kNumDummyCycles = (100ull - kTensorCoreUtilControl) * kNumUMMACycles / kTensorCoreUtilControl; + const auto start_clock = clock64(); + if (cute::elect_one_sync()) + while (clock64() - start_clock < kNumDummyCycles) {} + __syncwarp(); +} +``` + +注释：*Let tensor cores relax for lower possibility of frequency drop*。这是一个**主动降速**机制： + +- `kNumUMMACycles = 2·M·N·K / 8192` —— 一个 k_block 的浮点运算量除以 8192 FLOP/cycle（该 kernel 假定的 tensor core 峰值吞吐）得到理论执行周期。 +- `kNumDummyCycles = (100 - util)/util × kNumUMMACycles` —— 要让占空比降到 `util%`，需要在每 `kNumUMMACycles` 的实际计算后插入的空转周期。`util = 50` → 空转与计算等长。 +- 空转前先用 `tcgen05.commit` + `wait` **确认上一批 UMMA 真的做完了**，否则量出来的占空比不准。 +- 由 `device_runtime->set_tc_util()` / Python `deep_gemm.set_tc_util()` 控制，默认 100（整段 `if constexpr` 不生成代码，`tensor_core_full_barrier` 也不 `init`）。 + +用途是在功耗受限的集群上做可复现的性能对比、或验证「降频是否是真凶」。对纯性能跑分应始终保持 100。 + +--- + +## 10. 生产者-消费者同步 + +### 10.1 barrier 全景 + +全部同步对象都是 `cutlass::arch::ClusterTransactionBarrier`（8 B SMEM 驻留的 mbarrier），共 5 类： + +| barrier | 个数 | `init` 计数 | 等待者 | 到达者 | 语义 | +| --- | --- | --- | --- | --- | --- | +| `full_barriers[s]` | `kNumStages` | `kNumMulticast` | MMA warp（leader） | leader：`arrive_and_expect_tx(bytes×mc)`
peer：`arrive(0u)`
＋TMA 硬件的 tx 完成 | SMEM stage `s` 已装好，A/B 可读 | +| `empty_barriers[s]` | `kNumStages` | `1` | TMA warp（**每个** CTA） | `tcgen05.commit`（`umma_arrive`，2-CTA 时 multicast 到 pair 两边） | stage `s` 已被 UMMA 读完，可覆写 | +| `tmem_full_barriers[a]` | 2 | `1` | Epilogue warps（**每个** CTA） | 最后一个 k_block 上的 `tcgen05.commit` | TMEM 累加器 `a` 已写满 | +| `tmem_empty_barriers[a]` | 2 | `kNumMulticast × kNumUMMAStoreThreads` | MMA warp（**仅** leader） | 每个 epilogue 线程的 `arrive(0u)`（远程投递到 cluster rank 0） | 累加器 `a` 已读完，可重新累加 | +| `tensor_core_full_barrier` | 1 | `1`（仅 `tc_util < 100`） | MMA warp | `tcgen05.commit` | 上一批 UMMA 已完成（用于算占空比） | + +两条流水嵌套关系： + +``` +【A/B SMEM 环】 生产者 = TMA warp 消费者 = MMA warp + full : TMA → MMA empty : MMA → TMA + +【TMEM 累加器环】 生产者 = MMA warp 消费者 = Epilogue warps + tmem_full : MMA → Epilogue tmem_empty : Epilogue → MMA + +【C/D SMEM 环】 生产者 = Epilogue(STSM) 消费者 = TMA store 引擎 + 不用 mbarrier，而用 cp.async.bulk.wait_group.read + NamedBarrier +``` + +MMA warp 同时是上游的消费者和下游的生产者，是整条链的**唯一串行点**——这也是为什么它只用一个 warp（甚至一个 lane）：UMMA 是异步指令，发射本身极便宜，真正的瓶颈在 `tcgen05.commit` 的次数（§7.5 的 stage 合并就是为了摊薄它）。 + +### 10.2 parity 相位约定与「首轮免等」 + +`Barrier::wait(p)` 对应 `mbarrier.try_wait.parity.shared::cta.b64 P, [bar], p`，含义是「等待奇偶性为 `p` 的那个相位**完成**」。两类角色的相位推进方式不同： + +**① A/B 环：显式相位变量** + +```cpp +uint32_t stage_idx = 0, phase = 0; +auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + stage_idx = (stage_idx + 1) % kNumStages; + phase ^= stage_idx == 0; // 只在回绕时翻转 +}; + +// 生产者（TMA）：empty_barriers[stage_idx]->wait(phase ^ 1); +// 消费者（MMA） ：full_barriers[stage_idx]->wait(phase); +``` + +第一轮（`phase = 0`）： + +- TMA 等 `empty` 的 parity **1**。新建的 mbarrier 处于 phase 0 且未完成，对 parity 1 的 `try_wait` 立即成功——这就是「缓冲初始为空，生产者前 `kNumStages` 轮直接穿过」的标准手法。 +- MMA 等 `full` 的 parity **0**，必须等第一次 tx 完成。 + +回绕一次后 `phase = 1`，两边等待的 parity 同时翻转，协议自洽。 + +**② TMEM 环：从 `current_iter` 直接推导** + +```cpp +accum_stage_idx = scheduler.current_iter % kNumEpilogueStages; // 0,1,0,1,… +accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; // 0,0,1,1,0,0,… + +// 生产者（MMA）入口 ：tmem_empty_barriers[accum_stage_idx]->wait(accum_phase_idx ^ 1); +// 消费者（Epilogue）入口：tmem_full_barriers[accum_stage_idx]->wait(accum_phase_idx); +``` + +同样地，`iter = 0` 时 MMA 等 parity 1 → 立即穿过（累加器初始可用）。好处是**无需维护额外的状态变量**，两个角色各自从同一个 `current_iter` 算出完全一致的 `(stage, phase)`。 + +注意 `tensor_core_phase` 用的是手动 `^= 1`（每轮翻转），因为它只有一个 barrier、不区分 stage。 + +### 10.3 arrive 计数逐条推导 + +**`full_barriers[s]->init(kNumMulticast)`**（源码注释：*Arrive only at the leader CTA*） + +- leader CTA：`arrive_and_expect_tx(kNumArrivalBytes * kNumMulticast)` —— 本地 arrive 1 次 + 设置期望 tx。 +- peer CTA：`arrive(0u)` —— `mapa.shared::cluster` 把地址映射到 cluster rank 0 后 `mbarrier.arrive.shared::cluster`，即远程 arrive 到 **leader** 的那份 barrier。 +- TMA 硬件：`SM100_TMA_2SM_LOAD_2D` 的 `.cta_group::2` 语义使两个 CTA 各自 load 的 tx 都只记到 leader 的 barrier（*2-CTA function will send signals to the leader CTA only*）。 +- 合计：`kNumMulticast` 次 arrive + `kNumMulticast × (A+B)` 字节 tx → 相位完成。只有 leader 的 MMA warp 需要等（UMMA 只由 leader 发射），peer 本地那份 `full_barriers` 无人等待。 + +**`empty_barriers[s]->init(1)`**（注释：*Arrive at all CTAs*） + +- `kNumMulticast == 1`：`cutlass::arch::umma_arrive` → `tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64` → 本地 1 次 arrive。 +- `kNumMulticast == 2`：`umma_arrive_multicast_2x1SM(bar, 0b11)` → `tcgen05.commit.cta_group::2.…multicast::cluster.b64 [bar], mask` → **pair 内每个 CTA 的同名 barrier 各收 1 次 arrive**。 +- 所以两个 CTA 各自的 `empty_barriers[s]` 都是「计数 1、由 leader 的 commit-multicast 填满」，而两个 CTA 的 TMA warp 各自等自己本地的那份。注释 *Arrive at all CTAs* 正是这个意思。 + +**`tcgen05.commit` 是 warp 级指令（重要）** + +第 364–368 行的结构是： + +```cpp +if (cute::elect_one_sync()) { /* 只有 1 个 lane 发 UMMA */ } +__syncwarp(); +empty_barrier_arrive(...); // ← 在 elect 块之外，warp 1 的 32 个 lane 全部执行 umma_arrive +``` + +而 `empty_barriers[s]->init(1)`。若每个 lane 都产生一次 arrive，相位会被多推 31 次，parity 协议立即崩。因此可反推出：**`tcgen05.commit` 是 warp 级（收敛）指令，整 warp 执行只产生一次 arrive**，它代表的是「本 warp 此前发起的全部 `tcgen05.mma` 完成」这一事件，而非每线程事件。 + +旁证：`sm100_bmk_bnk_mn.cuh` 里同样是 `empty_barriers[i]->init(1)` + 在 warp 作用域直接调 `cutlass::arch::umma_arrive(...)`；`sm100_fp8_fp4_gemm_1d1d.cuh` 也是同样的 `empty_barrier_arrive` 结构。三个 SM100 kernel 一致，说明这是有意的写法。这也解释了 `empty_barrier_arrive` 末尾那个 `__syncwarp()`：保证 commit 在 warp 收敛状态下发出，并与下一轮 `stage_idx` 的使用隔开。 + +**`tmem_full_barriers[a]->init(1)`** + +注释：*the tensor memory accumulator pipeline has nothing to do with multicasting*——意思是这个环不属于 A/B 的 multicast 切分体系，但它仍然用同一个 `umma_arrive` 封装，因此 2-CTA 时同样 multicast 到两个 CTA，每边各 1 次 arrive，而两个 CTA 的 epilogue warp 各自等本地那份（因为两个 CTA 的 TMEM 各存自己那 128 行，都需要被读走）。 + +**`tmem_empty_barriers[a]->init(kNumMulticast * kNumUMMAStoreThreads)`** + +- epilogue 里 `tmem_empty_barrier->arrive(0u)` 没有 `elect_one_sync()` 包裹，所以**每个参与线程都 arrive 一次** → 单 CTA `kNumUMMAStoreThreads` 次。 +- `arrive(0u)` 的目标固定是 cluster rank 0，所以 pair 内两个 CTA 的到达全部汇到 leader → `× kNumMulticast`。 +- 只有 leader 的 MMA warp 等它（因为只有 leader 发 UMMA）。peer CTA 本地的 `tmem_empty_barriers` 是死对象。 + +### 10.4 fence 序列 + +异步代理（TMA / tensor core）与普通线程之间的可见性需要专用 fence，本 kernel 共四处： + +| 位置 | 指令 | 作用 | +| --- | --- | --- | +| barrier 初始化后（第 167 行） | `cutlass::arch::fence_barrier_init()`
= `fence.mbarrier_init.release.cluster` | 让初始化后的 mbarrier 对 **async proxy** 可见（TMA/UMMA 会直接读写它们），cluster 作用域保证 peer CTA 也看得到 | +| 每次等完 barrier、发 tcgen05 指令前（第 285/317 行） | `ptx::tcgen05_after_thread_sync()`
= `tcgen05.fence::after_thread_sync` | 把「同步点之后」的 TMEM/MMA 操作与同步点正确排序，防止 tensor core 异步管线越过 barrier 提前取数 | +| epilogue 读完 TMEM、通知 `tmem_empty` 前（`sm100_store_cd.cuh:113`） | `ptx::tcgen05_before_thread_sync()`
= `tcgen05.fence::before_thread_sync` | 保证所有 `tcgen05.ld` 已完成并排序在 arrive 之前。**MMA 侧不需要**这条，因为 `tcgen05.commit` 隐含了它（源码第 367 行注释） | +| TMEM load 之后、用寄存器值之前（`sm100_store_cd.cuh:91/99`） | `cutlass::arch::fence_view_async_tmem_load()`
= `tcgen05.wait::ld.sync.aligned` | `tcgen05.ld` 也是异步的，必须等它真正写回寄存器 | +| STSM 写完 SMEM、发 TMA store 前（`sm100_store_cd.cuh:118`） | `cute::tma_store_fence()`
= `fence.proxy.async.shared::cta` | generic proxy（`st.shared` / `stmatrix`）→ async proxy（TMA）的 SMEM 可见性 | + +这五条 fence 是 SM100 异步编程模型里最容易遗漏、且遗漏后只在特定时序下出错的部分。 + +### 10.5 稳态时序（一个输出块的完整生命周期） + +以 `kNumStages = 6`、`num_total_k_blocks = 128`（K=8192, BLOCK_K=64）为例，纵轴是时间： + +``` +TMA warp │ empty.wait │ TMA(s=0) │ empty.wait │ TMA(s=1) │ … │ TMA(s=5) │ empty.wait(阻塞) │ TMA(s=0) │ … + │ expect_tx │ │ expect_tx │ ▲ + ▼ ▼ ▼ ▼ │ +full[0] ────●───────────────────────────────────────────────────────────────────┐ │ +full[1] ──────────●──────────────────────────────────────────────────────────────┼─┐ │ + │ wait(phase=0) │ │ │ +MMA warp │ tmem_empty.wait(穿过) │ full[0].wait │ 4×UMMA │ commit→empty[0] │ full[1].wait │ 4×UMMA │ commit→empty[0]… + │ │ + └──── 最后一个 k_block 额外 commit → tmem_full[a] ────┘ + │ +Epilogue tmem_full.wait(a) + │ + ┌────────────────────────────┘ + ▼ + for each (w,s) store stage: + tma_store_wait<1> → NamedBarrier.sync + → tcgen05.ld ×8 → fence(wait::ld) → st.shared.v4 ×8 + → 【最后一次】tcgen05.fence::before_thread_sync + tmem_empty.arrive(0u) ← 早释放 + → tma_store_fence → NamedBarrier.sync → TMA store + commit_group +``` + +三个重叠关系： + +1. **TMA 超前 MMA 最多 6 个 k_block**（受 `kNumStages` 限）。 +2. **MMA 超前 Epilogue 最多 1 个输出块**（受 TMEM 双缓冲限）。块 i 的 epilogue 与块 i+1 的全部 128 个 k_block 重叠。 +3. **STSM 超前 TMA store 最多 1 个 store stage**（受 `kNumTMAStoreStages = 2` 与 `wait_group.read 1` 限）。 + +### 10.6 退出协议 + +```cpp +// MMA warp（仅 leader CTA），循环结束后： +const auto iter_idx = scheduler.current_iter - 1; +if (kNumMulticast > 1 and iter_idx >= 0) { + const auto accum_phase_idx = (iter_idx / kNumEpilogueStages) & 1; + tmem_empty_barriers[iter_idx % kNumEpilogueStages]->wait(accum_phase_idx); +} + +// 所有 warp： +kNumMulticast > 1 ? comm::cluster_sync_with_relaxed_arrive() : __syncthreads(); // 第 452 行 +if (warp_idx == 0) Allocator().free(0, kNumTmemCols); +``` + +源码注释把第 392–397 行称为 *another round of waits*，理由是 *To safely deconstruct barriers*。具体危险：peer CTA 的 epilogue 线程通过 `arrive(0u)` **远程**写 leader CTA SMEM 里的 mbarrier。如果 leader 已经跑完并退出（SMEM 释放），peer 那次远程 arrive 就是一次非法访问。因此 leader 必须等到**最后一次** `tmem_empty` 到达（因为 barrier 按序使用，等最后一个就隐含前面的都已到达）。 + +注意这里用的是 `wait(accum_phase_idx)` 而不是循环里的 `wait(accum_phase_idx ^ 1)`——相位不取反，因为现在要等的是「该相位**已完成**」（即 epilogue 真的 arrive 了），而不是循环入口那个「上一轮已释放」的语义。 + +随后的 `cluster_sync_with_relaxed_arrive()`（= `cluster_arrive_relaxed()` + `cluster_wait()`，注释说明比 `cute::cluster_sync` 略快但内存序保证更弱）保证两个 CTA 都到达后才释放 TMEM。第 451 行的 `// TODO: Remove redundant synchronization` 说明作者也意识到这里的同步可能多余。 + +--- + +## 11. Epilogue + +### 11.1 入口 + +```cpp +} else if (warp_idx >= kNumNonEpilogueThreads / 32 and + warp_idx < (kNumNonEpilogueThreads + kNumUMMAStoreThreads) / 32) { + const auto epilogue_warp_idx = warp_idx - (kNumNonEpilogueThreads / 32); + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(tmem_ptr_in_smem) == 0); + uint32_t tma_stage_idx = 0; // 跳块共享 + while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + accum_stage_idx / accum_phase_idx ← current_iter + tmem_full_barriers[accum_stage_idx]->wait(accum_phase_idx); + ptx::tcgen05_after_thread_sync(); + tmem_base_addr = accum_stage_idx * UMMA_N; + base_m_idx = scheduler.get_global_idx<(not is_m_grouped_contiguous(kGemmType)), MN>(shape_m, BLOCK_M, m_block_idx); + base_n_idx = n_block_idx * BLOCK_N; + kSwapAB ? sm100_store_cd_swap_ab<...>(...) : sm100_store_cd<...>(...); + } +} +``` + +两个细节： + +- `base_m_idx` 的 `kWithGroupOffset` 是 `not is_m_grouped_contiguous(...)`——m-grouped-contiguous 的 group 偏移已经由 `make_tma_a_desc` 把 `m * num_groups` 拼进了 gmem 外维，所以**不能**再加；而 masked / psum 变体需要加 `current_group_idx * shape_m`。 +- `base_n_idx` 直接用 `n_block_idx * BLOCK_N`，不走 `get_global_idx`。B 的 group 偏移在加载侧已经处理，存储侧的 `tensor_map_cd` 对 m-grouped 把 group 拼在 M 维上（`make_tma_cd_desc(d, m, n, …, num_groups, …)` → `gmem_outer = m * num_groups`）。 +- `tma_stage_idx` 声明在块循环外，注释 *Share store pipeline between blocks*：C/D 环也跳块连续，与 A/B 环同理。 + +### 11.2 非 swap-AB：TMEM → RF → SMEM 的 swizzle 逐行推导 + +```cpp +constexpr uint32_t kNumBankGroupBytes = 16; +constexpr uint32_t kNumElemsPerBankGroup = 16 / sizeof(cd_dtype_t); // bf16→8, fp32→4 +constexpr auto kNumMWaves = BLOCK_M / STORE_BLOCK_M; // 本 kernel 恒为 1 +constexpr uint32_t kNumStores = BLOCK_N / STORE_BLOCK_N; // = BLOCK_N / (128/sizeof) + +for (w = 0; w < kNumMWaves; ++w) + for (s = 0; s < kNumStores; ++s, advance_store_pipeline()) { + smem_base_ptr = smem_cd[tma_stage_idx]; + if (epilogue_warp_idx == 0) cute::tma_store_wait(); // wait_group.read 1 + NamedBarrier::sync(kNumUMMAStoreThreads, 0); + + for (i = 0; i < STORE_BLOCK_N / kNumElemsPerBankGroup; ++i) { + auto bank_group_index = i + lane_idx * (kSwizzleCDMode / kNumBankGroupBytes); // i + lane*8 + constexpr bool kHasShortcut = (kSwizzleCDMode / kNumBankGroupBytes) == 8; // 128B swizzle → true + auto row = kHasShortcut ? (i / 8 + lane_idx) : (bank_group_index / 8); + auto col = kHasShortcut ? (i) : (bank_group_index % 8); + col ^= row % (kSwizzleCDMode / 16); // col ^= row % 8 + + uint32_t tmem_addr = tmem_base_addr + w * BLOCK_N + s * STORE_BLOCK_N + i * kNumElemsPerBankGroup; + auto smem_ptr = smem_base_ptr + epilogue_warp_idx * 32 * kSwizzleCDMode + + row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; + uint32_t values[kNumElemsPerBankGroup]; + // fp32：SM100_TMEM_LOAD_32dp32b4x → 4 个值 → st.shared.v4.f32 + // bf16：SM100_TMEM_LOAD_32dp32b8x → 8 个值 → cast_into_bf16_and_pack ×4 → st.shared.v4.u32 + } + … + } +``` + +逐步解读： + +1. **行划分**：`epilogue_warp_idx * 32 * kSwizzleCDMode` 把 warp `w` 定位到 SMEM 的第 `32w` 行。`kNumUMMAStoreThreads = STORE_BLOCK_M` 个线程恰好覆盖 `STORE_BLOCK_M` 行，**一线程一行**。 +2. **TMEM 读取形状**：`32dp32b{4,8}x` = 32 个 datapath（对应 warp 的 32 个 lane）× 32-bit × 4/8 次重复。一个 warp 一条指令读走 `32 行 × {4,8} 列` 的 FP32 累加器，正好是 `kNumElemsPerBankGroup` 列 = 一个 16-B bank group 的宽度。 +3. **`i` 循环**：跑 `STORE_BLOCK_N / kNumElemsPerBankGroup` 次（bf16：`64/8 = 8`；fp32：`32/4 = 8`），恰好覆盖 128-B 行内的 8 个 bank group。所以两种 dtype 下都是 **8 次 TMEM load + 8 次 128-bit `st.shared`**。 +4. **swizzle 计算**：`kHasShortcut` 分支在 `kSwizzleCDMode == 128` 时成立，`row = i/8 + lane_idx = lane_idx`（因为 `i < 8`），`col = i ^ (lane_idx % 8)`。通用分支先把 `(i, lane)` 线性化成 `bank_group_index` 再除/模 8——两者在 128B swizzle 下等价，前者省掉一次除法。 +5. **每次写入量**：一次 `st.shared.v4.u32` = 16 B = 一个 bank group。一个 warp 一轮写完 `32 行 × 128 B = 4 KB`，8 轮写完整个 stage（`STORE_BLOCK_M × 128 B`）。 +6. **无 bank conflict**：同一轮里 32 个 lane 写 32 个**不同行**的同一逻辑 bank group，而 `col ^= row % 8` 把它们打散到 8 个不同的物理 bank group。 + +### 11.3 swap-AB：靠 `stmatrix.trans` 做转置 + +swap-AB 下 TMEM 里存的是 `D^T`（行 = n，列 = m），而 GMEM 里的 D 是行主序的 `M × N`，必须转置： + +```cpp +DG_STATIC_ASSERT(STORE_BLOCK_N == 128, "STORE_BLOCK_N must be 128 to match TMEM rows"); +DG_STATIC_ASSERT(kSwizzleCDMode == 128, "TMA D must be 128B swizzled"); +constexpr uint32_t STORE_BLOCK_N_ATOM = kSwizzleCDMode / sizeof(cd_dtype_t); // bf16 → 64 +constexpr uint32_t kNumSwizzleAtomRows = 8; + +const auto num_stores = effective_m / STORE_BLOCK_M; // 运行期！effective_m 可达 BLOCK_M/16 = 16 轮 +for (s = 0; s < num_stores; ++s, advance_store_pipeline()) { + for (i = 0; i < STORE_BLOCK_M / kNumSwizzleAtomRows; ++i) { // 16/8 = 2 + tmem_addr = tmem_base_addr + s * STORE_BLOCK_M + i * kNumSwizzleAtomRows; + constexpr uint32_t kNumWarpsPerAtom = STORE_BLOCK_N_ATOM / 32; // 64/32 = 2 + outer_atom_offset = (epilogue_warp_idx / kNumWarpsPerAtom) * STORE_BLOCK_M * kSwizzleCDMode; + inner_atom_offset = i * kNumSwizzleAtomRows * kSwizzleCDMode; + + // bf16：两次 16dp256b 加载（第二次 tmem_addr | 0x00100000 即行 +16）→ 8 个值 + // → cast_into_bf16_and_pack ×4 → SM90_U32x4_STSM_T::copy(...) ← .trans + // fp32：一次 32dp32b8x → 8 个值 → 逐行 st.shared.u32，row = lane%8、col = (warp%2)*4 + lane/8 + } +} +``` + +- **行覆盖**：`16dp256b` 一次只覆盖 16 个 datapath，所以用 `tmem_addr | 0x00100000`（行域 `+16`）发第二次，两次合起来 32 行×8 列 = 32 lane × 8 值。 +- **warp 分工**：`kNumWarpsPerAtom = 2`，warp 0/1 写 atom 0（n 方向 0–63）、warp 2/3 写 atom 1（n 方向 64–127），每 atom `STORE_BLOCK_M(16) × 128 B = 2 KB`，四 warp 共 4 KB = 一个 stage。 +- **转置本体**：`stmatrix.sync.aligned.x4.m8n8.shared.b16.trans`（`SM90_U32x4_STSM_T`）——四个 lane 组各自提供 4×32-bit（= 8 个 bf16），硬件在写入 SMEM 时完成 8×8 转置。这是避免在寄存器里做显式 shuffle 转置的关键。 +- **`num_stores` 是运行期的**：`effective_m = get_aligned_effective_m_in_block(m_block_idx)`，psum layout 的尾块可以小于 `BLOCK_M`，因此循环不能展开——这正是 swap-AB 省算力的地方（§7.4）。 +- **TMA store 拆成 `STORE_BLOCK_N / STORE_BLOCK_N_ATOM = 2` 条**，每条 box = `64 × 16`，坐标 `(n_idx = base_n_idx + i*64, m_idx = base_m_idx + s*16)`。 + +### 11.4 TMEM 早释放（关键优化） + +```cpp +// Notify tensor memory empty (only at the leader CTA) arrival ASAP +// NOTES: only the last stage needs to do this +if (w == kNumMWaves - 1 and s == BLOCK_N / STORE_BLOCK_N - 1) { + ptx::tcgen05_before_thread_sync(); + tmem_empty_barrier->arrive(0u); +} + +// 之后才是： +cute::tma_store_fence(); +NamedBarrier::sync(kNumUMMAStoreThreads, 0); +if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { …TMA store…; cute::tma_store_arrive(); } +``` + +释放点被刻意放在「所有 `tcgen05.ld` 已完成、但 TMA store 尚未发出」之间。合法性依据：到达此处时，`i` 循环已经把整个 `BLOCK_M × BLOCK_N` 累加器全部读进了寄存器并写入了 SMEM（`fence_view_async_tmem_load()` 保证 `tcgen05.ld` 真的完成），TMEM 不再需要。而后续的 `tma_store_fence` + NamedBarrier + TMA store 只涉 SMEM。 + +效果：MMA warp 可以**立刻**开始下一个输出块（往另一个 TMEM 缓冲累加），而 epilogue 还在慢慢把数据搬到 GMEM。若把 arrive 放到函数末尾，TMEM 双缓冲的收益会被 TMA store 的延迟吃掉一大半。 + +### 11.5 C/D 环：不用 mbarrier，用 `wait_group.read` + +```cpp +if (epilogue_warp_idx == 0) cute::tma_store_wait(); // = wait_group.read 1 +cutlass::arch::NamedBarrier::sync(kNumUMMAStoreThreads, 0); +…写 SMEM… +cute::tma_store_fence(); +cutlass::arch::NamedBarrier::sync(kNumUMMAStoreThreads, 0); +if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + SM90_TMA_STORE_2D::copy(&tensor_map_cd, smem_base_ptr, n_idx, m_idx); // 或 SM90_TMA_REDUCE_ADD_2D + cute::tma_store_arrive(); // cp.async.bulk.commit_group +} +__syncwarp(); +``` + +- **两级同步分工**：`NamedBarrier`（`barrier.sync.aligned id, num_threads`，id = 0，只绑 `kNumUMMAStoreThreads` 个线程）把「TMA 读完 SMEM」的信息从 warp 0 广播给全体；写完后再用一次 NamedBarrier 确保所有人的 `st.shared` 都可见，才让 warp 0 的单一 lane 发 TMA。 +- **为什么只 warp 0 等**：`cp.async.bulk.wait_group` 是**每线程**的计数，而 `commit_group` 只由 warp 0 的 elected lane 发出，所以只有它持有非零的 group 计数，也只有它能等。 +- **`.read` 后缀的语义差别**：cute 的 `tma_store_wait()` 展开为 `cp.async.bulk.wait_group.read N`，它只等到「TMA 引擎不再需要读这块 SMEM」，**不**等数据真正落到 GMEM。这正是覆写 stage 所需的最弱条件，比无 `.read` 的版本（等全局可见）早很多。对比 [ptx/tma.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/ptx/tma.cuh) 里手写的 `ptx::tma_store_wait()`，它特意注释 *this function does not have `.read`*——本 epilogue 用的是 cute 那个带 `.read` 的版本。 +- **`kWithAccumulation`**：当调用方传了 `c`（`with_accumulation = c.has_value()`），TMA 指令从 `SM90_TMA_STORE_*` 换成 `SM90_TMA_REDUCE_ADD_*`（`cp.reduce.async.bulk.tensor.…add`），在**写回路径上做 GMEM 原子累加**，实现 `D += A@B` 而不需要先把 C 读进来。Batched 时同理换成 3D 变体。 +- **`epilogue_type_t::apply_index_n(n_idx)`**：本 kernel 固定传 `EpilogueIdentity`（恒等）。另一个可选的 `EpilogueHeadSplits` 用于 attention 场景，把 Q/K/V 三段拼在一起的 N 轴索引跳过中间的 head 段，要求三段都能被 `STORE_BLOCK_N` 整除。 + +### 11.6 dtype 转换 + +```cpp +// fp32 输出：直接存 +cute::SM100_TMEM_LOAD_32dp32b4x::copy(tmem_addr, v0, v1, v2, v3); +fence_view_async_tmem_load(); +ptx::st_shared(smem_ptr, v0, v1, v2, v3); // st.shared.v4.f32 + +// bf16 输出：读 8 个 FP32，两两 pack +DG_STATIC_ASSERT(kNumElemsPerBankGroup == 8 and is_same_v); +cute::SM100_TMEM_LOAD_32dp32b8x::copy(tmem_addr, v0..v7); +fence_view_async_tmem_load(); +ptx::st_shared(smem_ptr, cast_into_bf16_and_pack(v0,v1), cast_into_bf16_and_pack(v2,v3), + cast_into_bf16_and_pack(v4,v5), cast_into_bf16_and_pack(v6,v7)); // st.shared.v4.u32 +``` + +`cast_into_bf16_and_pack` 的实现是 `__float22bfloat162_rn({x, y})` 后重解释为 `int`——**一次指令完成两个 FP32 → BF16 的舍入与拼接**（round-to-nearest-even）。所以 8 个 FP32 累加器 → 4 个 32-bit → 一条 128-bit `st.shared.v4.u32`，寄存器压力与指令数都是最优的。 + +累加器永远是 FP32（`make_instr_desc`），BF16 只出现在输入与最终输出，不存在中间累加降精度。 + +--- + +## 12. 其他重要机制 + +### 12.1 PDL：把 prologue 藏进前驱 kernel 的尾巴 + +Programmatic Dependent Launch 允许后一个 kernel 在前一个 kernel **尚未执行完**时就被派发到 SM 上跑，只要它不碰前驱的输出。DeepGEMM 的用法是把 `cudaGridDependencySynchronize()`（PTX: `griddepcontrol.wait`）**刻意推迟**到 prologue 的最后一步： + +```cpp +line 102 kNumMulticast > 1 ? cluster_sync_with_relaxed_arrive() : void(); // 2-CTA TMEM alloc 前对齐 +line 110 if (warp_idx == 0) { prefetch_tma_descriptor(&tensor_map_a/b/cd); } // tensormap 进 L2 +line 117 shape_* = SHAPE_* != 0 ? SHAPE_* : shape_*; // 编译期常量覆写 +line 122 extern __shared__ __align__(1024) uint8_t smem_buffer[]; // SMEM 指针算术 +line 148 if (warp_idx == 1 && elect_one_sync()) { …32 个 mbarrier init…; fence_barrier_init(); } +line 168 else if (warp_idx == 2) { Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem); } +line 172 cluster_sync / __syncthreads(); +line 175 cudaGridDependencySynchronize(); ◄── 真正的依赖点 +line 194 …角色分派，第一条 TMA load… +``` + +被提到依赖点之前的四件事，全都**只读写本 kernel 自己的资源**，与前驱的数据无因果关系： + +| prologue 工作 | 为什么可以越过 PDL 边界 | +| --- | --- | +| TMA descriptor prefetch | 三个 `tensor_map_*` 是 `__grid_constant__` 按值传入的 kernel 参数，住在常量内存里，由 host 在 launch 前填好，不是前驱的 GMEM 输出 | +| mbarrier init（`3S + 4` 个） | 纯 SMEM 写，SMEM 随 CTA 分配，天然私有 | +| `tcgen05.alloc` | TMEM 是 SM 上的独立资源池，分配走硬件仲裁，不依赖任何 GMEM 状态 | +| cluster sync / `fence_barrier_init` | 只是 CTA 之间的汇合与可见性，不涉及数据 | + +而第一条 `cp.async.bulk.tensor` 要读的 A/B 张量**很可能就是前驱 kernel（如 RMSNorm、量化、all-gather）刚写出来的**，因此必须留在 `cudaGridDependencySynchronize()` 之后。 + +收益量化：冷启动路径上，32 个 mbarrier 的逐个 `mbarrier.init` + `fence.mbarrier_init.release.cluster`、`tcgen05.alloc.sync.aligned`（2-CTA 时还要等 peer CTA 一起），以及三个 tensormap 的 L2 miss，加起来通常是数百到上千周期。把它们与前驱 kernel 的 drain 阶段重叠，等于把 GEMM 的「首字节延迟」压掉一截——对小 K、块数少的 shape 尤其明显。 + +三个使用注意点： + +1. **默认关闭**。`DeviceRuntime::enable_pdl = false`（[device_runtime.hpp](../third_party/DeepGEMM/csrc/jit/device_runtime.hpp) 第 16 行），需要 `deep_gemm.set_pdl(True)` 才生效。`LaunchArgs` 构造函数的默认值虽然是 `enable_pdl = true`，但 `KernelRuntime::launch()` 会在发射前无条件覆写： + + ```cpp + // Allow runtime override from Python. + // NOTES: the default is enabled. + launch_args.enable_pdl = device_runtime->get_pdl(); // kernel_runtime.hpp:146 + ``` + + 注释里的「the default is enabled」指的是 `LaunchArgs` 那个默认实参，而**实际生效的是全局开关**，它默认是关的。这处注释与代码的不一致很容易误读。设 `DG_JIT_DEBUG=1` 可以在 launch 日志里直接看到 `pdl: 0/1`。 +2. **不开 PDL 也正确**。`enable_pdl == false` 时不加 `CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION`，grid 按流序在前驱完全结束后才启动，此时 `cudaGridDependencySynchronize()` 的依赖已天然满足、立即返回，所以这一行不需要条件编译。 +3. 全仓库**没有任何一处调用 `cudaTriggerProgrammaticLaunchCompletion()`**（`griddepcontrol.launch_dependents`）。也就是说 DeepGEMM 只做「等待方」，不做「提前放行方」——它依赖前驱 kernel 自然结束来放行自己，而不主动帮后继 kernel 提前启动。这在 GEMM 之间级联时意味着 PDL 的收益是单向的。 + +### 12.2 `cluster_sync_with_relaxed_arrive()` 与它的安全性前提 + +```cpp +// comm/barrier.cuh +CUTLASS_DEVICE void cluster_sync_with_relaxed_arrive() { + // This is slightly faster than `cute::cluster_sync` but has weaker memory ordering guarantee + cute::cluster_arrive_relaxed(); // barrier.cluster.arrive.relaxed + cute::cluster_wait(); // barrier.cluster.wait +} +``` + +`cute::cluster_sync()` 展开的 arrive 带默认（release）内存序，会把「arrive 之前本 CTA 的所有写」提升到 cluster 作用域可见。relaxed 版省掉这一层，只做**控制流的汇合**，不承诺任何内存可见性——快一点，但用错就是隐蔽的竞态。 + +kernel 里三处使用，各自的正确性依据并不相同，值得逐条看清： + +| 位置 | 时机 | 为什么 relaxed 够用 | +| --- | --- | --- | +| 第 102 行 | `tcgen05.alloc` **之前** | 2-CTA 的 TMEM 分配是硬件级成对操作，只要求两个 CTA 都到达该点；此处**尚无任何需要跨 CTA 可见的写**（barrier 还没 init），所以连 release 都是多余的 | +| 第 172 行 | barrier init + TMEM alloc **之后** | 这一处的跨 CTA 可见性**是必需的**（peer CTA 会通过 `umma_arrive_multicast` 与 `arrive(0u)` 远程投递到本 CTA 的 mbarrier）。但承担 release 的不是 cluster arrive，而是它上面第 167 行的 `cutlass::arch::fence_barrier_init()`（`fence.mbarrier_init.release.cluster`）——这条 fence 专门为「mbarrier 初始化对 cluster 内 async proxy 可见」设计，比通用 release 更精确也更便宜。cluster arrive 于是退化为纯汇合，可以安全 relaxed | +| 第 452 行 | 所有角色退出之后 | 配合第 392–397 行的「额外一轮 `tmem_empty_barriers` wait」，确保 peer CTA 不会再向本 CTA 发远程 arrive，然后才能 `Allocator().free(0, kNumTmemCols)`。此处同样只需汇合语义 | + +`kNumMulticast == 1` 时三处分别退化为 `void()` / `__syncthreads()` / `__syncthreads()`——注意第一处是 `void()` 而非 `__syncthreads()`，因为单 CTA 下 `tcgen05.alloc` 本来就不需要跨 CTA 对齐，而 CTA 内部此时也还没有任何需要全 block 可见的状态。 + +### 12.3 三层断言防御体系 + +[common/exception.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/common/exception.cuh) 定义了三个宏，分工非常清楚： + +```cpp +#define DG_STATIC_ASSERT(cond, ...) static_assert(cond, __VA_ARGS__) // 编译期，零成本 +#define DG_DEVICE_ASSERT(cond) do { if (not (cond)) { printf(…); asm("trap;"); } } while (0) +#define DG_TRAP_ONLY_DEVICE_ASSERT(cond) do { if (not (cond)) asm("trap;"); } while (0) +``` + +`DG_DEVICE_ASSERT` 会生成 `printf` 的完整调用序列（占用寄存器、拉进 vprintf 的常量字符串），在热路径上代价可观；`DG_TRAP_ONLY_DEVICE_ASSERT` 只留一条条件 `trap`，几乎免费但出错时只知道「哪次 launch 挂了」，不知道具体条件。 + +本 kernel 内三者的分布： + +| 断言 | 行 | 类别 | 检查内容 | +| --- | --- | --- | --- | +| `cd_dtype_t` 是 `float` / `bfloat16_t` | 56 | static | 输出 dtype 白名单 | +| `BLOCK_K_ == 64` | 65 | static | 一个 swizzle atom 的 K 字节数必须是 128 B | +| `BLOCK_K % UMMA_K == 0`、`kKAlignment % UMMA_K == 0` | 66–67 | static | K 方向能被 UMMA 原子（16）整除 | +| `kNumMulticast ∈ {1,2}` | 68 | static | cluster 最多 2 CTA | +| `(kSwapAB and BLOCK_N == LAYOUT_AD_M) or …` | 69 | static | swap-AB 时 `BLOCK_N` 必须恰好 128 | +| `kNumUMMAStoreThreads % 32 == 0` | 81 | static | epilogue 必须以整 warp 参与 | +| 三个 SMEM 区都是 1024 B 的倍数 | 88 | static | `__align__(1024)` + swizzle-128B 的前提 | +| `kNumTMAStoreStages >= 1` | 90 | static | C/D 环非空 | +| `UMMA_A_SIZE_PER_STAGE <= SMEM_A + SMEM_B * kNumStages` | 94 | static | UMMA 的越界读不会冲出 A/B 区（见 §5.5） | +| `32 <= kNumTmemCols <= 512` | 99, 145 | static | TMEM 分配的硬件合法区间 | +| `kGemmType` 与 `kMajorA` 的组合合法 | 216 | static | m-grouped 必须 A K-major | +| `kNumStages <= 32` | 264 | static | lane 寄存器存描述符的前提（见 §9.3） | +| UMMA shape 合法（`M∈{64,128,256}` 且 N 的粒度/范围匹配） | 274–277 | static | 替代 CUTLASS MMA traits 的手写校验 | +| `kTensorCoreUtilControl > 0` | 371 | static | 利用率控制的分母非零 | +| `ptx::ld_shared(tmem_ptr_in_smem) == 0` | 405 | **trap-only** | TMEM 分配基址列号必须为 0（禁止 2 CTA 共 SM） | +| `false and "This kernel only support sm_100f"` | 460 | device | 位于 `#else`（非 SM100）分支，SM100 cubin 里不存在 | + +一个值得注意的结论：**在真正编出来的 SM100 cubin 里，`DG_DEVICE_ASSERT` 一条都没有**，运行期检查只剩第 405 行那一条 trap。所有配置合法性都在 JIT 编译阶段就被 `static_assert` 拦住了——这也意味着 host 端启发式（§2.2）如果生成了非法配置，失败形式是**编译报错**而不是运行期错误，调试信息直接指向那条 `static_assert` 的字符串。 + +第 405 行选 trap-only 而非 printf 版，是因为它位于 epilogue warp 的入口、每次 launch 只执行一次，但作者仍希望它不影响这一带的寄存器分配（printf 版本会引入额外的活跃值）。运行期索引合法性检查（依赖运行期 shape 的那些）被下沉到库里，例如 `mma::sm100::make_umma_desc` 的 `DG_DEVICE_ASSERT(mn_idx % BLOCK_MN_ATOM == 0)`。 + +### 12.4 编译期常量折叠的收益链 + +JIT 全量模板特化不是「代码洁癖」，它触发一串连锁的死代码消除。以典型调用（`compiled_dims = "nk"`）为例： + +``` +SHAPE_K != 0 + └─► line 119: shape_k 变成编译期常量 + └─► line 313: kMayHaveTailKBlock = (SHAPE_K == 0 or SHAPE_K % BLOCK_K != 0) = constexpr false + └─► line 344–359: 整个 tail-K 分支不生成 SASS + ├─ for_each_static_prefix 的 fold expression 展开消失 + ├─ 为「≤4 用 switch」准备的跳转表消失 + └─ issue_tail_k_block lambda 消失 + └─► num_total_k_blocks = ceil_div(shape_k, BLOCK_K) 成为编译期常量 + └─► K 主循环可完全展开，循环边界比较折叠 + └─► BLOCK_K/UMMA_K = 4 条 UMMA 全部直线排布 + └─► 每条的 a_desc.lo / b_desc.lo 增量是立即数 + （SMEM_A_SIZE_PER_STAGE/16、advance_umma_desc_lo 的 offset） + └─► 描述符推进折叠成一条 IADD3 +``` + +同时被 `if constexpr` 彻底消除的分支维度：`kMajorA`（2）× `kMajorB`（2）× `kSwapAB`（2）× `kNumMulticast`（2）× `kWithAccumulation`（2）× `kIsBatchedMM`（2）× `kGemmType`（7）× `cd_dtype_t`（2）。**一份 cubin 里只存在一条完全直线化的路径**，没有任何「运行期再决定走哪边」的开销，指令 cache 压力也最小。 + +反过来说，`SHAPE_M` 默认**不**编译进来是有意的取舍：推理场景下 M（token 数）几乎每次都变，若把 M 也特化，JIT 缓存会爆炸式增长。代价是 M 方向的 tail 处理必须保留为运行期逻辑——`get_aligned_effective_m_in_block()`、`kEnsureZeroPadding`、swap-AB 下的动态 `load_block_m`、epilogue 里运行期的 `num_stores = effective_m / STORE_BLOCK_M`，全都是为这个「M 未知」付出的成本。 + +JIT 缓存的组织方式也值得一提：每个 (shape, config) 对应一个独立目录，`load_kernel()` 会校验 `cuLibraryGetKernelCount == 1`，否则打印 *Corrupted JIT cache directory … please run `rm -rf `* 并断言失败（[handle.hpp](../third_party/DeepGEMM/csrc/jit/handle.hpp) 第 144–150 行）。一个目录一个 kernel 的约定让缓存失效的判断非常简单。 + +### 12.5 `kEnsureZeroPadding` 的真实作用范围 + +这个模板参数名字很唬人，但它**只在一个地方被读取**——[scheduler/gemm.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/scheduler/gemm.cuh) 第 192 行： + +```cpp +CUTLASS_DEVICE uint32_t get_aligned_effective_m_in_block(const uint32_t& m_block_idx) const { + constexpr uint32_t UMMA_STEP_N = 16; + DG_STATIC_ASSERT(BLOCK_M % UMMA_STEP_N == 0, "Invalid alignment"); + if constexpr (kGemmType == GemmType::MGroupedContiguousWithPsumLayout and not kEnsureZeroPadding) + return math::align(…当前 psum 块的实际剩余行数…, UMMA_STEP_N); + return BLOCK_M; +} +``` + +即：只有在 **psum layout 的 m-grouped contiguous** 且显式关掉 zero-padding 时，最后一个 M 块的有效行数才会被收缩到实际值（对齐到 UMMA 的 N 步长 16）。其余所有情形恒返回 `BLOCK_M`——整个块照算，越界行的 A 由 TMA 的 `CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE` 零填充，于是 D 的 padding 行自然是 0。 + +两条路径的差别在于「padding 区是否被写」：收缩路径下 epilogue 的 `num_stores` 变小，padding 行的 D **完全不写**（保留 GMEM 原值）；零填充路径下整个 `BLOCK_M × BLOCK_N` 都会被 TMA store 覆盖成 0。Python API 默认 `ensure_zero_padding = true`。对 BF16 Normal GEMM 而言这个参数是纯 dead code，读代码时不必为它费神。 + +--- + +## 13. 不变式、限制与已知坑 + +### 13.1 必须成立的不变式 + +下表是「改动这个 kernel 或其 host 启发式时不能破坏」的硬约束。左列任一被打破，结果要么是编译失败，要么是静默的数值错误。 + +| 不变式 | 由谁保证 | 破坏后的表现 | +| --- | --- | --- | +| `BLOCK_K_ == 64` | host `block_k = 128 / element_size` | `static_assert`（第 65 行） | +| `kSwizzleMode == BLOCK_ATOM_K × sizeof(dtype)`（K-major） | `get_swizzle_mode()` | `make_umma_desc` 里的 `DG_STATIC_ASSERT("Unexpected value")`；SBO/LBO 全错 | +| `kNumStages <= 32` | host `num_stages = min(…, 32)` | `static_assert`（第 264 行）。真正的原因是 lane 寄存器只有 32 个，用来存 per-stage 描述符低位 | +| 三个 SMEM 区各自 1024 B 对齐且尺寸为 1024 的倍数 | `SMEM_*_SIZE_PER_STAGE` 的构造 | `static_assert`（第 88 行）；swizzle-128B 的 atom 跨边界 → 数据错乱 | +| `2 × UMMA_N <= 512` | host TMEM 容量过滤 | host 直接跳过该候选；若绕过则 `tcgen05.alloc` 失败 | +| `32 <= kNumTmemCols <= 512` | `get_num_aligned_tmem_cols` | `static_assert`（第 99/145 行） | +| UMMA base address 必须是 TMEM 第 0 列 | 1 CTA/SM + `__launch_bounds__(256,1)` | 第 405 行 `trap`。所有 TMEM 地址都省掉了 base 偏移 | +| `UMMA_M == LAYOUT_AD_M × kNumMulticast`，`LAYOUT_AD_M == 128` | TMEM datapath 固定 128 行 | UMMA shape 非法（第 274 行） | +| cluster 只能沿 layout A/D 方向 | host 过滤：`swap_ab && cluster_m > 1` 跳过；`!swap_ab && cluster_n > 1` 跳过 | `cta_group::2` 的 M 维拼接方向错 | +| `kNumSMs % cluster_size == 0` | host 过滤 | cluster 跨 grid 边界，launch 失败 | +| `ceil_div(m, BLOCK_M) % cluster_m == 0`（N 同理） | host 过滤 | cluster 内两 CTA 落到不同边界外，`arrive` 计数永远凑不齐 → 死锁 | +| `kNum1DBlocksPerGroup % kNumMulticast == 0` | host `get_num_1d_blocks_per_group()` | cluster 内两 CTA 分到不同 L2 组，multicast 的收益归零 | +| `128 <= LOAD_BLOCK_M + LOAD_BLOCK_N × kNumStages` | 第 94 行断言的化简形式 | UMMA 读 A 时越过 B 区末尾，读到未初始化 SMEM | +| swap-AB ⇒ `BLOCK_N == 128` 且 `kSwizzleCDMode == 128` | host 过滤 + epilogue 的 `static_assert` | `STORE_BLOCK_N` 与 TMEM 行数不匹配 | +| m-grouped ⇒ A 必须 K-major | host `DG_HOST_ASSERT` + 第 216 行 `static_assert` | `get_global_idx` 的 group 偏移公式失效 | +| k-grouped contiguous ⇒ A/B 都必须 MN-major | host `DG_HOST_ASSERT`（impls 第 266 行） | K 方向的 group cumsum 索引算错 | +| `kNumEpilogueStages == 2` | 硬编码 | 退出协议（§10.6）与 `accum_phase_idx` 的 `& 1` 推导全部失效 | +| 每个 barrier 的 `init` 计数与 arrive 次数严格相等 | §10.3 的逐条推导 | 计数多 → 永久等待（死锁）；计数少 → parity 提前翻转（数据竞争） | + +### 13.2 功能与平台限制 + +- **仅 SM100**。`#if __CUDA_ARCH__ >= 1000` 之外只有 `DG_DEVICE_ASSERT(false and "This kernel only support sm_100f")`。SM90 走的是同仓库另一套 kernel（`sm90_*`），两者的流水组织完全不同（SM90 是 warpgroup MMA + `wgmma.mma_async`，没有 TMEM）。 +- **dtype 硬编码 BF16**。`smem_a` / `smem_b` 一律 cast 成 `cutlass::bfloat16_t*`，`make_instr_desc`。FP8/FP4 走 `sm100_fp8_fp4_gemm_1d1d.cuh`（需要 scale factor，流水结构不同）。输出 `cd_dtype_t` 只有 `float` 与 `bfloat16_t` 两种，没有 FP16。 +- **cluster 最多 2 CTA**（`kNumMulticast ∈ {1,2}`），不支持 SM100 理论上的更大 cluster 做 multicast。 +- **无 split-K**。K 方向由单个 CTA 串行跑完整个 `num_total_k_blocks`，K 很大而 M/N 很小的瘦长 shape 无法靠增加并行度填满 SM。`kWithAccumulation`（`cp.reduce.async.bulk.tensor…add`）提供了「多次调用累加到同一块 D」的能力，算是 host 层面的手工 split-K，但每次调用之间没有 kernel 内同步。 +- **persistent 调度是静态的**。`next_block_idx = (++current_iter) * kNumSMs + blockIdx.x`，纯算术映射，没有原子操作也没有工作窃取。好处是零同步开销、每个 CTA 都能独立推算出自己的块序列；代价是**块间代价不均时无法再平衡**——`MGroupedMasked` 下各 group 的 `masked_m` 差异很大时，末 wave 的长尾会直接暴露在关键路径上（host 的 `compare()` 里 `last_wave_util` 那一项就是在缓解它，但只能缓解不能消除）。 +- **D 必须 N-major**（`check_major_type_cd`）。想要列主序输出只能在 host 侧转置。 +- **`gridDim.x == num_sms`**，且 `deep_gemm.set_num_sms()` 会同时影响 grid 大小与调度器的 `kNumSMs`；如果设置成小于物理 SM 数，会有 SM 闲置；两者必须一致，否则块分配会漏。 + +### 13.3 已知坑与代码瑕疵 + +按「踩到的概率 × 排查成本」排序： + +1. **`kTensorCoreUtilControl < 100` 会主动降速**（§9.7）。这是给功耗受限集群做可复现对比用的旋钮，通过 `deep_gemm.set_tc_util(n)` 打开，默认 100。跑分或做性能回归前必须确认它是 100——否则 kernel 会在每个 stage 后 `clock64()` 自旋 `kNumDummyCycles`，测出来的 TFLOPS 毫无意义，而且现象很像「kernel 莫名变慢」，很难往这个方向想。 + +2. **`BLOCK_M ∈ {32, 64}` 时 MMA 仍按 M=128 发射**（§7.6）。`UMMA_M = LAYOUT_AD_M × kNumMulticast` 与 `BLOCK_M` 无关，所以小 `BLOCK_M` 的配置存在**有意的算力浪费**：TMEM 的 128 个 datapath 里只有 `BLOCK_M` 个装着有用数据。换来的是单一代码路径 + 避免小 M 块的 TMA L2 OOB。看 profiler 里「MMA 指令数 / 有效 FLOP」比值异常时，先确认 `BLOCK_M`。 + +3. **barrier 区有 `kNumStages` 大小的空洞**（§5.2）。`tensor_core_full_barrier` 的索引是 `3S + 4` 而不是 `2S + 4`，中间空出的 S 个 `Barrier` 是 FP8/FP4 1D1D kernel 的「每 stage 三组 barrier」约定（第三组是 with-SF full barriers）的遗留占位。host 侧 `smem_barriers = 32*8*3 + 2*8*2 + 8` 与它精确呼应，所以 SMEM 预留量对 `kNumStages <= 32` 恰好是**紧上界**。如果有人「优化」掉这个空洞但忘了同步改 host 的 `3`，在 `kNumStages` 接近 32 时就会写穿 SMEM。 + +4. **第 218 行 `uint32_t k_idx = k_block_idx * BLOCK_K;` 是未使用的残留变量**。真正传给 TMA 的是 `k_a_idx` / `k_b_idx`（由 `get_global_idx` 按 major 分别算出）。读代码时容易误以为 `k_idx` 参与了地址计算。 + +5. **第 451 行 `// TODO: Remove redundant synchronization`**。退出前的那次 cluster sync 作者自己也怀疑多余。但它目前是 `Allocator().free(0, kNumTmemCols)` 的唯一安全网（保证 peer CTA 不再引用本 CTA 的 TMEM / barrier），不要贸然删除。 + +6. **peer CTA 上有一批死对象**。`block_rank_in_cluster() != 0` 的 CTA 里：warp 1 完全不执行 MMA 分支（但它 prologue 里 init 的那些 `full_barriers` / `empty_barriers` 是**会被 leader CTA 远程 arrive 的**，不是死的）；真正死掉的是 `tmem_full_barriers`——它 `init(1)`，只由本 CTA 的 `umma_arrive` 触发，而 peer CTA 从不发 UMMA。用 nsys / compute-sanitizer 观察 barrier 计数时，不要把这些恒零的对象当成 bug。 + +7. **host 的 `store_block_n` ≠ kernel 的 `STORE_BLOCK_N`**（§2.3）。host 填 `layout.block_n`，kernel 非 swap 分支算的是 `kSwizzleCDMode / sizeof(cd_dtype_t)`。两者最终一致只是因为 `make_tma_2d_desc()` 里 `smem_inner_dim = swizzle_mode / elem_size` 又把 host 的值覆盖了一遍。对不上号时先想到这层覆盖。 + +8. **大量 warp 在空转**。256 线程 / 8 warp 里，真正干活的是：TMA 1 lane + MMA 1 lane（`commit` 时整 warp）+ TMEM alloc 时 32 lane（一次性）+ epilogue `kNumUMMAStoreThreads`。warp 3 恒空转；`BLOCK_M = 32` 时 warp 5–7 也空转。极端配置下 256 个线程里稳态活跃的只有约 34 个。这不是 bug（warp specialization 的固有形态），但会严重误导「占用率」类的性能分析——具体该看哪些指标、为何 `sm__warps_active` 在此无意义，见 **§3.4.6**。 + +9. **`__shfl_sync(0xffffffff, …)` 要求整 warp 收敛**。第 324 行的 `elect_one_sync()` 与 `__shfl_sync` 都在 `full_barriers[...]->wait()` 之后、warp-uniform 的分支内部，所以合法。若将来有人在 MMA warp 里引入依赖 lane 的分支（例如按 `lane_idx` 走不同路径），这两处会立刻变成未定义行为。 + +10. **`tcgen05.commit` 是 warp 级收敛指令**（§10.3）。`umma_arrive` 在 `elect` 块**之外**被整个 warp 1 调用，而 `empty_barriers[i]->init(1)`——这只有在「整 warp 执行 commit 只产生一次 arrival」的前提下才成立。任何试图把 `umma_arrive` 挪进 `elect_one_sync()` 块、或把 `init(1)` 改成 `init(32)` 的「修正」都会破坏协议。 + +--- + +## 14. 小结 + +这份 kernel 的核心思想可以压缩成三句话： + +1. **把同步全部外化成 mbarrier 的 parity 相位**。没有 `__syncthreads()` 出现在稳态路径上，五类 barrier（`full` / `empty` / `tmem_full` / `tmem_empty` / `tensor_core_full`）各自表达一条单向的生产-消费边，每个角色只等自己需要的那一条，因此三级流水（SMEM ring → TMEM 双缓冲 → C/D SMEM ring）能各自独立超前。 +2. **把能变成编译期常量的东西全部变成编译期常量**。JIT 全量特化换来的是内层 K 循环完全展开、描述符增量成即数、tail-K 整段消失、所有 `if constexpr` 分支塌缩成一条直线路径——这才是 DeepGEMM 相对通用库的真正护城河，比任何单个 PTX 技巧都重要。 +3. **让专用硬件做它最擅长的事**。TMA 负责所有 GMEM↔SMEM 搬运（含 swizzle、越界零填充、multicast、reduce-add），`tcgen05.mma` 负责所有乘加（操作数直接来自 SMEM，累加器直接在 TMEM），`tcgen05.ld` / `stmatrix` 负责 TMEM→RF→SMEM，通用寄存器与 `__syncthreads()` 只承担控制流。整个数据通路上，通用 SIMT 部分几乎不碰数据。 + +理解这三点之后，其余所有细节——为什么 `UMMA_M` 恒为 128、为什么用 lane 寄存器存描述符、为什么 TMEM 要早释放、为什么 barrier 区有个空洞、为什么 `cudaGridDependencySynchronize()` 放在第 175 行——都是它们的自然推论。 diff --git a/docs/deepgemm_sm90_bf16_gemm_design.md b/docs/deepgemm_sm90_bf16_gemm_design.md new file mode 100644 index 000000000..7e9388be0 --- /dev/null +++ b/docs/deepgemm_sm90_bf16_gemm_design.md @@ -0,0 +1,1619 @@ +# DeepGEMM `sm90_bf16_gemm` Kernel Detailed Design + +面向 Hopper（SM90，H100/H800）的 BF16 GEMM 内核详细设计。本文覆盖：warp specialization 与 pipeline 组织、多级 tiling、SMEM 排布、**寄存器累加器**分片、warpgroup MMA（`wgmma.mma_async`）调用方式、生产者-消费者 mbarrier 同步协议、TMA 指令（含 SM90 multicast）的发射与完成语义，以及若干容易被忽略但对正确性/性能关键的设计点。 + +> 本文与 [deepgemm_sm100_bf16_gemm_design.md](./deepgemm_sm100_bf16_gemm_design.md) 是姊妹篇。两者共享同一套 host 启发式框架、同一个 persistent 调度器、同一份 TMA 封装，但 **device 侧的流水组织因架构差异而完全不同**：SM90 用 warpgroup MMA + 寄存器累加器，SM100 用单线程 UMMA + Tensor Memory。凡涉及两代差异处，本文都会显式对照，并在 §1.4 给出一张总表。 + +## 0. Code Index + +| 层次 | 文件 | 职责 | +| --- | --- | --- | +| Device 主体 | [sm90_bf16_gemm.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh) | kernel 本体：SMEM 布局、TMA/math 两类 warpgroup、流水推进、epilogue | +| WGMMA 描述符 | [mma/sm90.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/mma/sm90.cuh) | `BF16MMASelector`、`GmmaDescriptor` 构造、SBO/LBO 推导、K 方向描述符推进 | +| wgmma PTX | [ptx/wgmma.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/ptx/wgmma.cuh) | `wgmma.fence/commit_group/wait_group`、累加器操作数 fence | +| TMA load | [common/tma_copy.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/common/tma_copy.cuh) | swizzle-atom 循环、1CTA / SM90-multicast / SM100-2SM 分支 | +| LD/ST & STSM | [ptx/ld_st.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/ptx/ld_st.cuh) | `SM90_U32x2_STSM_N`、`st_shared`、`mapa_shared` | +| 调度器 | [scheduler/gemm.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/scheduler/gemm.cuh) | persistent 块分配、L2 swizzle、SM90 multicast 合法性、grouped/batched 索引 | +| Epilogue transform | [epilogue/transform.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/epilogue/transform.cuh) | `EpilogueIdentity` / `EpilogueHeadSplits` 的 N 索引映射 | +| 通用工具 | [common/utils.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/common/utils.cuh)
[common/math.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/common/math.cuh) | `PatternVisitor`、编译期循环展开、`ceil_div`/`align`、`cast_into_bf16_and_pack` | +| Cluster 同步 | [comm/barrier.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/comm/barrier.cuh) | `cluster_sync_with_relaxed_arrive` | +| Host JIT | [impls/sm90_bf16_gemm.hpp](../third_party/DeepGEMM/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp) | 模板实参拼装、TMA descriptor 构造、launch（5 个入口 API） | +| Host 启发式 | [heuristics/sm90.hpp](../third_party/DeepGEMM/csrc/jit_kernels/heuristics/sm90.hpp)
[heuristics/common.hpp](../third_party/DeepGEMM/csrc/jit_kernels/heuristics/common.hpp)
[heuristics/utils.hpp](../third_party/DeepGEMM/csrc/jit_kernels/heuristics/utils.hpp)
[heuristics/config.hpp](../third_party/DeepGEMM/csrc/jit_kernels/heuristics/config.hpp) | BLOCK_M/N/K、cluster、swizzle、stage 数、线程数推导、L1/L2 周期打分 | +| TMA desc 构造 | [impls/runtime_utils.hpp](../third_party/DeepGEMM/csrc/jit_kernels/impls/runtime_utils.hpp) | `make_tma_{a,b,cd,3d}_desc`、`get_compiled_dim`、swizzle→tensormap 映射 | + +> **注**：本仓库的 `third_party/DeepGEMM/third-party/cutlass` 子模块为空目录，CUTLASS/CuTe 头文件未随仓库落地。文中涉及 `cute::GmmaDescriptor` 的精确位域、`cute::SM90::GMMA::MMA_64xNx16_F32BF16BF16_SS::fma` 展开出的 `wgmma.mma_async` PTX 文本、`cutlass::arch::warpgroup_reg_{alloc,dealloc}` 的 `setmaxnreg` 指令、`SM90_TMA_LOAD_MULTICAST_2D` 的 `cp.async.bulk.tensor` 修饰串等，均标注为「依据 CUTLASS 约定 / 由本仓库调用方式反推」，不做逐字断言。 + +--- + +## 1. 设计总览 + +### 1.1 一句话概括 + +一个 **persistent + warp-specialized + TMA 全异步、但 MMA 半同步** 的两段流水线内核。数据通路为： + +``` +GMEM ──(cp.async.bulk.tensor / TMA，可 multicast)──► SMEM(A,B ring) + │ + ▼ wgmma.mma_async.sync.aligned.m64nNk16（整 warpgroup 协同） + RF(累加器 accum[]，分散在 128 线程) + │ + ▼ stmatrix.x2.b16 (bf16) / st.shared.v2.f32 (fp32) + SMEM(C/D 单缓冲) + │ + ▼ cp.async.bulk.tensor store / reduce.add + GMEM(D) +``` + +与 SM100 最大的不同在于中段：**SM90 没有 Tensor Memory，累加器从头到尾住在 math warpgroup 128 个线程的寄存器里**。这带来两个连锁后果： + +1. `wgmma.mma_async` 是 **warpgroup 级协同指令**，128 个线程必须全部参与发射（不是 SM100 的单 lane 控制指令）——因为每个线程都要「拿着」自己那一份累加器寄存器。 +2. **发射 MMA 的 warpgroup、持有累加器的 warpgroup、做 epilogue 的 warpgroup 是同一批线程**。没有 SM100 那种「MMA warp / epilogue warp 分离 + TMEM 双缓冲」的重叠，累加器的生命周期把一个输出块的 compute 与 epilogue 串在一起。 + +A/B 的 GMEM→SMEM 搬运仍然全异步（TMA + mbarrier），这一段与 SM100 同构，可以超前 MMA 多达 `kNumStages` 个 k_block。 + +### 1.2 框图（单个 CTA 内） + +以典型配置 `BLOCK_M=128`（⇒ `kNumMathThreads=256`，2 个 math warpgroup）为例，共 384 线程 / 12 warp： + +``` + ┌──────────────────── kNumTMAThreads(128) + kNumMathThreads(256) = 384 threads ───────────────────┐ + │ │ + math warp-group 0 (warp0..3, 128 thr) math warp-group 1 (warp4..7, 128 thr) TMA warp-group (warp8..11) │ + ┌───────────────────────────────────┐ ┌───────────────────────────────────┐ ┌──────────────────────────────┐│ + │ WGMMA + EPILOGUE │ │ WGMMA + EPILOGUE │ │ w8: prefetch tensormap (1 lane)││ + │ rows 0..63 (math_wg_idx=0) │ │ rows 64..127 (math_wg_idx=1) │ │ w9: init barriers (1 lane) ││ + │ accum[64] in RF │ │ accum[64] in RF │ │ w10: TMA LOAD k-loop(1 lane) ││ + │ reg_alloc<224> │ │ reg_alloc<224> │ │ w11: 空转 ││ + └──────┬──────────────────▲──────────┘ └──────┬──────────────────▲──────────┘ │ reg_dealloc<48> ││ + │ wait full[s] │ wait full[s] │ │ └──────┬──────────────▲────────┘│ + ▼ │ ▼ │ wait │ empty[s] │ │ + ┌────────────────────────┴──────────────────────────────────────────┴──────────────────── ▼ ─────────────┴───────┐ │ + │ SMEM: D(单缓冲) │ A ring(kNumStages) │ B ring(kNumStages) │ full/empty barriers │ │ + └───────────────────────────────────────────────────────────────────────────────────────────────────────────────┘ │ + └─────────────────────────────────────────────────────────────────────────────────────────────────┘ +``` + +- **math warpgroup**（`warp_idx < kNumMathThreads/32`）：既发 `wgmma`、又持有累加器、又做 epilogue（STSM/st.shared → SMEM → TMA store）。`BLOCK_M ≤ 64` 时只有 1 个 math warpgroup（128 线程），`BLOCK_M > 64` 时有 2 个（256 线程），每个 warpgroup 负责 `WGMMA::M = 64` 行。 +- **TMA warpgroup**（`warp_idx ≥ kNumMathThreads/32`，恒 128 线程 = 4 warp）：只有第三个 warp（w10）的单个 lane 真正发 TMA；w8 prefetch descriptor、w9 初始化 barrier、w11 空转。选第三个 warp 的原因见 §3.1 注释——`BLOCK_M==32` 时 warp0/1 可能正忙于 WGMMA。 + +注意 w10 发 TMA 用的是 **单个 elected lane**（`cute::elect_one_sync()`），这与 SM100 一致：TMA 本身是单线程指令。真正与 SM100 不同的是 **MMA 侧**：SM90 的 math warpgroup 是 128 个线程一起发 `wgmma`，而 SM100 是 1 个 lane 发 `tcgen05.mma`。 + +### 1.3 关键设计选择 + +| 设计点 | 取值 | 理由 | +| --- | --- | --- | +| 编译方式 | 每个 (shape, config) 组合 JIT 生成一份全量模板特化的 `.cu` 编到 cubin | BLOCK_*、swizzle、stage 数、甚至 N/K 本身都成为编译期常量 → 内层 K 循环完全展开、描述符偏移常量折叠 | +| Grid | `gridDim.x = num_sms`，`__launch_bounds__(kNumTMAThreads+kNumMathThreads, 1)` | persistent kernel，1 CTA/SM；配合近满额 SMEM，物理上排除 2 CTA 共 SM | +| Cluster | 1 或 2（`cluster_m×cluster_n ≤ 2`） | 用 **TMA multicast** 让 pair 内两个 CTA 共享一份 A 或 B，省 L2/GMEM 带宽（**注意：不省 SMEM，见 §7.2**） | +| MMA 指令 | `wgmma.mma_async.sync.aligned.m64nNk16.f32.bf16.bf16`（`_SS` 变体） | Hopper warpgroup MMA；A/B 均来自 SMEM descriptor，累加器在 **寄存器** | +| `WGMMA::M` | 恒为 **64** | Hopper wgmma 的 M 原子固定 64（一个 warpgroup 4 warp × 16 行）；`BLOCK_M ∈ {16,32}` 时也照发 M=64（见 §7.5） | +| `WGMMA::K` | 恒为 **16** | BF16 的 wgmma K 原子固定 16 | +| `WGMMA::N` | `= BLOCK_N ∈ {8,16,…,256}` | 由 `BF16MMASelector` 选出对应的 `MMA_64xNx16` | +| 累加器 | 寄存器数组 `accum[kNumAccum × waves]`，`kNumAccum = 64×BLOCK_N/128 = BLOCK_N/2` | 没有 TMEM；每线程持 `BLOCK_N/2` 个 FP32（×wave 数） | +| C/D 缓冲 | SMEM **单缓冲**（`SMEM_D_SIZE`） | 靠 `tma_store_wait<0>()` 串行复用；不像 SM100 双缓冲重叠 | +| A/B 环 | 尽可能多的 stage（host 端按 SMEM 预算反解，上限 **16**） | 掩盖 HBM 延迟；stage ≥ 10 且 NT-Normal-单 math warpgroup 时再做「stage 合并」把 `BLOCK_K` 放大（见 §7.4） | +| 寄存器再分配 | TMA warpgroup `dealloc<48>`，math warpgroup `alloc<224 或 248>` | 累加器吃寄存器，把 TMA warpgroup 的配额让给 math warpgroup | +| PDL | `cudaGridDependencySynchronize()` 放在 barrier init + cluster sync **之后** | barrier 初始化、descriptor prefetch 与前驱 kernel 的尾巴重叠 | +| 不支持 | swap-AB、Tensor Core 利用率控制、tail-K 专用分支 | SM90 路径明确断言 `swap_ab==0`；tail-K 靠 TMA 零填充自然处理（见 §7.6） | + +### 1.4 SM90 ↔ SM100 对照总表 + +这张表是理解本 kernel 的钥匙——很多 SM90 的「为什么这么写」只有在与 SM100 对照时才清晰。 + +| 维度 | SM90（本文） | SM100（姊妹篇） | +| --- | --- | --- | +| Tensor Core 指令 | `wgmma.mma_async.m64nNk16` | `tcgen05.mma.cta_group::{1,2}.kind::f16` | +| 发射宽度 | **整 warpgroup（128 线程）协同** | **单线程**（ elected lane） | +| A 操作数 | SMEM descriptor（`_SS`） | SMEM descriptor（`_SS`） | +| B 操作数 | SMEM descriptor | SMEM descriptor | +| **D 累加器** | **寄存器**（分散在 128 线程，每线程 `BLOCK_N/2` 个 f32） | **TMEM**（独立 256 KB 存储，128 行 × 512 列） | +| MMA 的 M | 固定 64 | `128 × kNumMulticast`（可达 256） | +| 为什么需要那么多线程 | 累加器每个元素都得有个线程「拿着」 | 没人需要「拿着」任何东西 | +| compute 与 epilogue | **同一批 math 线程串行**（累加器占用寄存器整个 k-loop） | MMA warp 与 epilogue warp **分离**，TMEM 双缓冲重叠 | +| cluster 协作 | TMA **multicast**（复制操作数，省带宽不省 SMEM，两 CTA 各算各的块） | **2-CTA UMMA**（切分操作数，省带宽也省 SMEM，两 SM tensor core 合算一块） | +| C/D SMEM | 单缓冲 | 双缓冲（`kNumTMAStoreStages=2`） | +| 最大 stage 数 | 16 | 32 | +| 描述符 per-stage 存储 | 单个 `a_desc_lo`/`b_desc_lo`，每 stage 加常量步长 | 32 个 lane 各存一个 stage 的描述符低位，`__shfl` 取用 | +| barrier 种类 | 2 类（`full` / `empty`） | 5 类（`full`/`empty`/`tmem_full`/`tmem_empty`/`tensor_core_full`） | +| swap-AB | 不支持（`DG_HOST_ASSERT(swap_ab==0)`） | 支持（MoE 小 M 场景） | +| 寄存器再配置 | 有（`setmaxnreg`） | 无（TMEM 不吃寄存器） | +| **线程填充率是否是有效指标** | **是**——math warpgroup 的 128/256 线程稳态都在算或搬 | **否**——256 线程稳态仅约 34 活跃，须看 tensor pipe 指标 | + +最后一行尤其重要：SM90 的 `wgmma` 要 128 线程是**存储约束**（累加器在寄存器）而非算力约束，但这些线程在 k-loop 里确实都参与发射、在 epilogue 里确实都参与搬运，所以「活跃线程数」对 SM90 是有意义的健康指标；而对 SM100 用同一指标会得出完全错误的「低效」结论。评估 SM90 时仍可结合 `sm__pipe_tensor_cycles_active`，但 `sm__warps_active` 不再像 SM100 那样具有误导性。 + +--- + +## 2. Host 侧：JIT 特化与配置推导 + +Kernel 的全部行为由 22 个模板实参决定，它们在 [sm90_bf16_gemm.hpp](../third_party/DeepGEMM/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp) 的 `generate_impl()` 里被格式化成一段只包含 `__instantiate_kernel()` 的 `.cu` 源码，再交给 nvcc 编成 cubin。 + +### 2.1 模板实参来源 + +device 侧模板签名（[sm90_bf16_gemm.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh) 第 28–38 行）： + +```cpp +template +``` + +各实参在 host 侧的来源（`generate_impl` 第 50–65 行）： + +``` +kMajorA / kMajorB ← to_string(gemm_desc.major_a/b)：由 a/b 的 stride 推断 +SHAPE_M / SHAPE_N / SHAPE_K ← get_compiled_dim(dim, 'm'/'n'/'k', compiled_dims)： + 在 compiled_dims 里 → 填真实值；否则填 0（= 运行期参数） +kNumGroups ← gemm_desc.num_groups +BLOCK_M / BLOCK_N / BLOCK_K_ ← gemm_config.layout.{block_m, block_n, block_k} +kSwizzle{A,B,D}Mode ← gemm_config.storage_config.swizzle_{a,b,cd}_mode +kNumStages_ ← gemm_config.pipeline_config.num_stages +kNumTMAThreads ← launch_config.num_tma_threads（恒 128） +kNumMathThreads ← launch_config.num_math_threads（block_m≤64 ? 128 : 256） +kNumTMAMulticast ← layout.get_cluster_size()（= cluster_m × cluster_n） +kIsTMAMulticastOnA ← (layout.cluster_n > 1) +kNumSMs ← launch_config.num_sms（= gridDim.x） +kGemmType / kWithAccumulation / cd_dtype_t ← gemm_desc +``` + +与 SM100 的 26 个实参相比，SM90 **少了** `kSwapAB`、`kEnsureZeroPadding`、`kKAlignment`、`kTensorCoreUtilControl`、`kNumNonEpilogueThreads/kNumEpilogueThreads` 这一批——因为 SM90 路径不支持 swap-AB、没有 tensor core 利用率旋钮、epilogue 与 math 是同一批线程。**多了** `kNumTMAThreads/kNumMathThreads`（SM100 用的是 `kNumNonEpilogueThreads/kNumEpilogueThreads`，两套线程划分模型不同）。 + +Python 侧 `compiled_dims` 默认 `"nk"`（grouped contiguous/masked 也是 `"nk"`，两个 batched einsum 变体是 `"mn"`）。因此典型场景下 `SHAPE_M == 0`、`SHAPE_N`/`SHAPE_K` 为编译期常量，kernel 内第 68–70 行的覆写： + +```cpp +shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; +shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n; +shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k; +``` + +会把 N/K 变成常量。注意 SM90 **没有** SM100 那个 `kMayHaveTailKBlock` 的编译期分支——即便 K 是常量，SM90 也不生成专门的 tail-K 代码，而是靠 TMA 的 OOB 零填充把不足 `BLOCK_K` 的尾巴补 0（见 §7.6）。这是 SM90 相对 SM100 少掉的一整块复杂度。 + +### 2.2 Layout 候选枚举 + +`SM90ArchSpec::get_layout_candidates()`（[heuristics/sm90.hpp](../third_party/DeepGEMM/csrc/jit_kernels/heuristics/sm90.hpp) 第 16–118 行）： + +- **`block_k` 恒定**：`128 / element_size(BF16=2) = 64`。即 `BLOCK_K_ == 64`，物理含义是一个 swizzle atom 的 K 方向字节数固定 128 B。 +- **`block_m` 候选**（按 GemmType 分流）： + - Normal / Batched / KGroupedContiguous：基础 `{64, 128}`；`m ≤ 16` 追加 `16`，`m ≤ 32` 追加 `32`（注释：*smaller block M can avoid TMA L2 OOB bound*）；若输出非 FP32 再追加 `256`（*BF16 output GEMM supports 256*）。 + - MGroupedContiguous{,WithPsumLayout}：强制 `= get_mk_alignment_for_contiguous_layout()`。 + - MGroupedMasked：`{64, 128}`。 +- **`block_n` 候选**：步长 `lcm(16, block_n_multiple_of)`，从 `step` 到 `end` 枚举；`end` 按 kernel 类型收紧（NoSF→256、1D2D→192、1D1D→160，注释：*Register spills*）。BF16 GEMM 走 NoSF，故 `block_n ∈ {16,32,…,256}`。 +- **`disable_multicast`**：k-grouped 且 `num_groups > 4`，或 Batched → 禁用 multicast（`cluster` 只枚举到 1）。 +- **逐条过滤**（第 71–110 行）： + - `cluster_m × cluster_n > 2`、`num_sms % cluster_size != 0` → 跳过。 + - `block_m > 128 && block_n > 128` → 跳过（*for enough registers, at least one dim less than 128*）——这是寄存器容量约束，因为累加器吃 `block_n/2 × waves` 个寄存器。 + - masked / psum 布局要求 `ceil_div(n, block_n) % cluster_size == 0`（multicast 合法性）。 + - `swizzle_a_mode % 64 != 0 || swizzle_b_mode % 64 != 0` → 跳过（*32B's performance is low*）。 + - **stage 数下限**：`num_stages < 3`，或（`block_m×block_n < 128×192` 且 `num_stages < 4`）→ 跳过（*To hide TMA latency*）。 +- **打分 `compare()`**：SM90 用的是一个**解析带宽模型**（`get_layout_info`，第 201–238 行），而非 SM100 的「wave 优先」字典序： + +```cpp +num_bytes_l2_ab = expected_k * (block_m/cluster_n + block_n/cluster_m) * elem_ab; // multicast 省 L2 +num_bytes_l1_ab = expected_k * (block_m + block_n) * elem_ab; +num_bytes_l1_tc = expected_k * (max(64, block_m) + block_n) * elem_ab + block_m*block_n*elem_cd; +num_l2_cycles = (num_bytes_l2_ab + num_bytes_l1_l2_cd) * num_blocks / l2_bw_per_cycle; +num_l1_cycles = (num_bytes_l1_ab + num_bytes_l1_tc + num_bytes_l1_l2_cd) * num_blocks / l1_bw_per_cycle; +num_cycles = max(num_l1_cycles, num_l2_cycles) / wave_efficiency; +// compare: a.num_cycles < b.num_cycles +``` + +即 SM90 显式建模 L1（128 B/cycle/SM）与 L2（`min(64×num_sms, 8e6/1.3e3)` B/cycle）带宽，取二者瓶颈周期、再除以 wave 效率作为代价。`num_bytes_l2_ab` 里的 `block_m/cluster_n + block_n/cluster_m` 正是 **multicast 把某一维的 L2 流量对半**的体现。若只有 1 个 wave（`num_waves ≤ 1`），multicast 直接被判定为无穷大代价（*Disable multicasting if only one wave exists*）——单 wave 时 multicast 没有跨块复用，只剩开销。 + +### 2.3 StorageConfig + +```cpp +constexpr int wgmma_m = 64; +DG_HOST_ASSERT(layout.swap_ab == 0); // SM90 不支持 swap-AB +load_block_m = layout.block_m; // 注意：不除 cluster！ +load_block_n = layout.block_n; +store_block_m = (kernel_type == 1D1D) ? 64 : layout.block_m; // BF16 GEMM 走 block_m +store_block_n = layout.block_n; + +swizzle_mode_a = get_swizzle_mode(major_a == K ? block_k : load_block_m, sizeof(a)); +swizzle_mode_b = get_swizzle_mode(major_b == K ? block_k : load_block_n, sizeof(b)); +swizzle_mode_cd = (cd_dtype != float) ? get_swizzle_mode(store_block_n, sizeof(cd)) : 0; +``` + +两个与 SM100 的关键差异： + +1. **`load_block_m/n` 不除以 cluster**。SM100 是 `load_block_m = block_m / cluster_n`（2-CTA UMMA 各存一半）；SM90 是 `load_block_m = block_m`（multicast 各存**整份**）。这正是 §1.4「multicast 省带宽不省 SMEM」的根源，详见 §7.2。 +2. **FP32 输出时 `swizzle_mode_cd = 0`**（*We only enable swizzling for non-FP32 outputs*）。因为 FP32 走 `st.shared.v2.f32` 而非 STSM，且 `TMA_D_BLOCK_N` 在 `kSwizzleDMode==0` 时退化为整个 `BLOCK_N`（单条 TMA store）。BF16 输出时 `swizzle_cd = get_swizzle_mode(block_n, 2)`，`block_n ≥ 64` 时恒为 128。 + +`get_swizzle_mode()`（[heuristics/utils.hpp](../third_party/DeepGEMM/csrc/jit_kernels/heuristics/utils.hpp)）从 `{128,64,32,16}` 里挑第一个能整除 `block_size × elem_size` 的值。BF16 + `block_k=64` → 128 B，恒为 `swizzle=128`。 + +### 2.4 PipelineConfig（SMEM 预算 → stage 数） + +```cpp +constexpr int kNumMaxStages = 16; // SM100 是 32 +const int smem_cd = align(block_m * block_n * elemsize_cd, 1024); // 单缓冲！ +const int smem_barriers = kNumMaxStages * 8 * 2; // = 256 B（只有 full+empty 两组） +const int smem_a_per_stage = load_block_m * block_k * elemsize_a; +const int smem_b_per_stage = load_block_n * block_k * elemsize_b; +// BF16 GEMM 无 SF、无 extra tensormap +const int smem_extra = smem_cd + smem_barriers; +const int smem_per_stage = smem_a_per_stage + smem_b_per_stage; +const int num_stages = min((smem_capacity - smem_extra) / smem_per_stage, kNumMaxStages); +smem_size = smem_extra + num_stages * smem_per_stage; +``` + +`smem_capacity = 232448`（227 KB，H100 上限）。三处与 SM100 的差异精确对应 device 布局（§5）： + +- **`smem_cd` 单缓冲**：SM100 是 `× 2`（双缓冲），SM90 没有 `× 2`。对应 device 侧只有一个 `smem_d`。 +- **`smem_barriers = 16×8×2`**：只有 `full` + `empty` 两组 barrier，每组按 `kNumMaxStages=16` 预留。SM100 是 `32×8×3 + 2×8×2 + 8`（三组 + tmem 两组 + tensor_core），BF16 那里第三组是死占位。SM90 干净得多——没有 TMEM，自然没有 `tmem_full/tmem_empty`。 +- **`kNumMaxStages = 16`**：上限只有 SM100 的一半。原因不是 lane 寄存器（SM90 不像 SM100 那样用 32 个 lane 存描述符），而是纯粹的经验/SMEM 预算取舍。 + +### 2.5 LaunchConfig（线程数模型） + +```cpp +const int num_tma_threads = 128; // 恒 1 个 warpgroup +const int num_math_threads = layout.block_m <= 64 ? 128 : 256; // 1 或 2 个 warpgroup +return { num_sms, cluster_size, num_tma_threads + num_math_threads, + num_tma_threads, num_math_threads, 0, 0 }; // 后两项 SM90 无意义 +``` + +于是总线程数 ∈ `{256, 384}`： + +| `BLOCK_M` | math warpgroup 数 | `kNumMathThreads` | 总线程 | warp 总数 | math warp | TMA warp | +| --- | --- | --- | --- | --- | --- | --- | +| 16 / 32 / 64 | 1 | 128 | 256 | 8 | w0–w3 | w4–w7 | +| 128 / 256 | 2 | 256 | 384 | 12 | w0–w7 | w8–w11 | + +`num_math_threads` 由 `BLOCK_M` 决定的本质：每个 math warpgroup 覆盖 `WGMMA::M = 64` 行，`BLOCK_M > 64` 就需要第二个 warpgroup。`BLOCK_M = 256` 时仍是 2 个 warpgroup，但要做 2 个「wave」（§6.2），而不是 4 个 warpgroup。 + +`__launch_bounds__(kNumTMAThreads + kNumMathThreads, 1)` 的第二个参数 `1` 声明每 SM 最多 1 个 block——与近满额 SMEM 一起，物理上排除 2 CTA 共 SM。 + +### 2.6 TMA descriptor + +三个 descriptor 都以 `__grid_constant__ cute::TmaDescriptor` 按值传参（128 B 常量内存，避免走 GMEM）。构造见 [runtime_utils.hpp](../third_party/DeepGEMM/csrc/jit_kernels/impls/runtime_utils.hpp)： + +| descriptor | gmem (inner, outer) | smem box (inner, outer) | 备注 | +| --- | --- | --- | --- | +| A | K-major: `(k, m×G)`；MN-major: `(m×G, k)` | K-major: `(block_k→64, block_m)`；MN-major: `(block_m, block_k)` | `num_groups > 1` 时强制 K-major；box 内维被 swizzle 覆写为 `swizzle/elem` | +| B | K-major: `(k, n)`；MN-major: `(n, k)` | 同上，`block_n` 换 `load_block_n` | `num_groups` 只作用在外维：`gmem_outer × num_groups` | +| C/D | `(n, m×G)` | `(store_block_n→swizzle/elem, store_block_m)` | D 必须 N-major；FP32 时 swizzle=0，box 内维=`block_n` | + +公共属性：`CU_TENSOR_MAP_INTERLEAVE_NONE`、`CU_TENSOR_MAP_L2_PROMOTION_L2_256B`、`CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE`（越界元素零填充——这是 tail-K 与 M/N 非整除时结果仍正确的硬件保证，见 §7.6）、swizzle 由 `mode_into_tensor_map_swizzle()` 映射到 `CU_TENSOR_MAP_SWIZZLE_{NONE,32B,64B,128B}`。 + +Batched（`sm90_bf16_bhr_hdr_bhd` / `bhd_hdr_bhr`）走 `make_tma_3d_desc()`，第三维是 head，device 侧 `kIsBatchedMM`（= `kGemmType == Batched`）打开 `SM90_TMA_LOAD_3D` 分支。SM90 一共暴露 5 个 host 入口：`sm90_bf16_gemm`（Normal）、`sm90_m_grouped_bf16_gemm_contiguous`、`sm90_bf16_m_grouped_gemm_masked`、`sm90_bf16_k_grouped_gemm`、两个 batched einsum。 + +### 2.7 Launch 属性 + +`LaunchArgs` 携带 `num_sms / num_threads / smem_size / cluster_size`，最终经 `launch_kernel` → `cuLaunchKernelEx`： + +1. `cuFuncSetAttribute(CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size)` —— 动态 SMEM 远超 48 KB 静态上限，必须显式抬。 +2. `cluster_size > 1` → `CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION = {cluster_size, 1, 1}`。 +3. `enable_pdl` → `CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION = 1`。默认关闭，需 `deep_gemm.set_pdl(True)` 打开（与 SM100 共用 `DeviceRuntime`）。 +4. 编译 flag：`--gpu-architecture=sm_90a`（wgmma 与 TMA multicast 都需要 `a` 后缀的架构特性），`-O3 --expt-relaxed-constexpr --expt-extended-lambda`，产物 `-cubin`。kernel 内 `#if __CUDA_ARCH__ >= 900` 之外的分支只留 `DG_DEVICE_ASSERT(false and "This kernel only support sm_90a")`（第 386–389 行）。 + +--- + +## 3. 线程组织与 warp 角色 + +### 3.1 角色表（以 `BLOCK_M=128` ⇒ 384 线程 / 12 warp 为例） + +| warp | 条件 | 角色 | 实际活跃线程 | +| --- | --- | --- | --- | +| 0–3 | `warp_idx < kNumMathThreads/32` | **math warpgroup 0**：发 `wgmma`（rows 0–63）、持累加器、做 epilogue | 128（全 warpgroup 协同发 wgmma） | +| 4–7 | 同上（`kNumMathThreads=256` 时） | **math warpgroup 1**：rows 64–127 | 128 | +| `kNumMathThreads/32`（=8） | `warp_idx == kNumMathThreads/32 && elect_one_sync()` | **prefetch** 三个 TMA descriptor（prologue 一次） | 1 lane | +| `+1`（=9） | `warp_idx == kNumMathThreads/32+1 && elect_one_sync()` | **init barriers**（prologue 一次）+ `fence_barrier_init` | 1 lane | +| `+2`（=10） | `warp_idx == kNumMathThreads/32+2 && elect_one_sync()` | **TMA load**：persistent 遍历所有块，每块跑完整 K 循环，发 A/B 的 TMA | **1 lane**（TMA 单线程指令） | +| `+3`（=11） | — | **空转** | 0 | + +`BLOCK_M ≤ 64`（`kNumMathThreads=128`）时，math warpgroup 只有 w0–w3，TMA warpgroup 是 w4–w7，角色一一对应下移。 + +源码第 152–154 行的注释解释了为什么 TMA load 用**第三个** warp 而非第一个： + +```cpp +// NOTES: only one thread (or warp) will be used +// We use the third warp, as warp 0/1 may be doing WGMMA with `BLOCK_M == 32` +if (warp_idx == kNumMathThreads / 32 + 2 and cute::elect_one_sync()) { +``` + +这里的「warp 0/1」指的是 math warpgroup 内部的相对编号——当 `BLOCK_M=32 < WGMMA::M=64` 时，只有前 2 个 math warp（w0/w1）参与 WGMMA store（§3.4），把 TMA 的三个 prologue 角色错开到 math warpgroup 之后、且集中在同一个 TMA warpgroup 内，避免与 math warp 争用。 + +### 3.2 寄存器再配置（`setmaxnreg`） + +Hopper 引入的 warpgroup 级寄存器再分配，本 kernel 用它把寄存器从「不怎么用寄存器的 TMA warpgroup」挪给「累加器吃寄存器的 math warpgroup」： + +```cpp +// 第 128–129 行：编译期常量 +constexpr uint32_t kNumTMARegisters = 48; +constexpr uint32_t kNumMathRegisters = kNumMathThreads == 128 ? 248 : 224; + +// TMA warpgroup 分支（第 150 行） +cutlass::arch::warpgroup_reg_dealloc(); // setmaxnreg.dec.sync.aligned.u32 48 + +// math warpgroup 分支（第 210 行） +cutlass::arch::warpgroup_reg_alloc(); // setmaxnreg.inc.sync.aligned.u32 224/248 +``` + +`kNumMathRegisters` 的取值是被 **64K 寄存器/SM 的总预算**反解出来的（`setmaxnreg` 要求 8 的倍数、范围 `[24, 256]`）： + +| 配置 | 总线程 | math 配额 | TMA 配额 | 合计 | 是否 ≤ 65536 | +| --- | --- | --- | --- | --- | --- | +| 1 math wg（`BLOCK_M≤64`） | 256 | `128 × 248 = 31744` | `128 × 48 = 6144` | 37888 | ✓（还有余量，但 248 已接近单线程上限 255） | +| 2 math wg（`BLOCK_M>64`） | 384 | `256 × 224 = 57344` | `128 × 48 = 6144` | 63488 | ✓（贴着上限，故 math 只能给到 224 而非 248） | + +这解释了「为什么 2 个 math warpgroup 时每线程寄存器反而更少（224 < 248）」：线程总数从 256 涨到 384，为了塞进同一个 64K 寄存器文件，单线程配额必须下调。这也从侧面印证了 §2.2 那条 `block_m>128 && block_n>128 → 跳过` 的过滤——累加器 `block_n/2 × waves` 个寄存器 + 描述符 + 地址，在 224/248 的预算里放不下两个都大的维度。 + +**SM100 没有这一步**：累加器在 TMEM 不吃寄存器，无需再分配。寄存器再配置是 SM90「累加器在寄存器」这一根本约束的直接衍生。 + +### 3.3 为什么 WGMMA 需要 128 个线程 + +`wgmma.mma_async.sync.aligned` 是 **warpgroup 级协同指令**：`.sync.aligned` 要求 warpgroup 内全部 128 线程收敛执行，每个线程贡献自己那一份 A（若 `_RS`）/持有自己那一份 D 累加器寄存器。本 kernel 用 `_SS` 变体（A、B 都来自 SMEM descriptor），所以线程不参与提供 A/B 数据，但**必须全部到场**，因为： + +- 一条 `m64nNk16` 的 D 累加器是 `64 × N` 个 FP32，硬件把它按固定图案**分散到 128 个线程的寄存器**里，每线程恰好 `64×N/128 = N/2 = kNumAccum` 个。 +- `WGMMA::wgmma(desc_a, desc_b, shifted_accum, 1)` 的 `shifted_accum` 就是本线程持有的那 `kNumAccum` 个寄存器的指针；`call_fma_impl` 用 `cute::make_index_sequence` 把它们逐个展开成 `MMA::fma(desc_a, desc_b, d[0], d[1], …, scale)` 的操作数。 + +对照 SM100：`tcgen05.mma` 的 D 在 TMEM，没有任何线程需要「拿着」累加器，于是发射退化成单 lane 的控制指令。**SM90 的 128 线程是被累加器的存储位置逼出来的，与算力无关**——但和 SM100 不同的是，这些线程在 epilogue 阶段确实要干活（把寄存器里的累加器搬进 SMEM），所以它们在整个输出块生命周期里都是「有事做」的，不是空转。 + +### 3.4 math warpgroup 数量、wave 与 store 线程 + +三个编译期量共同决定了 math 侧的形状（第 223–230 行）： + +```cpp +constexpr uint32_t WAVE_BLOCK_M = BLOCK_M <= WGMMA::M ? BLOCK_M : WGMMA::M * 2; // ≤64→BLOCK_M；>64→128 +DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes"); +float accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; // 每线程的累加器 + +constexpr uint32_t kNumWGMMAStoreThreads = WAVE_BLOCK_M * (128 / WGMMA::M); // = WAVE_BLOCK_M × 2 +const bool do_wgmma_store = BLOCK_M >= 64 or warp_idx < kNumWGMMAStoreThreads / 32; +``` + +| `BLOCK_M` | math wg 数 | `WAVE_BLOCK_M` | wave 数 `BLOCK_M/WAVE_BLOCK_M` | 每线程 accum | `kNumWGMMAStoreThreads` | 参与 store 的 warp | +| --- | --- | --- | --- | --- | --- | --- | +| 16 | 1 | 16 | 1 | `(16/2)×1=8`… 见下注 | 32 | w0（`BLOCK_M<64`，只前 1 warp） | +| 32 | 1 | 32 | 1 | 16 | 64 | w0–w1 | +| 64 | 1 | 64 | 1 | 32 | 128 | w0–w3（`BLOCK_M≥64` 全 store） | +| 128 | 2 | 128 | 1 | 64 | 256 | 全 8 warp | +| 256 | 2 | 128 | 2 | 128 | 256 | 全 8 warp | + +> 注：每线程 accum 数 = `kNumAccum × waves = (BLOCK_N/2) × waves`，上表按 `BLOCK_N=64` 举例（`kNumAccum=32`）；实际值随 `BLOCK_N` 线性变化，`BLOCK_N=256` 时每线程可高达 `128 × waves` 个 FP32，这正是 §2.2 寄存器过滤的由来。 + +关键区分： + +- **wave**（`local_idx`）：一个 math warpgroup（覆盖 64 行）需要跑几轮才能覆盖 `WAVE_BLOCK_M`。`BLOCK_M=256` 时 `WAVE_BLOCK_M=128`、2 个 warpgroup各覆盖 64 行合起来 128 行 = 1 个 wave，需要 2 个 wave 才够 256 行。wave 之间累加器数组 `accum` 分段（`shifted_accum = accum + kNumAccum × local_idx`）。 +- **`do_wgmma_store`**：`BLOCK_M ≥ 64` 时全部 math warp 都参与 epilogue；`BLOCK_M < 64`（16/32）时，因为 WGMMA 按 M=64 发射、但只有前 `BLOCK_M` 行有效，只有前 `kNumWGMMAStoreThreads/32` 个 warp 的结果需要写回，其余 warp `continue` 跳过（第 289–290 行）。 + +`a_desc` 在构造时就用 `math_wg_idx * WGMMA::M` 把每个 warpgroup 定位到自己那 64 行（第 217 行），wave 内再用 `local_idx * WAVE_BLOCK_M` 推进（第 262–263 行）。两者叠加：warpgroup `g`、wave `w` 覆盖的全局行区间是 `[g×64 + w×WAVE_BLOCK_M, +64)`。 + +--- + +## 4. Persistent 调度器 + +SM90 与 SM100 **共用同一个** `sched::Scheduler`（[scheduler/gemm.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/scheduler/gemm.cuh)），差异只在模板实参与少数 `#if __CUDA_ARCH__ < 1000` 的 SM90 专属分支。 + +### 4.1 复制式状态机，不是共享工作队列 + +第 135–136 行在**角色分派之前**构造 `scheduler`： + +```cpp +auto scheduler = sched::Scheduler( + shape_m, shape_n, shape_k, grouped_layout); +``` + +它是个**寄存器里的值对象**，每个线程各持一份私有副本。TMA warp（w10）与每个 math warpgroup 各自独立调用 `get_next_block()`，靠 `next_block_idx = (++current_iter) * kNumSMs + blockIdx.x`（第 198 行）这一条纯算术式子得到**完全相同**的块序列： + +``` +iter 0 → blockIdx.x +iter 1 → blockIdx.x + kNumSMs +iter 2 → blockIdx.x + 2*kNumSMs +... +``` + +因此： + +- **零原子操作、零全局 ticket 计数器**，调度开销是几条整数指令； +- TMA warp 与 math warpgroup 天然锁步，无需为「现在在处理哪一块」建立任何额外通信； +- 与 SM100 不同的是，SM90 **不需要**用 `current_iter` 去推导 TMEM 累加器缓冲的 stage/phase（没有 TMEM）。SM90 里 math warpgroup 与 TMA warp 的耦合只通过 A/B 环的 `full`/`empty` barrier，累加器是每块就地 `{0}` 重置的寄存器数组。 + +代价同样是负载不能动态窃取：尾波的空闲 CTA 只能干等。但 SM90 的打分模型（§2.2）用解析带宽模型 + `wave_efficiency` 提前量化了这件事。 + +### 4.2 L2 swizzle 分组 + +`get_swizzled_block_idx()`（第 117–153 行）把线性的 `block_idx` 重映射成 `(m_block_idx, n_block_idx)`，目的是让**同时在飞的 kNumSMs 个 CTA 尽量共享 L2 里的 A/B**。这段逻辑 SM90/SM100 完全一致： + +```cpp +kNum1DBlocksPerGroup = get_num_1d_blocks_per_group<...>(); // 编译期，∈ {8, 16} +primary_num_blocks = kIsMulticastOnA ? num_n_blocks : num_m_blocks; +secondary_num_blocks = kIsMulticastOnA ? num_m_blocks : num_n_blocks; +num_blocks_per_group = secondary_num_blocks * kNum1DBlocksPerGroup; +group_idx = block_idx / num_blocks_per_group; +first_block_idx = group_idx * kNum1DBlocksPerGroup; +in_group_idx = block_idx % num_blocks_per_group; +num_blocks_in_group = min(kNum1DBlocksPerGroup, primary_num_blocks - first_block_idx); + +// kIsMulticastOnA == false（在 M 上分组，组内 M 变化最快） +m_block_idx = first_block_idx + in_group_idx % num_blocks_in_group; +n_block_idx = in_group_idx / num_blocks_in_group; +``` + +组大小选 `{8, 16}` 中最小化 L2 工作集者（`get_num_1d_blocks_per_group`，第 14–26 行）。`DG_STATIC_ASSERT(kNum1DBlocksPerGroup % kNumMulticast == 0)` 保证一个 cluster 的两个 CTA 不跨组边界。 + +### 4.3 SM90 专属：multicast 的动态关闭 + +这是 SM90 调度器与 SM100 最实质的分歧。SM100 的 2-CTA UMMA **不能动态关闭**（硬件成对发射），只能靠 host 侧整除性过滤；而 **SM90 的 TMA multicast 可以在运行期逐块关闭**，因为发不发 multicast load 只是 TMA warp 的一个分支选择，两个 CTA 始终各算各的块。为此调度器提供了两个 SM90-only 方法： + +**① `is_tma_multicast_valid(m_block_idx)`**（第 290–307 行）： + +```cpp +if (num_blocks_in_group == 1) return false; // 组内只剩 1 块，无 peer 可 multicast +if constexpr (Normal / Masked / KGrouped / Batched / MGroupedPsum) return true; +else /* MGroupedContiguous */ { + if constexpr (kIsMulticastOnA) return true; + else return grouped_layout[m_block_idx*BLOCK_M] == grouped_layout[(m_block_idx^1)*BLOCK_M]; // peer 同组才能共享 B +} +``` + +TMA warp 在第 161 行读它，决定 `num_tma_multicast_a/b` 是取 `kNumTMAMulticast` 还是退回 `1`（第 162–163 行）。对 m-grouped contiguous，若相邻两块（`m_block_idx` 与 `m_block_idx ^ 1`）属于**不同 group**，它们的 B 就不是同一份，multicast 非法 → 退回单播。 + +**② `is_peer_cta_alive`**（第 281–283 行，仅 Normal 路径设置）： + +```cpp +is_peer_cta_alive = num_n_blocks % kNumMulticast == 0 or // N 恒对齐（常量短路） + num_m_blocks % kNumMulticast == 0 or // M 恒对齐（常量短路） + (next_block_idx ^ 1) < num_blocks; // peer CTA 的块仍在界内 +``` + +它服务于 `empty_barrier_arrive` 的目标 CTA 选择（§10.3）：当 cluster 的 peer CTA 因为落在矩阵边界外而没有有效块时，本 CTA 的 math warp 不能把 empty 信号远程投递给「不存在的 peer」，否则会写到一个没人等待、甚至已释放的 barrier 上。`is_peer_cta_alive == false` 时，两个 lane 都投递给本 CTA（`target_cta = block_rank_in_cluster()`）。 + +第 132–142 行还有一段 `#if __CUDA_ARCH__ < 1000` 的「修正不对齐的 TMA multicast」，注释点明：*for SM90 only, as SM90 can dynamically disable TMA multicast while SM100 uses 2-CTA, which can not be dynamically disabled*。当 `num_blocks_in_group` 是奇数时，它把最后一个落单的块单独成组（`num_blocks_in_group = 1`），从而让 `is_tma_multicast_valid` 对它返回 false、退回单播。 + +### 4.4 GemmType 变体 + +| GemmType | `num_blocks` | 额外状态 | 说明 | +| --- | --- | --- | --- | +| `Normal` | `num_m_blocks × num_n_blocks` | — | `get_global_idx` 退化为 `block_idx × block_size`；设置 `is_peer_cta_alive` | +| `Batched` | 同上，`× kNumGroups` | `current_group_idx` 作 batch_idx | 不走 swizzle，按 `kIsMulticastOnA` 决定 m/n 谁变化快；TMA 走 3D；host 侧禁用 multicast | +| `MGroupedContiguous` | 同上 | `grouped_layout[m]` = 每行所属 group | B 的外维加 `group × shape_dim` 偏移；multicast 需 peer 同组 | +| `MGroupedMasked` | 逐 group 累加 | `current_m_cumsum` | 边扫边把 `next_block_idx` 落到对应 group，`num_m_blocks` 每 group 重算 | +| `MGroupedContiguousWithPsumLayout` | 逐 group 累加 | `last_psum_m`/`current_psum_m`/`current_m_block_cumsum` | group 边界按 psum 偏移切分，`m_block_idx += last_psum_m / BLOCK_M` | +| `KGroupedContiguous{,WithPsumLayout}` | 同上 | `current_shape_k`/`current_k_cumsum`/`current_k_start,end` | 每 group 的 K 长度不同 → `num_total_k_blocks` 逐块变化；要求 A/B 都 MN-major | + +对本 kernel 最关键的约束在第 177 行： + +```cpp +DG_STATIC_ASSERT(kGemmType == GemmType::Normal or kGemmType == GemmType::KGroupedContiguous + or kMajorA == cute::UMMA::Major::K, "Invalid major"); +``` + +即所有 m-grouped 变体的 A 必须 K-major（group 偏移加在外维上）。host 侧对应 `DG_HOST_ASSERT(major_a == K)`。 + +### 4.5 跨块连续的流水线状态 + +第 139–146 行： + +```cpp +uint32_t stage_idx = 0, phase = 0; +auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + // Flip phases only if reach the next first stage + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; // 只在回绕到 stage 0 时翻转相位 +}; +``` + +`stage_idx`/`phase` 声明在**块循环之外**，意味着 A/B 环**跨输出块连续运转**：TMA warp 可以在 math warpgroup 还在算第 i 块最后一个 k_block 时，就开始往刚被释放的 stage 里灌第 i+1 块的 k=0 数据。块边界上没有「排空-重启」的开销，这是 persistent kernel 相较 grid-per-tile 的核心优势。 + +`advance_pipeline` 同时被 TMA warp（第 167 行 `for` 的递增式）和 math warpgroup（第 244 行）使用，两边独立维护但推进规则一致，因此 `stage_idx`/`phase` 序列天然对齐。注意 SM90 的 `phase` 翻转写成 `phase ^= stage_idx == 0`（先更新 stage_idx 再判断是否回绕到 0），与 SM100 的 `stage_idx = (stage_idx+1) % kNumStages; phase ^= stage_idx == 0` 语义等价，只是三元写法不同。 + +与 SM100 的一个显著区别：SM100 的 C/D 环也是跨块连续的（`tma_stage_idx` 双缓冲）；SM90 的 D 是**单缓冲**，跨块复用同一个 `smem_d`，靠 epilogue 开头的 `tma_store_wait<0>()` 串行化（§11.1）。所以 SM90 的「跨块连续」只体现在 A/B 环，D 侧是块间串行的。 + +--- + +## 5. 共享内存布局 + +### 5.1 线性布局 + +```cpp +extern __shared__ __align__(1024) uint8_t smem_buffer[]; // 1024 B 对齐，服务于 swizzle-128B +``` + +三个区段尺寸都是编译期常量（第 73–75 行）： + +```cpp +static constexpr uint32_t SMEM_D_SIZE = constexpr_align(BLOCK_M * BLOCK_N * sizeof(cd_dtype_t), 1024u); +static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_bfloat16); +static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_bfloat16); +``` + +注意 A/B 用的是 **`BLOCK_M`/`BLOCK_N` 而非 SM100 的 `LOAD_BLOCK_M/N`**——因为 SM90 的 multicast 不切分 SMEM（§7.2），每个 CTA 都存整份 A、整份 B，所以「load block」与「block」是同一个量，源码里干脆不引入 `LOAD_*` 记号。 + +`utils::PatternVisitor`（[common/utils.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/common/utils.cuh) 第 11–22 行）是个零开销的「下标 → 指针」闭包包装器（`operator[](i)` 直接调 lambda），用它替代指针数组，避免在 SMEM/寄存器里存 stage 指针表： + +| 区段 | 起址 | 大小 | 访问器 | +| --- | --- | --- | --- | +| **D（单缓冲）** | `smem_buffer + 0` | `SMEM_D_SIZE` | `smem_d`（裸指针，非 ring） | +| A ring | `smem_buffer + SMEM_D_SIZE` | `kNumStages * SMEM_A_SIZE_PER_STAGE` | `smem_a[i]` | +| B ring | `+ kNumStages * SMEM_A_SIZE_PER_STAGE` | `kNumStages * SMEM_B_SIZE_PER_STAGE` | `smem_b[i]` | +| Barriers | `+ kNumStages * SMEM_B_SIZE_PER_STAGE` | 见 §5.2 | `full_barriers[i]` / `empty_barriers[i]` | + +```cpp +auto smem_d = reinterpret_cast(smem_buffer); +auto smem_a = PatternVisitor([&](uint32_t i){ return (bf16*)(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE); }); +auto smem_b = PatternVisitor([&](uint32_t i){ return (bf16*)(smem_buffer + SMEM_D_SIZE + + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); }); +``` + +三个 `DG_STATIC_ASSERT(... % 1024 == 0)`（第 95–96 行）保证每个区段起点都 1024 B 对齐——这是 swizzle-128B 的硬件要求（一个 swizzle atom 是 8 行 × 128 B = 1 KB，若基址不按 1 KB 对齐，TMA 写入的 swizzle 图案与 WGMMA 描述符解读的图案会错位）。`SMEM_D_SIZE` 显式 `constexpr_align(…, 1024)`，A/B 因 `BLOCK_* × 64 × 2` 天然是 1024 的倍数（`BLOCK_K=64`、bf16=2B ⇒ 每行 128 B，8 行 1 KB）。 + +**与 SM100 的三处结构差异**： + +1. **D 在最前、且单缓冲**。SM100 是 `C/D ring`（`× kNumTMAStoreStages=2`）在最前；SM90 只有一个 `smem_d`，没有 ring 下标。 +2. **没有 TMEM 基址槽**。SM100 在 barriers 之后还要放一个 4 B 的 `tmem_ptr_in_smem`（`tcgen05.alloc` 的结果）；SM90 无 TMEM，barriers 之后就是末尾。 +3. **barriers 只有两组**（§5.2），SM100 有五组 + 一个空洞。 + +### 5.2 Barrier 区（干净的两组） + +以 `Barrier`（= `cutlass::arch::ClusterTransactionBarrier`，8 B）为单位，`barrier_start_ptr` 起： + +```cpp +auto barrier_start_ptr = (Barrier*)(smem_buffer + SMEM_D_SIZE + + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE)); +auto full_barriers = PatternVisitor([=](uint32_t i){ return barrier_start_ptr + i; }); +auto empty_barriers = PatternVisitor([=](uint32_t i){ return barrier_start_ptr + kNumStages + i; }); +``` + +| 索引区间 | 名称 | 个数 | `init()` 计数 | +| --- | --- | --- | --- | +| `[0, S)` | `full_barriers` | `kNumStages` | `1` | +| `[S, 2S)` | `empty_barriers` | `kNumStages` | `kNumTMAMulticast × kNumMathThreads/32` | + +（`S = kNumStages`） + +host 侧 `smem_barriers = kNumMaxStages * 8 * 2 = 256 B`（§2.4）与这里的 `2 × kNumStages` 个 Barrier 精确呼应。**没有 SM100 那个 `kNumStages` 大小的空洞**——SM100 的索引按「每 stage 三组 barrier」排布（第三组是 FP8/FP4 的 with-SF full barriers，BF16 下空占位），SM90 只有 full + empty 两组，索引连续，不留死区。 + +初始化由 `warp_idx == kNumMathThreads/32 + 1`（TMA warpgroup 的第二个 warp）的单个 lane 完成（第 113–122 行），随后： + +```cpp +cutlass::arch::fence_barrier_init(); // fence.mbarrier_init.release.cluster +``` + +注释写明目的：*Make initialized barrier visible in async proxy*。mbarrier 会被 TMA 这个**异步代理**访问（`arrive_and_expect_tx` 的 tx 计数由 TMA 硬件回填），普通的 `__syncthreads()` 不足以建立 generic proxy → async proxy 的可见性，必须用这条 cluster 作用域的 release fence。之后第 125 行做 `cluster_sync_with_relaxed_arrive()`（multicast）或 `__syncthreads()`（单 CTA），确保 peer CTA 也能看到 leader 初始化的 barrier 状态。 + +### 5.3 A/B stage 内部排布 + +**K-major（最常见）**：TMA box 是 `(inner = BLOCK_K = 64 elem = 128 B, outer = BLOCK_M)`，swizzle atom = 8 行 × 128 B。一个 stage 就是 `BLOCK_M/8` 个 atom 沿 M 方向线性堆叠： + +``` +smem_a[s] (BLOCK_M=128, BLOCK_K=64, bf16 → 16 KB) +┌──────────────────────── atom 0 : rows 0.. 7 ────────────────────────┐ ← SBO = 1024 B +│ row r: 128 B = 8 个 16-B bank group，物理位置 g' = g ^ (r % 8) │ +├──────────────────────── atom 1 : rows 8..15 ────────────────────────┤ +│ ... │ +├─────────────────────── atom 15 : rows 120..127 ──────────────────────┤ +└───────────────────────────────────────────────────────────────────────┘ +``` + +`^ (r % 8)` 的 bank-group 置换就是 `CU_TENSOR_MAP_SWIZZLE_128B` / `cute::SM90::GMMA::LayoutType::B128` 定义的图案，由 TMA 硬件在写入时施加、由 WGMMA 硬件在读取时反解，软件两侧都不参与。作用是把「同一列的 8 个元素」打散到 8 个不同 bank group，消除 tensor core 按列取数时的 SMEM bank conflict。 + +**MN-major**：TMA box 变成 `(inner = BLOCK_MN, outer = BLOCK_K)`，`BLOCK_INNER_ATOM = swizzle/elem = 64`，`tma::copy` 内部循环 `BLOCK_MN / 64` 次，每次目的地址 `smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM`。即 SMEM 布局是「**K 外、MN-atom 内**」。WGMMA 描述符的 `stride_k` 也随之从 K-major 的 `1` 变成 `get_inner_block_atom_size<...>()`（§9.4）。**Stage 合并**（§7.4）触发时，一个 stage 内是 `kNumStagesPerMerge` 个 64-宽 K-atom 沿「MN 外、K-atom 内」排列，描述符构造用 `BLOCK_ATOM_K = BLOCK_K / kNumStagesPerMerge` 而非合并后的 `BLOCK_K`。 + +### 5.4 D 单缓冲内部排布 + +D 只有一块 `SMEM_D_SIZE = align(BLOCK_M × BLOCK_N × sizeof(cd), 1024)`，覆盖整个输出块（不像 A/B 分 stage）： + +- **BF16 输出**（`kSwizzleDMode > 0`，恒 128）：一块是 `BLOCK_M` 行 × 128 B，行内 8 个 16-B bank group 的 `^ (r % 8)` 置换。每个 swizzle atom 覆盖 `TMA_D_BLOCK_N = kSwizzleDMode / sizeof(bf16) = 64` 个 N 元素，故 `BLOCK_N` 宽被切成 `BLOCK_N / 64` 个 atom，对应 `BLOCK_N / TMA_D_BLOCK_N` 条 TMA store。 +- **FP32 输出**（`kSwizzleDMode == 0`）：不 swizzle，D 就是行主序的 `BLOCK_M × BLOCK_N` 个 float；`TMA_D_BLOCK_N` 退化为整个 `BLOCK_N`，一条 TMA store 覆盖整块（§11.3）。 + +因为单缓冲，**下一个输出块的 epilogue 必须等上一块的 TMA store 把 SMEM 读完**才能覆写——这就是 epilogue 开头 `tma_store_wait<0>()` 的作用（§11.4），也是 SM90 无法像 SM100 那样让 compute 与 epilogue 重叠的根因之一。 + +### 5.5 WGMMA 越界读的静态防护（含一处代码瑕疵） + +第 77–79 行： + +```cpp +// NOTES: Make sure we have enough shared memory for WGMMA padding +static constexpr uint32_t WGMMA_A_SIZE_PER_STAGE = WGMMA::M * BLOCK_K * sizeof(__nv_fp8_e4m3); +DG_STATIC_ASSERT(WGMMA_A_SIZE_PER_STAGE <= SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE * kNumStages, + "Memory Out of bound for WGMMA"); +``` + +动机与 SM100 §5.5 同源：`WGMMA::M` 恒为 64，当 `BLOCK_M = 16/32` 时 `SMEM_A_SIZE_PER_STAGE < 64 行`，WGMMA 仍**按 64 行去读 A**，越过的部分落在后续 A stage、乃至 B ring 上。那些行算出的 D 在 epilogue 里不会被读（`do_wgmma_store` 只放行前 `BLOCK_M` 行，§3.4），读到垃圾无害——但**必须落在本 CTA 已分配的动态 SMEM 内**。 + +**瑕疵**：`sizeof(__nv_fp8_e4m3)` 是 **1 字节**，而本 kernel 的 A 是 bf16（2 字节）。WGMMA 实际读取的 A footprint 是 `64 × BLOCK_K × 2` 字节，断言左值却只算了 `64 × BLOCK_K × 1`——**恰好少算一半**。这几乎可以肯定是从 FP8 kernel（`sm90_fp8_gemm.cuh`，A 为 e4m3、1 字节）移植到 BF16 时漏改的 `sizeof`。 + +为什么目前不出事：断言右值 `SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE × kNumStages` 由 `kNumStages`（≥3，通常 5–6）个 B stage 主导，裕量极大。以 `BLOCK_M=16, BLOCK_N=256, BLOCK_K=64, kNumStages=5` 为例，真实需求 `64×64×2 = 8192` B，右值 `16×64×2 + 256×64×2×5 = 2048 + 163840 = 165888` B——无论用 1 字节还是 2 字节，断言都轻松通过。所以这是一个**潜伏的、被裕量掩盖的松检查**：它把安全边界放宽了 2×，一旦将来出现「极小 `BLOCK_M` + 极小 `kNumStages` + 极小 `BLOCK_N`」的组合，理论上可能漏判一次真实越界。修正方式是把 `sizeof(__nv_fp8_e4m3)` 改成 `sizeof(__nv_bfloat16)`（或 `sizeof(cd 的输入 dtype)`），代价为零。详见 §13.3。 + +### 5.6 Worked Example：8192 × 8192 × 8192，BF16→BF16，H100（132 SM） + +取一组代表性配置：`BLOCK_M=128, BLOCK_N=128, BLOCK_K=64, cluster=(1,2)`（multicast on A），`swizzle A/B/D = 128`。 + +推导链： + +``` +kNumTMAMulticast = 2 (cluster_m=1, cluster_n=2) +kIsTMAMulticastOnA = true (cluster_n > 1) +load_block_m = BLOCK_M = 128 (不除 cluster！multicast 各存整份 A) +load_block_n = BLOCK_N = 128 +WGMMA::M = 64, WGMMA::N = 128, WGMMA::K = 16, kNumAccum = 64×128/128 = 64 +kNumMathThreads = 256 (BLOCK_M=128 > 64) ⇒ 2 个 math warpgroup，总线程 384 +WAVE_BLOCK_M = 128 (BLOCK_M>64 ⇒ WGMMA::M×2)，wave 数 = 128/128 = 1 +每线程 accum = kNumAccum × waves = 64 × 1 = 64 个 FP32 +kNumMathRegisters = 224 (2 wg)，kNumTMARegisters = 48 +kDoMergeStages = false (num_stages=6 < 10) +kNum1DBlocksPerGroup: kIsMulticastOnA ⇒ 组在 N + cand 8 → 8×128 + ceil(132/8)×128 = 1024 + 2176 = 3200 + cand 16 → 16×128 + ceil(132/16)×128 = 2048 + 1152 = 3200 ⇒ 平手取 8 +``` + +num_stages 反解（`smem_capacity = 232448`）： + +``` +smem_cd = align(128×128×2, 1024) = 32768 (单缓冲，无 ×2) +smem_barriers = 16×8×2 = 256 +smem_per_stage = 128×64×2 (A) + 128×64×2 (B) = 16384 + 16384 = 32768 +num_stages = min((232448 − 32768 − 256) / 32768, 16) = min(6.08, 16) = 6 +smem_size = 32768 + 256 + 6×32768 = 229632 +``` + +SMEM 字节表（单 CTA）： + +| 偏移 | 大小 | 内容 | +| --- | --- | --- | +| 0 | 32 768 | D 单缓冲：128 行 × 128 B（bf16 输出，swizzle-128B） | +| 32 768 | 98 304 | A ring：6 × (128 × 64 × 2 B = 16 KB) | +| 131 072 | 98 304 | B ring：6 × (128 × 64 × 2 B = 16 KB) | +| 229 376 | 96 | barriers：`full[6] + empty[6] = 12` 个 Barrier | +| **device 合计** | **229 472** | | +| **host 申请** | **229 632** | `33 024 (extra) + 6 × 32 768` | + +**关键对照**：本例开了 multicast（cluster=2），但每个 CTA 仍存**整份** A（16 KB/stage），SMEM 占用与不开 multicast 时**完全相同**。multicast 省下的是 L2/GMEM 带宽——两个 CTA 的 A 来自同一次 GMEM 读（§7.2、§2.2 打分模型里 `block_m/cluster_n` 那一项）。反观 SM100 的 2-CTA UMMA，同样 cluster=2 时 `LOAD_BLOCK_M = BLOCK_M/2`，A 的 SMEM 直接减半，stage 数能翻倍——这是「复制 vs 切分」的本质差别。 + +每个 k_block：单 CTA 载入 32 KB（A 16 KB + B 16 KB），计算 `128 × 128 × 64 × 2 = 2.10 MFLOP`，由 2 个 math warpgroup 各发 `BLOCK_K/WGMMA::K = 4` 条 `m64n128k16` WGMMA 完成自己那 64 行。 + +--- + +## 6. 寄存器累加器布局 + +> 本章对应 SM100 文档的「§6 Tensor Memory 布局」。SM90 没有 TMEM，累加器住在 math warpgroup 的寄存器里，因此这一章讲的是**一条 `wgmma` 的 D 累加器如何按硬件固定图案分散到 128 个线程的寄存器**，以及 kernel 如何用 `accum[]` 数组 + wave/local_idx 索引去命中它。 + +### 6.1 累加器数组与分片规则 + +math warpgroup 分支第 223–225 行： + +```cpp +constexpr uint32_t WAVE_BLOCK_M = BLOCK_M <= WGMMA::M ? BLOCK_M : WGMMA::M * 2; // ≤64→BLOCK_M；>64→128 +DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes"); +float accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; // 每线程的累加器 +``` + +- `WGMMA::kNumAccum = WGMMA::M × WGMMA::N / 128 = 64 × BLOCK_N / 128 = BLOCK_N / 2`（[mma/sm90.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/mma/sm90.cuh) 第 92 行）。物理含义：一条 `m64nNk16` WGMMA 产生 `64 × N` 个 FP32 累加器，硬件把它们**均匀分散到发射它的 128 个线程**，每线程恰好 `64×N/128 = N/2` 个。 +- `BLOCK_M / WAVE_BLOCK_M` 是 **wave 数**（§3.4）：一个 math warpgroup 覆盖 `WGMMA::M = 64` 行，`WAVE_BLOCK_M` 行需要 `WAVE_BLOCK_M/64` 个 warpgroup 并排；`BLOCK_M` 行需要 `BLOCK_M/WAVE_BLOCK_M` 个 wave 串起来。累加器数组按 wave 分段，每段 `kNumAccum` 个。 +- `= {0}`：整个数组在**每个输出块开头**清零。这是 SM90 与 SM100 的一个关键差异——SM100 靠首条 UMMA 的 `scale_c=0` 覆写来「隐式清零 TMEM」，SM90 靠寄存器数组初始化显式清零（对寄存器数组而言 `{0}` 是免费的，编译成若干 `mov` 或直接被后续 WGMMA 覆盖）。 + +### 6.2 每线程寄存器占用（对照表） + +每线程累加器数 = `kNumAccum × waves = (BLOCK_N/2) × (BLOCK_M/WAVE_BLOCK_M)`： + +| `BLOCK_M` | `BLOCK_N` | wave 数 | `kNumAccum` | 每线程 accum(FP32) | `kNumMathRegisters` | accum 占比 | +| --- | --- | --- | --- | --- | --- | --- | +| 64 | 128 | 1 | 64 | 64 | 248 | 26% | +| 128 | 128 | 1 | 64 | 64 | 224 | 29% | +| 128 | 256 | 1 | 128 | 128 | 224 | 57% | +| 256 | 128 | 2 | 64 | 128 | 224 | 57% | +| 256 | 256 | 2 | 128 | 256 | — | **超出，被 host 过滤** | + +最后一行正是 §2.2 那条 `block_m > 128 && block_n > 128 → 跳过` 的由来：`BLOCK_M=256` 且 `BLOCK_N=256` 时每线程要 256 个 FP32 存累加器，加上描述符、地址、循环变量，224 的寄存器配额根本放不下（会 spill 到 local memory，性能崩塌）。host 侧用这条过滤把「两个维度都大」的组合直接排除。 + +**与 SM100 的根本对照**：SM100 的累加器在 TMEM（256 KB 独立存储，128 行 × 512 列），**一个寄存器都不占**，所以 SM100 能同时开 `UMMA_M=256`（双缓冲占满 512 列 TMEM）而不受寄存器约束。SM90 把累加器放寄存器，代价是：① 需要 128 个线程「拿着」它（§3.3）；② `BLOCK_M × BLOCK_N` 的乘积被寄存器容量卡死；③ 需要 `setmaxnreg` 从 TMA warpgroup 抢寄存器（§3.2）。 + +### 6.3 WGMMA 的累加目标与 scale_d + +内层 K 循环第 257–266 行： + +```cpp +for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { + auto shifted_accum = accum + WGMMA::kNumAccum * local_idx; // 定位到本 wave 那段 + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + /* 更新 a_desc.reg32_[0] / b_desc.reg32_[0]（§9.4） */ + WGMMA::wgmma(a_desc, b_desc, shifted_accum, 1); // scale_d 恒为 1 + } +} +``` + +`WGMMA::wgmma`（[mma/sm90.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/mma/sm90.cuh) 第 85–87 行）： + +```cpp +static void wgmma(uint64_t desc_a, uint64_t desc_b, float* d, bool scale_d) { + call_fma_impl(desc_a, desc_b, d, scale_d, cute::make_index_sequence{}); +} +// call_fma_impl: MMA::fma(desc_a, desc_b, d[Idx]..., scale_d ? ScaleOut::One : ScaleOut::Zero); +``` + +`cute::make_index_sequence` 把本线程的 `kNumAccum = N/2` 个寄存器 `d[0..N/2-1]` **逐个展开**成 `MMA::fma` 的操作数——这就是 §3.3「每个线程必须到场」的具体体现：`fma` 的 D 操作数就是本线程寄存器里的 `shifted_accum[i]`。 + +`scale_d` 在本 kernel **恒传 1**（累加）。清零不靠 scale_d，而靠 §6.1 的 `accum[] = {0}`。对照 SM100：那里 `scale_c` 是运行期谓词（首条 UMMA 传 0 覆写、其余传 1 累加），因为 TMEM 无法像寄存器数组那样「一行代码清零」，只能借第一条 MMA 的覆写语义。SM90 有寄存器数组，直接用语言层面的 `= {0}` 更简单。 + +### 6.4 wave 与 warpgroup 的行定位 + +两级偏移共同确定「warpgroup `g`、wave `w` 覆盖哪些全局行」： + +```cpp +// 构造时（第 217 行）：warpgroup 定位到自己那 64 行 +auto a_desc = make_gmma_desc( + smem_a[0], math_wg_idx * WGMMA::M, 0); // mn_idx = math_wg_idx × 64 +// K 循环内（第 262–263 行）：wave 内再推进 local_idx × WAVE_BLOCK_M +a_desc.reg32_[0] = advance_gmma_desc_lo<...>(a_desc_base_lo, + local_idx * WAVE_BLOCK_M, // mn_idx：wave 偏移 + (k * WGMMA::K) % BLOCK_ATOM_K, // k_idx：atom 内 K 偏移 + atom_k_idx * BLOCK_M * BLOCK_ATOM_K); // offset：第几个 K-atom（合并时才 >0） +``` + +叠加后，warpgroup `g`、wave `w` 读的 A 行区间是 `[g×64 + w×WAVE_BLOCK_M, +64)`。累加器侧用 `shifted_accum = accum + kNumAccum × local_idx`（`local_idx` 即 wave 号 `w`）命中对应段。A 描述符的行偏移与累加器数组的 wave 分段**严格对齐**——这是「同一批线程既发 WGMMA 又持累加器又做 epilogue」能够正确工作的前提。 + +| `BLOCK_M` | math wg 数 | `WAVE_BLOCK_M` | wave 数 | wg0 覆盖行 | wg1 覆盖行 | +| --- | --- | --- | --- | --- | --- | +| 64 | 1 | 64 | 1 | 0–63（w0） | — | +| 128 | 2 | 128 | 1 | 0–63（w0） | 64–127（w0） | +| 256 | 2 | 128 | 2 | 0–63(w0)+128–191(w1) | 64–127(w0)+192–255(w1) | + +--- + +## 7. 矩阵 Tiling 层次 + +### 7.1 六级 tiling + +以 §5.6 的例子（`BLOCK_M=128, BLOCK_N=128, BLOCK_K=64, cluster=2`）为例，从全局到指令共六级： + +| 级别 | 尺度 | 承载者 | 说明 | +| --- | --- | --- | --- | +| L0 全局 | `M × N × K` | grid | persistent，`gridDim.x = kNumSMs` | +| L1 wave | `kNumSMs` 个输出块 | grid 的一轮 | `current_iter` 递增一次；L2 swizzle 在此层重排 | +| L2 cluster | **2 个独立输出块**（共享 A 或 B） | CTA-pair | **与 SM100 本质不同**：两个 CTA 各算各的块，仅共享一份操作数（§7.2） | +| L3 CTA tile | `BLOCK_M × BLOCK_N = 128 × 128` | 单 CTA | 落在本 CTA math warpgroup 的**寄存器累加器**（分散在 256 线程） | +| L4 k_block / stage | `BLOCK_M × BLOCK_K`（A）+ `BLOCK_N × BLOCK_K`（B） | SMEM ring 的一格 | 流水的调度单位，`kNumStages` 格在飞 | +| L5 WGMMA atom | `WGMMA::M × WGMMA::N × WGMMA::K = 64 × 128 × 16` | 一条指令 | 每 warpgroup 每 wave 发 `BLOCK_K / WGMMA::K = 4` 条消费一个 stage | +| L6 swizzle atom | `8 行 × 128 B` | SMEM 物理布局 | TMA 写入与 WGMMA 读出共用的最小图案单位 | + +K 方向总迭代数：`num_total_k_blocks = ceil_div(scheduler.current_shape_k, BLOCK_K)`（第 166/243 行）。用的是 `scheduler.current_shape_k` 而非 `shape_k`——k-grouped 变体里每个 group 的 K 长度不同，且这个值是**运行期**的，所以 K 主循环不能整体展开，只能展开内层的 `BLOCK_K/WGMMA::K` 次 WGMMA。 + +L2 那一行是 SM90 与 SM100 最大的 tiling 差异：SM100 的 L2 是「一条 `cta_group::2` UMMA 覆盖的 `2×BLOCK_M × BLOCK_N`单一块」，SM90 的 L2 是「两个本应独立、只是搭伴共享一份 A/B 的 `BLOCK_M × BLOCK_N` 块」。下面展开。 + +### 7.2 SM90 multicast：复制操作数，不切分输出块 + +`kNumTMAMulticast == 2` 时，一对 CTA（cluster rank 0 / rank 1）**各算一个完整的输出块**，仅通过 TMA multicast 共享其中一份操作数。共享 A 还是 B 由 `kIsTMAMulticastOnA`（= `cluster_n > 1`）决定： + +**情形：`kIsTMAMulticastOnA == true`（multicast A，`cluster_n = 2`，组在 N）** + +``` + n_block j n_block j+1 + ┌──────────────┐ ┌──────────────┐ + m_block i │ CTA0 算 │ │ CTA1 算 │ ← 两个独立输出块 + │ D[i, j] │ │ D[i, j+1] │ + └──────────────┘ └──────────────┘ + ▲ ▲ + │ 各自存整份 B（不同 n_block） + ┌────┴───────────────────┴────┐ + │ A[i] 由 rank0 一条 multicast │ ← 一份 GMEM 读 + │ 同时写进 CTA0/CTA1 的 smem_a │ 服务两个 CTA + └─────────────────────────────┘ +``` + +- 调度器把**相邻的 `n_block_idx`** 分给 cluster 内两个 CTA（`kIsMulticastOnA` → 组内 N 变化最快，§4.2），它们的 `m_block_idx` 相同 → A block 相同。 +- **A**：`tma::copy` 内部只由 `block_rank_in_cluster() == 0` 发一条 `SM90_TMA_LOAD_MULTICAST_2D`（CTA mask `0b11`），一份 GMEM 读同时写进两个 CTA 的 `smem_a[stage]`，tx 信号也同时打到两个 CTA 的 `full_barriers[stage]`。 +- **B**：每个 CTA 发自己的 `SM90_TMA_LOAD_2D`（`num_tma_multicast_b == 1`），各存自己 `n_block` 的整份 B。 +- **计算**：每个 CTA 用自己的 A 副本 + 自己的 B，独立算出一个 `128 × 128` 块，写进自己的寄存器累加器、自己的 `smem_d`、自己的 GMEM D。 + +**省什么、不省什么**：multicast 省的是 **A 的 L2/GMEM 读带宽**（一次读服务两个 CTA），这正是 §2.2 打分模型里 `num_bytes_l2_ab` 用 `block_m/cluster_n` 而非 `block_m` 的原因。但**不省 SMEM**：`load_block_m = BLOCK_M`（§2.3），每个 CTA 仍存整份 A。也不省算力：两个 CTA 各发自己的 WGMMA，tensor core 负载与不开 multicast 时一样。 + +**与 SM100 2-CTA UMMA 的对照**（这是理解两代 cluster 协作的钥匙）： + +| 维度 | SM90 multicast | SM100 2-CTA UMMA | +| --- | --- | --- | +| cluster 内两 CTA | 算**两个不同**输出块 | 合算**同一个**输出块（`UMMA_M=256`） | +| A 的 SMEM | 各存**整份**（`load_block_m = BLOCK_M`） | 各存**一半**（`load_block_m = BLOCK_M/cluster_n`） | +| tensor core | 各自发 `wgmma`（独立） | 两 SM 台成一条 `cta_group::2` UMMA | +| 省带宽 | ✓（一份读服务两 CTA） | ✓ | +| 省 SMEM | ✗ | ✓（stage 数可翻倍） | +| 能否运行期关闭 | ✓（`is_tma_multicast_valid`，§4.3） | ✗（硬件成对发射） | +| 累加器位置 | 各自的寄存器 | 各自的 TMEM（各存自己 128 行） | + +情形 `kIsTMAMulticastOnA == false`（multicast B，`cluster_m = 2`，组在 M）角色互换：两 CTA 拿相邻 `m_block_idx`、共享同一份 B，A 各存自己的。 + +### 7.3 索引计算：`get_global_idx` 与「无 rank 偏移」 + +SM90 与 SM100 共用 `get_global_idx`（§4.4），四个索引的双模板开关语义一致： + +```cpp +uint32_t m_idx = scheduler.get_global_idx<(kGemmType == MGroupedMasked), MN>(shape_m, BLOCK_M, m_block_idx); +uint32_t n_idx = scheduler.get_global_idx<(kMajorB == K), MN>(shape_n, BLOCK_N, n_block_idx, m_block_idx); +uint32_t k_a_idx = scheduler.get_global_idx<(kMajorA == MN), K >(shape_k, BLOCK_K, k_block_idx, m_block_idx); +uint32_t k_b_idx = scheduler.get_global_idx<(kMajorB == MN), K >(shape_k, BLOCK_K, k_block_idx, m_block_idx); +``` + +**关键差异：SM90 没有 SM100 那段「2-CTA 偏移」**。SM100 在算完 `m_idx/n_idx` 后会叠加 `block_rank_in_cluster() * load_block_m/n`（因为两 CTA 共算一块，各负责一半）；SM90 **完全不叠加**，因为两个 CTA 的 `m_block_idx/n_block_idx` 本就是调度器分好的**不同块**，各自的 `m_idx/n_idx` 已经是最终值。multicast 的「共享」只发生在 TMA 层（一份读写两处），不发生在索引层。 + +`n_idx` 的 `kWithGroupOffset = (kMajorB == K)`：B 的 group 维总是拼在外维（`make_tma_b_desc` 的 `gmem_outer_dim * num_groups`），K-major 时外维是 N 故加 `group * shape_n`；MN-major 时外维是 K，group 偏移由 `IndexType::K` 那条处理。`k_a_idx/k_b_idx` 的 `kWithGroupOffset = (major == MN)`：MN-major 时 K 是外维，k-grouped 的 `current_k_cumsum`/`current_k_start` 加在 K 上。 + +### 7.4 Stage 合并：用更大的 `BLOCK_K` 摊薄 `warpgroup_wait<0>` + +第 46–57 行： + +```cpp +// NOTES: this is for reducing the `warpgroup_wait<0>()` overhead +constexpr uint32_t kDoMergeStages = + kNumStages_ >= 10 and kGemmType == GemmType::Normal and + kMajorA == K and kMajorB == K and kNumMathThreads == 128; // ← 比 SM100 多了这个条件 +constexpr uint32_t kNumMinStages = 5; +constexpr uint32_t kNumStagesPerMerge = kDoMergeStages ? kNumStages_ / kNumMinStages : 1; +constexpr uint32_t BLOCK_K = BLOCK_K_ * kNumStagesPerMerge; // 64 → 128/192 +constexpr uint32_t kNumStages = kNumStages_ / kNumStagesPerMerge; +``` + +动机（注释）：*reducing the `warpgroup_wait<0>()` overhead*。每个 k_block 末尾 math warpgroup 都要做一次 `warpgroup_commit_batch()` + `warpgroup_wait<0>()`（等本批全部 WGMMA 完成），这是 **math warpgroup 级的串行同步**。把 2–3 个 64-宽的 stage 合成 1 个 128/192-宽的 stage 后，同步次数减半/减三分之二、每次 WGMMA 连发数从 4 增到 8/12，而**总 SMEM 占用和流水深度（字节数）不变**。 + +两处与 SM100 的差异： + +1. **多一个 `kNumMathThreads == 128` 条件**（即 `BLOCK_M ≤ 64`，单 math warpgroup）。SM100 无此限制。原因：单 math warpgroup 时并行度低、`warpgroup_wait<0>` 的串行开销更难被掩盖，合并收益最大；双 math warpgroup（`BLOCK_M > 64`）时寄存器压力已高，不再合并。 +2. **`kNumMinStages = 5`（SM100 是 8）、触发阈值 `kNumStages_ ≥ 10`（SM100 是 ≥ 8）**。SM90 的 `kNumMaxStages = 16`（SM100 是 32），阈值相应下调。 + +合并后 SMEM 布局的关键点在三处保持一致（与 SM100 同构）： + +1. **TMA**：`tma::copy` 内部 `BLOCK_INNER_ATOM = 128/2 = 64`，循环 2 次，第 i 次写到 `smem + i * BLOCK_M * 64`。 +2. **WGMMA 描述符**：构造时用 `BLOCK_ATOM_K = BLOCK_K / kNumStagesPerMerge = 64`（**不是** `BLOCK_K`，第 216 行），保证 `DG_STATIC_ASSERT(kSwizzleMode == BLOCK_ATOM_K * sizeof(dtype))` 即 `128 == 64×2` 仍成立；推进时 `atom_k_idx = k * WGMMA::K / BLOCK_ATOM_K`，偏移 `atom_k_idx * BLOCK_M * BLOCK_ATOM_K`。 +3. **stage 步长**：`a_desc_lo` 用 `SMEM_A_SIZE_PER_STAGE`（按合并后的 `BLOCK_K` 算）作为 stage 间的步长（第 245 行）。 + +举例：`kNumStages_ = 10` → `kNumStagesPerMerge = 2`、`BLOCK_K = 128`、`kNumStages = 5`；每个 k_block 发 `128/16 = 8` 条 WGMMA，`atom_k_idx ∈ {0,0,0,0,1,1,1,1}`，`(k*16) % 64 ∈ {0,16,32,48,0,16,32,48}`。 + +### 7.5 `BLOCK_M < 64` 时的算力浪费（有意为之） + +`WGMMA::M` 恒为 64，**与 `BLOCK_M` 无关**。当启发式因 `m ≤ 16`/`≤ 32` 选出 `BLOCK_M = 16`/`32` 时： + +- WGMMA 仍按 M=64 发射，读 64 行 A（其中 48/32 行是 SMEM 越界垃圾，§5.5），往 128 个线程的寄存器写 64 行的 D； +- `kNumWGMMAStoreThreads = WAVE_BLOCK_M × (128/WGMMA::M) = BLOCK_M × 2`（因 `WAVE_BLOCK_M = BLOCK_M`），只有前 `kNumWGMMAStoreThreads/32` 个 warp 的结果会写回； +- 其余 warp 在 epilogue 入口 `if (not do_wgmma_store) continue;`（第 289–290 行）直接跳过，寄存器里的垃圾 D 被丢弃。 + +| `BLOCK_M` | `WGMMA_M_PER_WARP=16` 行/warp | `kNumWGMMAStoreThreads` | 参与 store 的 warp | 有效行 | +| --- | --- | --- | --- | --- | +| 16 | warp0→0–15 | 32 | 仅 w0 | 0–15（全部有效） | +| 32 | w0→0–15, w1→16–31 | 64 | w0–w1 | 0–31（全部有效） | +| 64 | w0–w3→0–63 | 128 | w0–w3 | 0–63 | + +这是一个**明确的取舍**（与 SM100 §7.6 同理）：小 M 场景本来就是访存/延迟受限，选小 `BLOCK_M` 的目的正是注释里的 *avoid TMA L2 OOB bound*（不去 GMEM 白读 64 行），MMA 吞吐富余，用一条统一代码路径换掉「M=16/32 的另一套 WGMMA + 描述符/断言」的复杂度是划算的。 + +### 7.6 Tail-K：靠 TMA 零填充，无专用分支 + +**SM90 没有 SM100 那个 `kMayHaveTailKBlock` 编译期分支**。即便 K 是编译期常量，SM90 也不生成专门的 tail-K 代码，而是： + +```cpp +const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); // 向上取整 +for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + /* 每个 k_block 都发足 BLOCK_K/WGMMA::K 条 WGMMA，不区分是否尾块 */ +} +``` + +正确性由两层硬件保证： + +1. **TMA 零填充**：`CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE`（§2.6）使超出 `shape_k` 的元素在写入 SMEM 时自动填 0。尾块中不足 `BLOCK_K` 的部分被 0 填满。 +2. **tx 计数不变**：尾块的 TMA 仍搬整个 box（零填充也算字节），所以 `arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE)` 的期望值与实际到货字节始终匹配，barrier 协议无需特例。 + +于是尾块的 WGMMA 在零填充的 A/B 上照算，`0 × B = 0` 对累加器无贡献，D 结果正确。代价是尾块白发了一些「乘 0」的 WGMMA（最多浪费 `BLOCK_K-1` 个 K 元素的算力）。SM100 选择用 `for_each_static_prefix` 跳过这些无效 UMMA（省一点算力），SM90 选择不管（省一整块代码复杂度）——这是两代在「尾块处理」上的哲学差异。 + +> k-grouped 路径下 `current_k_start` 按 `kKAlignment=128` 对齐（§4.4），`BLOCK_K` 整除 128，所以尾块总是对齐的，`ceil_div` 不会真的向上取。 + +--- + +## 8. TMA 加载路径 + +### 8.1 调用形态 + +TMA warp（`kNumMathThreads/32 + 2` 的单个 lane）在每个 k_block 上按 `kMajorA/kMajorB` 四选一发射（第 186–197 行）： + +```cpp +if constexpr (kMajorA == cute::UMMA::Major::K) + tma::copy( + &tensor_map_a, &full_barrier, smem_a[stage_idx], k_a_idx, m_idx, num_tma_multicast_a, batch_idx); +if constexpr (kMajorA == cute::UMMA::Major::MN) + tma::copy( + &tensor_map_a, &full_barrier, smem_a[stage_idx], m_idx, k_a_idx, num_tma_multicast_a, batch_idx); +// B 同理，num_tma_multicast_b +``` + +模板参数序是 ``，函数参数序是 `(desc, barrier, smem_dst, inner_idx, outer_idx, num_multicast, batch_idx)`。**inner 恒为 SMEM 里连续的那一维**，所以 K-major 时 `(inner, outer) = (k, mn)`，MN-major 时 `(mn, k)`，两个 `if constexpr` 分支只是把实参顺序换了一下。与 SM100 唯一的形参差别：SM90 传的是 `num_tma_multicast_a/b`（每侧可能是 1 或 2，由 `is_tma_multicast_valid` 逐块决定，§4.3），SM100 传固定的 `kNumMulticast`。 + +### 8.2 swizzle-atom 循环 + +[common/tma_copy.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/common/tma_copy.cuh) 第 25–35 行： + +```cpp +constexpr uint32_t BLOCK_INNER_ATOM = get_inner_block_atom_size(); +// = kSwizzleMode == 0 ? BLOCK_INNER : kSwizzleMode / sizeof(dtype_t) +#pragma unroll +for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) + SM90_TMA_LOAD_2D::copy(desc_ptr, (uint64_t*)barrier_ptr, EVICT_NORMAL, + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, + inner_idx + i * BLOCK_INNER_ATOM, outer_idx); +``` + +TMA box 的内维被 host 压到 `swizzle/elem = 64` 个元素（128 B），所以一个逻辑上 `BLOCK_INNER` 宽的块要拆成 `BLOCK_INNER / 64` 条 TMA： + +| 场景 | `BLOCK_INNER` | atom | TMA 条数 | SMEM 目的地址步进 | +| --- | --- | --- | --- | --- | +| K-major A，未合并（`BLOCK_K=64`） | 64 | 64 | 1 | — | +| K-major A，合并后（`BLOCK_K=128`） | 128 | 64 | 2 | `BLOCK_M * 64` | +| MN-major A，`BLOCK_M=128` | 128 | 64 | 2 | `BLOCK_K * 64` | + +目的地址步进 `BLOCK_OUTER * BLOCK_INNER_ATOM` 正是 §5.3 描述的「atom 沿外维堆叠」布局，与 `make_gmma_desc` 的 SBO/LBO 推导严格对偶。 + +### 8.3 三种 TMA 变体（SM90 重点：rank0-only multicast） + +```cpp +if (num_tma_multicast == 1) { + cute::SM90_TMA_LOAD_2D::copy(...); // cp.async.bulk.tensor.2d…（单 CTA） +} else { + #if __CUDA_ARCH__ >= 1000 + cute::SM100_TMA_2SM_LOAD_2D::copy(...); // 带 .cta_group::2（SM90 不走） + #elif __CUDA_ARCH__ >= 900 + if (cute::block_rank_in_cluster() == 0) + cute::SM90_TMA_LOAD_MULTICAST_2D::copy(..., (1 << num_tma_multicast) - 1, ...); + #endif +} +``` + +- **1-CTA**：`SM90_TMA_LOAD_2D`，Hopper 就有的 `cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint`。`num_tma_multicast == 1` 时（未开 cluster，或 `is_tma_multicast_valid` 逐块退回单播）走这里。 +- **SM90 multicast**：`SM90_TMA_LOAD_MULTICAST_2D`，**只由 `block_rank_in_cluster() == 0` 发一条**带 CTA mask（`(1 << num_tma_multicast) - 1 = 0b11`）的 multicast load，一份 GMEM 读同时写进 cluster 内所有 CTA 的 SMEM，tx 信号也同时打到每个 CTA 的同名 mbarrier。rank 1 根本不发 A 的 multicast（被 `if (block_rank_in_cluster() == 0)` 拦下），但仍会在自己的 `full_barrier` 上收到 tx。 +- **SM100 2-CTA**：`SM100_TMA_2SM_LOAD_2D`（带 `.cta_group::2`）在 `__CUDA_ARCH__ >= 1000` 才编译，SM90 cubin 里不存在。两者语义不同：SM100 是「两 CTA 各自发射、tx 只记到 leader」；SM90 是「只 rank0 发射、tx 记到所有接收方」。 + +3D 变体（Batched）同理：`SM90_TMA_LOAD_3D` / `SM90_TMA_LOAD_MULTICAST_3D`，多一个 `batch_idx` 坐标。Cache hint 统一 `EVICT_NORMAL`，开头有一条静态断言确保 SM90/SM100 两套枚举值一致。 + +### 8.4 Descriptor prefetch + +```cpp +if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { // TMA warpgroup 的第一个 warp，单 lane + cute::prefetch_tma_descriptor(&tensor_map_a); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_cd); +} +__syncwarp(); +``` + +在 kernel 最开头、**任何同步之前**执行。descriptor 在常量内存里，首次 TMA 访问会有冷启动延迟，提前 prefetch 可以把它藏进 barrier 初始化的时间里。与 SM100 的差别：SM100 用 `warp_idx == 0` 全体（`prefetch.tensormap` 不是单线程指令），SM90 用 TMA warpgroup 首 warp 的**单个 elected lane**（`elect_one_sync()`）——因为 SM90 把 math warpgroup 排在前面（warp 0..kNumMathThreads/32-1），TMA 角色全在后面的 warpgroup，prefetch 自然也用 TMA warpgroup 的 warp。 + +### 8.5 `expect_tx` 字节数 + +```cpp +full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); // 第 198 行 +``` + +- `SMEM_*_SIZE_PER_STAGE` 用的是**合并后**的 `BLOCK_K`，与 TMA 实际搬运的字节数一致。 +- **不乘 `kNumTMAMulticast`**（对比 SM100 乘）：因为 SM90 每个 CTA 的 `full_barrier` 只统计**本 CTA 收到的字节**——multicast 的 A 字节已由硬件直接送到本 CTA 的 barrier，自己的 B 字节由自己的 load 送。两者相加恰好是一个 stage 的 A+B。 +- `full_barriers[i]->init(1)`（§5.2）：只有**一次** arrive（就是这条 `arrive_and_expect_tx`），由本 CTA 的单个 TMA lane 发出。 +- **顺序**：TMA 先发射（第 186–197 行），`arrive_and_expect_tx` 后执行（第 198 行）。合法——mbarrier 的 tx-count 是「期望值累加」，只要在该相位完成前把期望值补上即可。 + +### 8.6 完整的 TMA warp 循环 + +```cpp +while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx); // 逐块判定 + const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; + const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1; + const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); // ① 等消费者释放 + /* ② 算 m_idx / n_idx / k_a_idx / k_b_idx / batch_idx（无 rank 偏移） */ + /* ③ 发 A、B 的 TMA（各 1~2 条指令；multicast 侧只 rank0 发） */ + full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); // ④ + } +} +// 退出前（仅 multicast）：再等一轮 empty，确保 peer 不会再远程 arrive +if constexpr (kNumTMAMulticast > 1) + for (uint32_t i = 0; i < kNumStages; advance_pipeline(i)) + empty_barriers[stage_idx]->wait(phase ^ 1); +``` + +整个 TMA warp 就是这四步的无限重复，没有任何计算。它的推进速度只受 `empty_barriers` 的释放节奏限制，因此可以超前 WGMMA 多达 `kNumStages` 个 k_block。末尾那段「额外一轮 empty wait」（第 202–206 行）是 multicast 的**退出协议**：peer CTA 的 math warp 会通过 `arrive(target_cta)` 远程写本 CTA 的 `empty_barriers`，本 CTA 必须等这些远程 arrive 全部落地后才能退出（否则 SMEM 释放后 peer 的远程 arrive 是非法访问），详见 §10.6。 + +--- + +## 9. WGMMA 发射路径 + +> 本章对应 SM100 文档的「§9 UMMA 发射路径」。两者都用 SMEM 描述符（`_SS`），但 SM90 是 **warpgroup 级协同的 `wgmma.mma_async`**，累加器在寄存器；SM100 是 **单线程的 `tcgen05.mma`**，累加器在 TMEM。 + +### 9.1 WGMMA 类型选择 + +```cpp +using WGMMA = typename mma::sm90::BF16MMASelector::type; +``` + +`BF16MMASelector`（[mma/sm90.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/mma/sm90.cuh) 第 101–148 行）把 `BLOCK_N` 映射到一个 `BF16MMA>`，`N` 从 8 到 256、步进 8，共 32 个**硬编码的 C++ 类型特化**。每个特化的静态常量： + +```cpp +static constexpr int M = 64; // Hopper wgmma 的 M 原子固定 64 +static constexpr int N = N_; // = BLOCK_N +static constexpr int K = 16; // BF16 的 wgmma K 原子固定 16 +static constexpr int kNumAccum = M * N / 128; // = BLOCK_N / 2 +``` + +**与 SM100 的根本差异**：SM100 把 MMA 形状编成一个运行期可改的 `runtime_instr_desc`（32-bit），`UMMA_N` 甚至能逐块改写（swap-AB）；SM90 把形状直接固化成**编译期的 C++ 类型**，每个 `BLOCK_N` 对应一个不同的 `MMA_64xNx16_..._SS` 类，展展成一条固定的 `wgmma.mma_async.sync.aligned.m64nNk16.f32.bf16.bf16` PTX。所以 SM90 的 `BLOCK_N` 必须是编译期常量（它是模板参），不能运行期变——这也是为什么 SM90 没有 swap-AB（swap-AB 需要运行期改 `UMMA_N`，而 wgmma 的 N 固定）。 + +`_SS` 后缀 = A、B 都来自 Shared memory 描述符（对比 `_RS` 变体 A 来自寄存器）。BF16 GEMM 永远走 `_SS`。 + +### 9.2 SMEM 描述符（`GmmaDescriptor`） + +`make_gmma_desc()`（[mma/sm90.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/mma/sm90.cuh) 第 243–279 行）→ `make_smem_desc()` 填充 `cute::GmmaDescriptor` 的位域： + +| 字段 | 值 | 含义 | +| --- | --- | --- | +| `start_address_` | `__cvta_generic_to_shared(p) >> 4` | SMEM 地址，**16 B 为单位**（bit 0–13） | +| `layout_type_` | `to_gmma_layout_type<...>()` | `INTERLEAVE`（swizzle 0/16）/ `B32` / `B64` / `B128` | +| `leading_byte_offset_` (LBO) | `lbo >> 4` | atom 间在某一维的字节步长（bit 16–29） | +| `stride_byte_offset_` (SBO) | `sbo >> 4` | atom 间在另一维的字节步长（高 32 位） | +| `base_offset_` | `0` | — | + +**K-major** 的 SBO/LBO 推导： + +```cpp +DG_STATIC_ASSERT(kSwizzleMode == BLOCK_K * sizeof(dtype_t), "Unexpected value"); // 128 == 64 * 2 +const uint32_t stride_byte_offset = num_non_contiguous * BLOCK_K * sizeof(dtype_t); // 8*64*2 = 1024 +const uint32_t leading_byte_offset = 0; +``` + +注释解释为什么 LBO 是 0：*on K, there is only 1 atom as asserted previously*——那条静态断言保证每个 block 在 K 方向恰好只有一个 swizzle atom，于是「K 方向 atom 间步长」无意义。SBO = 一个 atom 的字节大小 = `8 行 × 128 B = 1024 B`，即 MN 方向相邻 atom 的步长。`num_non_contiguous = 128 / 16 = 8`（常量，**无 SM100 那个 base32 特例**，BF16 永走 8）。 + +**MN-major**： + +```cpp +constexpr uint32_t BLOCK_MN_ATOM = get_inner_block_atom_size(); // 64 +DG_DEVICE_ASSERT(mn_idx % BLOCK_MN_ATOM == 0); // 不允许 atom 内的 MN 偏移 +uint32_t stride_byte_offset = num_non_contiguous * BLOCK_MN_ATOM * sizeof(dtype_t); // 8*64*2 = 1024 +uint32_t leading_byte_offset = BLOCK_K * BLOCK_MN_ATOM * sizeof(dtype_t); +if constexpr (kSwizzleMode == 16) math::swap(stride_byte_offset, leading_byte_offset); +``` + +语义约定：swizzle 时 `{SBO, LBO}` 是 atom 在 `{K, MN}` 上的步长；非 swizzle（`kSwizzleMode == 16`，*means non-swizzling but interleaving*）时是 `{MN, K}`，所以要 swap。 + +### 9.3 单个 `a_desc_lo`/`b_desc_lo` + 算术推进（对照 SM100 的 32-lane） + +这是 SM90 与 SM100 描述符管理最直观的差别。math warpgroup 分支第 217–220 行： + +```cpp +constexpr uint32_t BLOCK_ATOM_K = BLOCK_K / kNumStagesPerMerge; +auto a_desc = make_gmma_desc(smem_a[0], math_wg_idx * WGMMA::M, 0); +auto b_desc = make_gmma_desc(smem_b[0], 0, 0); +const uint32_t a_desc_lo = __shfl_sync(0xffffffff, a_desc.reg32_[0], 0); // 从 lane 0 广播 +const uint32_t b_desc_lo = __shfl_sync(0xffffffff, b_desc.reg32_[0], 0); +``` + +然后在 K 循环内，每个 stage 的基地址用**算术加**得到（第 245–246 行）： + +```cpp +const auto a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16); +const auto b_desc_base_lo = b_desc_lo + stage_idx * (SMEM_B_SIZE_PER_STAGE / 16); +``` + +| | SM90（本文） | SM100 | +| --- | --- | --- | +| per-stage 描述符存储 | **单个** `a_desc_lo`（lane0 广播），每 stage 加常量步长 | **32 个 lane** 各存一个 stage 的 `a_desc.lo`，`__shfl_sync(..., stage_idx)` 取用 | +| stage 上限 | 无硬约束（算术加，`kNumMaxStages=16` 是 SMEM 预算定的） | `kNumStages <= 32`（lane 数硬约束） | +| 为何可行 | 所有 stage 连续且等大，`start_address` 随 stage 线性递增 | 同左，但选择用 lane 寄存器存而非算术加 | + +`__shfl_sync(0xffffffff, a_desc.reg32_[0], 0)` 把 lane 0 的值广播到全 warp。注释点明目的：*use `__shfl_sync` to encourage NVCC to use unified registers*——因为 `make_gmma_desc` 的输入（`smem_a[0]`、`math_wg_idx * 64`）本就是 warp-uniform 的，所有 lane 算出的 `a_desc` 相同，这次 shfl 在值上是恒等的，但向编译器**声明了“这个值是 warp 统一的”**，从而把它放进统一寄存器（uniform register），降低每-lane 寄存器压力。 + +注意只广播/存了 `.reg32_[0]`（低 32 位，含 `start_address_` 与 LBO），而 `.reg32_[1]`（高 32 位，含 SBO / base_offset / layout_type）是 **stage/k 无关的常量**，保留在 `a_desc` 结构体里不变。循环里只改 `.reg32_[0]`。 + +### 9.4 K 内层展开与描述符推进 + +```cpp +for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { // wave + auto shifted_accum = accum + WGMMA::kNumAccum * local_idx; + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { // 4 条（未合并） + const uint32_t atom_k_idx = k * WGMMA::K / BLOCK_ATOM_K; + a_desc.reg32_[0] = advance_gmma_desc_lo( + a_desc_base_lo, local_idx * WAVE_BLOCK_M, (k * WGMMA::K) % BLOCK_ATOM_K, atom_k_idx * BLOCK_M * BLOCK_ATOM_K); + b_desc.reg32_[0] = advance_gmma_desc_lo( + b_desc_base_lo, 0, (k * WGMMA::K) % BLOCK_ATOM_K, atom_k_idx * BLOCK_N * BLOCK_ATOM_K); + WGMMA::wgmma(a_desc, b_desc, shifted_accum, 1); + } +} +``` + +`advance_gmma_desc_lo` 的算式（第 237–241 行）： + +```cpp +return base + (((offset + mn_idx * BLOCK_K + k_idx * stride_k) * sizeof(dtype_t)) >> 4u); +// stride_k = (major == K) ? 1 : get_inner_block_atom_size() +// 注：模板形参名 `BLOCK_K` 在调用处实例化为 `BLOCK_ATOM_K` +``` + +- **K-major**：`mn_idx * BLOCK_ATOM_K` 是行偏移（每行 `BLOCK_ATOM_K` 个 K 元素连续），`k_idx * 1` 是行内 K 偏移，`offset = atom_k_idx * BLOCK_M * BLOCK_ATOM_K` 是跳到第 `atom_k_idx` 个 K-atom（仅合并时 > 0）。合计 `× 2 B >> 4` 换成 16-B 单位。 +- **MN-major**：`stride_k = BLOCK_MN_ATOM = 64`，因为 MN-major 下 K 是**外维**，K 前进 1 要跨过一整个 64 宽的 MN atom。 + +举例（未合并、`BLOCK_M=128`、`BLOCK_ATOM_K=64`、K-major、`local_idx=0`）：`k ∈ {0,1,2,3}`，`k_idx = (k*16)%64 ∈ {0,16,32,48}`，`atom_k_idx = 0`，于是 `a_desc.reg32_[0]` 依次 `+0, +2, +4, +6`（每条 WGMMA 消耗 `16 × 2 B = 32 B = 2` 个 16-B 单位）。与 SM100 的 `{0,2,4,6}` 完全一致。 + +因为 `local_idx`、`k`、`atom_k_idx`、`k_idx` 全是编译期常量（`#pragma unroll` + 常量 `BLOCK_*`），`advance_gmma_desc_lo` 的返回值被常量折叠，每条 WGMMA 的描述符增量是**立即数**（一条 IADD3）。 + +### 9.5 fence / arrive / commit / wait 序列 + +一个 k_block 内的 WGMMA 发射被四条指令包住（第 251–276 行），[ptx/wgmma.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/ptx/wgmma.cuh)： + +```cpp +full_barriers[stage_idx]->wait(phase); // ① 等 TMA 到货（A/B 已入 SMEM） +#pragma unroll +for (i = 0; i < kNumAccum*waves; ++i) ptx::warpgroup_fence_operand(accum[i]); // ② +ptx::warpgroup_arrive(); // ③ wgmma.fence.sync.aligned +for (local_idx) for (k) { /* 改 desc */ WGMMA::wgmma(a_desc, b_desc, shifted_accum, 1); } // ④ +ptx::warpgroup_commit_batch(); // ⑤ wgmma.commit_group.sync.aligned +#pragma unroll +for (i = 0; i < kNumAccum*waves; ++i) ptx::warpgroup_fence_operand(accum[i]); // ⑥ +ptx::warpgroup_wait<0>(); // ⑦ wgmma.wait_group.sync.aligned 0 +empty_barrier_arrive(stage_idx); // ⑧ 通知 TMA：本 stage 读完 +``` + +| 步 | PTX | 作用 | +| --- | --- | --- | +| ②⑥ | `asm volatile("" : "+f"(reg))` | **编译器屏障**（无指令）：把累加器寄存器标为读+写，防止编译器把 `accum` 的初始化/读取跨越 wgmma fence/commit/wait 重排（wgmma 异步读写累加器，编译器看不到） | +| ③ | `wgmma.fence.sync.aligned` | 把 fence 之前的 SMEM 读（full 等到后 A/B 可见）与累加器状态排序进 async wgmma proxy；标志「进入 wgmma 区」 | +| ⑤ | `wgmma.commit_group.sync.aligned` | 把此前发的所有 `wgmma.mma_async` 打包成一个 commit group，供 wait 等待 | +| ⑦ | `wgmma.wait_group.sync.aligned 0` | **阻塞**直到本 group 全部 wgmma 完成（结果落到累加器寄存器）。`<0>` = 等全部 | + +**为何每 k_block 都 `wait<0>`（而不是更深流水）**：`wgmma.mma_async` 从 SMEM 异步读 A/B，必须等它读完才能通过 `empty_barrier_arrive` 释放这个 SMEM stage（否则 TMA 可能在 wgmma 还在读时覆写 A/B）。所以顺序是 `wait<0>` → `empty_arrive`，每 k_block 一次。这就是 §1.1 说的「MMA **半同步**」：math warpgroup 在每个 k_block 的 `wait<0>` 处阻塞，暴露该 k_block 的 MMA 完成延迟（同一 group 内的 4 条 wgmma 彼此在 tensor pipe 里重叠，但跨 k_block 不重叠）。真正被重叠掉的是 **SMEM 预取**（TMA 超前 `kNumStages` 个 k_block，§10.5）。Stage 合并（§7.4）把每次 `wait<0>` 覆盖的 wgmma 从 4 增到 8/12、同步次数减半，正是为了摊薄这个半同步开销。 + +对照 SM100：`tcgen05.commit` 让 empty_barrier **异步跟踪** MMA 完成，MMA warp 不阻塞（发射完就继续）；且 SM100 的累加器在 TMEM，SMEM stage 的释放与 MMA 完成解耦。SM90 因累加器在寄存器 + `wait<0>` 的阻塞语义，必须同步等待。 + +### 9.6 math warpgroup 的完整循环 + +```cpp +while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + float accum[kNumAccum * waves] = {0}; // ① 本块累加器清零 + auto empty_barrier_arrive = [&](uint32_t s) { ... }; // ② 定义 arrive lambda + const auto num_total_k_blocks = ceil_div(current_shape_k, BLOCK_K); + for (k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { + a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE/16); // ③ 算术推 stage + full_barriers[stage_idx]->wait(phase); // ④ 等 TMA + /* ⑤ fence_operand → warpgroup_arrive → waves×4 条 wgmma → commit → fence_operand → wait<0> */ + empty_barrier_arrive(stage_idx); // ⑥ 释放 stage + } + if (not do_wgmma_store) continue; // ⑦ BLOCK_M<64 时无效 warp 跳过 + /* ⑧ epilogue：tma_store_wait<0> → NamedBarrier → STSM/st.shared → tma_store_fence → TMA store（§11） */ +} +``` + +注意 epilogue（⑧）**在同一个 math warpgroup 的同一个块循环内**，紧接 K 循环之后——这就是 §1.1「compute 与 epilogue 串行」的代码体现。一个输出块的 K 循环跑完（累加器已全部在寄存器），同一批线程立刻做 epilogue 把它们搬出；期间不发新 wgmma。下一块的 K 循环要等 epilogue 结束才开始（累加器数组要重置 `{0}`）。SM100 则是 MMA warp 与 epilogue warp 分离、TMEM 双缓冲重叠，两者调度形态截然不同。 + +--- + +## 10. 生产者-消费者同步 + +### 10.1 barrier 全景（只有两类） + +全部同步对象都是 `cutlass::arch::ClusterTransactionBarrier`（8 B SMEM 驻留的 mbarrier），**共 2 类**（SM100 是 5 类）： + +| barrier | 个数 | `init` 计数 | 等待者 | 到达者 | 语义 | +| --- | --- | --- | --- | --- | --- | +| `full_barriers[s]` | `kNumStages` | `1` | math warpgroup（每 CTA 全体 math 线程） | TMA warp 单 lane 的 `arrive_and_expect_tx(A+B)` + TMA 硬件 tx 完成 | SMEM stage `s` 已装好，A/B 可读 | +| `empty_barriers[s]` | `kNumStages` | `kNumTMAMulticast × kNumMathThreads/32` | TMA warp（每 CTA 单 lane） | 每个 math warp 的 `empty_barrier_arrive`（multicast 时跨 CTA） | stage `s` 已被 WGMMA 读完，可覆写 | + +只有一条流水（SM100 有三条）： + +``` +【A/B SMEM 环】 生产者 = TMA warp 消费者 = math warpgroup + full : TMA → math empty : math → TMA + +【累加器】 无环——每块就地 {0} 重置的寄存器数组，compute 与 epilogue 串行（§9.6） +【D SMEM】 单缓冲，不用 mbarrier，靠 tma_store_wait<0> + NamedBarrier 串行复用（§11.4） +``` + +**SM90 没有 SM100 的 TMEM 环（`tmem_full`/`tmem_empty`）**——因为累加器不是跨块复用的共享资源，而是每块私有的寄存器数组，compute→epilogue 在同一批线程内串行，不需要生产-消费握手。这是 SM90 同步结构比 SM100 简单得多的根本原因。 + +math warpgroup 既是 A/B 环的消费者，又（在 epilogue 里）是 D 单缓冲的生产者，是整条链的串行点。但它用整个 warpgroup（128/256 线程）而非 SM100 的单 warp：wgmma 是 warpgroup 协同指令（§3.3），且 epilogue 要把寄存器里的累加器搬进 SMEM，也需要这么多线程。 + +### 10.2 parity 相位约定与「首轮免等」 + +`Barrier::wait(p)` 对应 `mbarrier.try_wait.parity.shared::cta.b64 P, [bar], p`，含义是「等待奇偶性为 `p` 的那个相位**完成**」。A/B 环用显式相位变量（§4.5）： + +```cpp +uint32_t stage_idx = 0, phase = 0; +auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; // 只在回绕到 stage 0 时翻转 +}; +// 生产者（TMA） ：empty_barriers[stage_idx]->wait(phase ^ 1); +// 消费者（math）：full_barriers[stage_idx]->wait(phase); +``` + +第一轮（`phase = 0`）： + +- TMA 等 `empty` 的 parity **1**。新建的 mbarrier 处于 phase 0 且未完成，对 parity 1 的 `try_wait` 立即成功——这就是「缓冲初始为空，生产者前 `kNumStages` 轮直接穿过」的标准手法。 +- math 等 `full` 的 parity **0**，必须等第一次 tx 完成。 + +回绕一次后 `phase = 1`，两边等待的 parity 同时翻转，协议自洽。**SM90 没有 SM100 的「TMEM 环从 `current_iter` 推导 phase」那一套**（§10.1），因为无 TMEM 环；`current_iter` 在 SM90 只用于算块序号，不参与任何 barrier 相位。 + +### 10.3 arrive 计数逐条推导 + +**`full_barriers[s]->init(1)`** + +- 每个 CTA 的 TMA warp 单 lane 发一次 `arrive_and_expect_tx(A+B)` → 本地 arrive 1 次 + 设置期望 tx。这就是 `init(1)` 的那 1 次。 +- tx 字节来自：multicast 的 A（rank0 一条 load 写两个 CTA，tx 同时打到两 CTA 的 barrier）+ 自己的 B。 +- **对比 SM100**：SM100 `init(kNumMulticast)`（两 CTA 各 arrive 一次到 leader）；SM90 `init(1)`（每 CTA 自己的 barrier 只被自己 arrive，multicast 的 A 由硬件送 tx 而非送 arrive）。 + +**`empty_barriers[s]->init(kNumTMAMulticast × kNumMathThreads/32)`** + +`kNumMathThreads/32` = math warp 数（4 或 8）。`empty_barrier_arrive`（第 233–240 行）： + +```cpp +auto empty_barrier_arrive = [&](uint32_t s) { + if constexpr (kNumTMAMulticast == 1) { + lane_idx == 0 ? empty_barriers[s]->arrive() : void(); // 每 warp 1 次，共 math_warps 次 + } else { + auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster(); + lane_idx < kNumTMAMulticast ? empty_barriers[s]->arrive(target_cta) : void(); // 每 warp 2 次，分投两 CTA + } +}; +``` + +- **`kNumTMAMulticast == 1`**：每个 math warp 的 lane 0 arrive 一次，单 CTA 共 `kNumMathThreads/32` 次 → `init(1 × kNumMathThreads/32)`。✓ +- **`kNumTMAMulticast == 2`**：每个 math warp 的 lane 0/1 各 arrive 一次，`target_cta = lane_idx` 即 lane0→CTA0、lane1→CTA1。于是每个 CTA 的 `empty_barriers[s]` 收到：本 CTA 所有 math warp 投向自己的 1 次 + peer CTA 所有 math warp 投向自己的 1 次 = `2 × kNumMathThreads/32` → `init(2 × kNumMathThreads/32)`。✓ +- **为何 multicast 时 peer 的 math warp 要 arrive 本 CTA 的 empty**：multicast-on-A 时，本 CTA 的 `smem_a[s]` 是 rank0 一条 multicast load 写的（两个 CTA 各一份副本）。rank0 的 TMA warp 要重发下一个 multicast、覆写两个 CTA 的 `smem_a[s]`，必须知道**两个 CTA 的 math warp 都读完了各自的副本**。所以两 CTA 的 math warp 都向两 CTA 的 empty arrive。 +- **`is_peer_cta_alive == false` 的回退**（§4.3）：peer CTA 因落在矩阵边界外而无有效块时，两个 lane 都投向**本 CTA**（`block_rank_in_cluster()`）。于是本 CTA 的 empty 从自己的 math warp 收到 2 次/warp = `2 × kNumMathThreads/32`，恰好凑足 `init` 计数（peer 贡献 0，但 peer 已退出、不等自己的 empty）。 + +`arrive(target_cta)` 底层是 `mapa.shared::cluster` 把地址映射到目标 CTA 后 `mbarrier.arrive.shared::cluster`，即跨 CTA 远程 arrive（[ptx/ld_st.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/ptx/ld_st.cuh) 第 156–160 行的 `mapa_shared`）。 + +### 10.4 fence 序列 + +SM90 需要两类 fence：mbarrier 初始化可见性（async proxy）与 wgmma 的累加器/排序。共四处： + +| 位置 | 指令 | 作用 | +| --- | --- | --- | +| barrier 初始化后（第 121 行） | `fence_barrier_init()` = `fence.mbarrier_init.release.cluster` | 让 init 后的 mbarrier 对 **async proxy**（TMA）可见，cluster 作用域保证 peer CTA 也看得到 | +| wgmma 前（第 255 行） | `warpgroup_arrive()` = `wgmma.fence.sync.aligned` | 把此前 SMEM 读（full 等到后 A/B 可见）与累加器状态排序进 async wgmma proxy | +| 累加器读写两侧（第 254/272 行） | `warpgroup_fence_operand(reg)` = `asm("" : "+f"(reg))` | 编译器屏障，防止累加器访问跨越 wgmma fence/commit/wait 重排（§9.5） | +| STSM 写完 SMEM、TMA store 前（第 360 行） | `tma_store_fence()` = `fence.proxy.async.shared::cta` | generic proxy（stmatrix/st.shared）→ async proxy（TMA）的 SMEM 可见性 | + +**SM90 不需要 SM100 的 `tcgen05.fence::before/after_thread_sync` 与 `tcgen05.wait::ld`**——那些都是为 TMEM 异步读写服务的，SM90 无 TMEM。SM90 独有 `wgmma.fence/commit_group/wait_group` 三件套（为寄存器累加器的异步 wgmma 服务）与 `warpgroup_fence_operand`（编译器屏障）。两代的 fence 集合恰好互补，各自对应自己的异步代理（SM90：TMA + wgmma；SM100：TMA + tcgen05）。 + +### 10.5 稳态时序（一个输出块的完整生命周期） + +以 `kNumStages = 6`、`num_total_k_blocks = 128`（K=8192, BLOCK_K=64）、`BLOCK_M=128`（2 个 math warpgroup）为例，纵轴是时间： + +``` +TMA warp │ empty.wait │ TMA(s=0) │ empty.wait │ TMA(s=1) │ … │ TMA(s=5) │ empty.wait(阻塞) │ TMA(s=0) │ … + │ expect_tx │ │ expect_tx │ ▲ + ▼ ▼ ▼ ▼ │ +full[0] ────●───────────────────────────────────────────────────┐ │ +full[1] ──────────●──────────────────────────────────────────┼─┐ │ + │ wait(phase=0) │ │ │ +math wg │ full[0].wait │ fence→4×wgmma→commit→wait<0>→empty[0] │ full[1].wait │ 4×wgmma │ … │ (K循环完) │ EPILOGUE │ + ▼ + tma_store_wait<0> → NamedBarrier + → STSM/st.shared 写 smem_d + → tma_store_fence → NamedBarrier → TMA store +``` + +两个重叠关系（比 SM100 少一个）： + +1. **TMA 超前 math 最多 `kNumStages` 个 k_block**（受 A/B 环深度限）。epilogue 期间 TMA 仍在为下一块预取。 +2. **compute 与 epilogue 不重叠**（同一批 math 线程串行，§9.6）。这是 SM90 相对 SM100 缺的一个重叠：SM100 的块 i epilogue 与块 i+1 的全部 MMA 重叠（TMEM 双缓冲），SM90 的块 i epilogue 期间 math warpgroup 不算，只能靠 TMA 预取把下一块的 A/B 先搬好。 +3. **D 单缓冲不重叠**：下一块的 epilogue 必须等上一块 TMA store 把 `smem_d` 读完（`tma_store_wait<0>`，§11.4）。 + +### 10.6 退出协议（TMA drain，无最终 cluster sync） + +SM90 的退出比 SM100 简单：**没有最终的 `cluster_sync`，也没有 TMEM free**。唯一的安全网是 TMA warp 在 multicast 时的「额外一轮 empty wait」（第 202–206 行）： + +```cpp +// To safely deconstruct distributed shared barriers, we need another round of empty waits +if constexpr (kNumTMAMulticast > 1) { + for (uint32_t i = 0; i < kNumStages; advance_pipeline(i)) + empty_barriers[stage_idx]->wait(phase ^ 1); +} +``` + +危险场景：multicast 时 peer CTA 的 math warp 通过 `arrive(target_cta)` **远程**写本 CTA SMEM 里的 `empty_barriers`。如果本 CTA 的 TMA warp 已跑完并退出（SMEM 释放），peer 那次远程 arrive 就是一次非法访问。因为 TMA 超前 math 最多 `kNumStages` 个 k_block，TMA 跑完所有块时 math 可能还有至多 `kNumStages` 个 k_block 的 empty arrive 未发出；多等一轮 `kNumStages` 次 empty，就把这些尾巴的远程 arrive 全部排空（drain），确保退出后不会再有人写本 CTA 的 barrier。 + +这与 SM100 退出前的「额外一轮 `tmem_empty_barriers` wait」同构，只是 SM90 drain 的是 A/B 环的 empty（唯一会被远程 arrive 的 barrier），SM100 drain 的是 TMEM 环的 tmem_empty。SM90 不需要 SM100 那次最终 `cluster_sync_with_relaxed_arrive()`，因为没有 TMEM 需要两 CTA 对齐后释放；prologue 的 cluster sync（第 125 行）已保证 barrier 初始化互见。 + +--- + +## 11. Epilogue + +### 11.1 入口：与 math warpgroup 同一批线程 + +SM90 的 epilogue **不是独立的 warp 分支**，而是接在 math warpgroup 的 K 循环之后（第 279–383 行），由同一批线程串行完成（§9.6）。入口先做三个编译期检查与两个门控： + +```cpp +constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16); +constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes); // bf16→64, fp32→BLOCK_N +constexpr uint32_t WGMMA_M_PER_WARP = WGMMA::M / 4; // = 16 +DG_STATIC_ASSERT(BLOCK_M % 8 == 0, "Invalid swizzling atom"); +DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32, "..."); +DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N"); + +if (not do_wgmma_store) continue; // BLOCK_M<64 时，无效 warp 直接跳过（§7.5） + +if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) // 只由将发 TMA store 的线程等 + cute::tma_store_wait<0>(); // 等上一块的 TMA store 读完 smem_d（§11.4） +cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 0); // 广播“smem_d 已释放”（§12.3） +``` + +`WGMMA_M_PER_WARP = 16`：一个 `m64` WGMMA 的 64 行按 4 个 warp 平分，每 warp 16 行。epilogue 里 `warp_idx * WGMMA_M_PER_WARP` 把每个 warp 定位到自己那 16 行。 + +### 11.2 BF16：STSM 写回 + swizzle 逐行推导 + +BF16 输出走 `stmatrix`（第 297–344 行）： + +```cpp +DG_STATIC_ASSERT(kSwizzleDMode > 0, "Invalid swizzling type"); +DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization"); +for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { // wave + auto m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = accum + WGMMA::kNumAccum * local_idx; + for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) { // = BLOCK_N/8 次 + uint8_t* smem_ptr = /* swizzle 地址计算，见下 */; + // NOTES: only 16 lanes' addresses are used + ptx::SM90_U32x2_STSM_N::copy( + __float22bfloat162_rn({shifted_accum[i*4+0], shifted_accum[i*4+1]}), // 2 fp32 → 1 bf162 + __float22bfloat162_rn({shifted_accum[i*4+2], shifted_accum[i*4+3]}), + smem_ptr); + } +} +``` + +swizzle 地址计算（`kSwizzleDMode > 0`）： + +```cpp +constexpr uint32_t kNumBankGroupBytes = 16; +auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8); +auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes); +constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8; // 128B swizzle → true +auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8); +auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8); +col ^= row % (kSwizzleDMode / 16); // col ^= row % 8 +smem_ptr = (uint8_t*)smem_d + + warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) + // Warp 偏移（每 warp 16 行） + m_offset * kSwizzleDMode + // Wave 偏移 + atom_offset * BLOCK_M * kSwizzleDMode + // Swizzle atom 偏移（n 方向第几个 atom） + row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // atom 内偏移 +``` + +这与 SM100 §11.2 的 bank-group 置换逻辑**同构**（`col ^= row % 8` 把同列的8 行打散到 8 个物理 bank group，消除 bank conflict），差异在于： + +1. **写回指令是 STSM 而非 `st.shared`**：`SM90_U32x2_STSM_N`（[ptx/ld_st.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/ptx/ld_st.cuh) 第 43–52 行）= `stmatrix.sync.aligned.x2.m8n8.shared.b16`，一条写 2 个 8×8 bf16 矩阵（= 4 个 bf16/lane）。注释 *only 16 lanes' addresses are used*：`stmatrix.x2` 只用 lane 0–15 提供的 16 个行地址（两个 8×8 = 16 行），lane 16–31 的地址被忽略。 +2. **无需 TMEM load**：SM100 要先 `tcgen05.ld` 把累加器从 TMEM 读到寄存器、再 `st.shared`；SM90 的累加器**本就在寄存器**，直接 `__float22bfloat162_rn` 打包后 STSM。少一层搬运、少一个异步等待（无 `fence_view_async_tmem_load`）。 +3. **循环次数**：`i` 跑 `kNumAccum/4 = BLOCK_N/8` 次（每次 STSM 消耗 4 个累加器），恰好覆盖本线程的 `kNumAccum` 个寄存器。 + +### 11.3 FP32：`st.shared` 直存 + +FP32 输出时 `kSwizzleDMode == 0`（§2.3），不走 STSM（`stmatrix.b16` 只适用于 16-bit），而是逐行 `st.shared.v2.f32`（第 345–359 行）： + +```cpp +for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) { + auto m_offset = local_idx * WAVE_BLOCK_M; + auto shifted_accum = accum + WGMMA::kNumAccum * local_idx; + auto smem_d_0 = (float2*)(smem_d + (m_offset + warp_idx*WGMMA_M_PER_WARP + lane_idx/4 + 0) * BLOCK_N + (lane_idx%4)*2); + auto smem_d_1 = (float2*)(smem_d + (m_offset + warp_idx*WGMMA_M_PER_WARP + lane_idx/4 + 8) * BLOCK_N + (lane_idx%4)*2); + for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) { + ptx::st_shared(smem_d_0 + i*4, make_float2(shifted_accum[i*4+0], shifted_accum[i*4+1])); // 行 lane/4 + ptx::st_shared(smem_d_1 + i*4, make_float2(shifted_accum[i*4+2], shifted_accum[i*4+3])); // 行 lane/4+8 + } +} +``` + +这正是 `wgmma` 累加器 fragment 的标准布局：每 lane 持 `kNumAccum = N/2` 个值，对应两行（`lane/4` 与 `lane/4+8`，在本 warp 的 16 行内），列 `(lane%4)*2 + i*8 + {0,1}`。`smem_d_0 + i*4`（float2*）= `+i*8` 个 float，与列偏移一致。不 swizzle，D 就是行主序的 `BLOCK_M × BLOCK_N` 个 float，`TMA_D_BLOCK_N = BLOCK_N`（一条 TMA store 覆盖整块）。 + +### 11.4 单缓冲 D 的串行化：`tma_store_wait<0>` + +SM90 的 D 是**单缓冲**（§5.4），下一块的 epilogue 必须等上一块的 TMA store 把 `smem_d` 读完才能覆写： + +```cpp +if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) cute::tma_store_wait<0>(); // = cp.async.bulk.wait_group.read 0 +cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 0); +``` + +- **`wait<0>`（而非 SM100 的 `wait<1>`）**：`.read 0` = 等到「未完成的 TMA store group 数 ≤ 0」，即上一块的 store 全部把 SMEM 读完。SM100 用双缓冲 + `wait = wait<1>`，允许一个 store 在飞；SM90 单缓冲只能 `wait<0>` 完全排空。 +- **`.read` 后缀的语义**：cute 的 `tma_store_wait()` 展开为 `cp.async.bulk.wait_group.read N`，只等到「TMA 引擎不再需要读这块 SMEM」，**不**等数据真正落到 GMEM。这正是覆写 `smem_d` 所需的最弱条件。 +- **只部分线程等**：`cp.async.bulk.wait_group` 是**每线程**计数，而 `commit_group`（`tma_store_arrive`）只由 `threadIdx.x < BLOCK_N/TMA_D_BLOCK_N` 的线程发出（§11.5），所以只有它们持有非零的 group 计数、也只有它们能等。随后 `NamedBarrier` 把「smem_d 已释放」广播给全体 store 线程。 + +这是 SM90 epilogue 相对 SM100 的一个性能损失：单缓冲使块的 store 与下一块的 compute 无法在 D 侧重叠（§10.5）。 + +### 11.5 TMA store / reduce.add + +写完 SMEM、`tma_store_fence()` + `NamedBarrier` 后，前 `BLOCK_N / TMA_D_BLOCK_N` 个线程各发一条 TMA store（第 363–382 行）： + +```cpp +const auto m_idx = scheduler.get_global_idx<(not is_m_grouped_contiguous(kGemmType)), MN>(shape_m, BLOCK_M, m_block_idx); +DG_STATIC_ASSERT(kNumWGMMAStoreThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks"); +if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) { + auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N; + auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M; // 第 threadIdx.x 个 n-atom + using cute_tma_t = cute::conditional_t; // Batched 走 3D 变体 + cute_tma_t::copy(&tensor_map_cd, smem_ptr, n_block_idx * BLOCK_N + in_block_n_offset, m_idx); + cute::tma_store_arrive(); // cp.async.bulk.commit_group +} +__syncwarp(); +``` + +- **D 在 SMEM 里按 n-atom 分块**：`smem_d` 布局是 `[BLOCK_N/TMA_D_BLOCK_N 个 n-atom][每 atom BLOCK_M 行 × TMA_D_BLOCK_N 列]`，第 `threadIdx.x` 个 atom 在 `smem_d + threadIdx.x * TMA_D_BLOCK_N * BLOCK_M`。bf16 时 `TMA_D_BLOCK_N = 64`，`BLOCK_N=128` 拆成 2 条 store（由 thread 0/1 发）；fp32 时 `TMA_D_BLOCK_N = BLOCK_N`，一条 store。 +- **`m_idx` 的 `kWithGroupOffset = not is_m_grouped_contiguous(...)`**：m-grouped-contiguous 的 group 偏移已由 `make_tma_a_desc` 把 `m * num_groups` 拼进了 gmem 外维，故**不能**再加；masked/psum 变体需要加 `current_group_idx * shape_m`。 +- **`kWithAccumulation`**：当调用方传了 `c`，TMA 从 `SM90_TMA_STORE_*` 换成 `SM90_TMA_REDUCE_ADD_*`（`cp.reduce.async.bulk.tensor.…add`），在**写回路径上做 GMEM 原子累加**，实现 `D += A@B` 而不需先读 C。Batched 时同理换成 3D 变体（多传 `scheduler.current_group_idx` 作 batch 坐标）。 +- **`epilogue::transform`**：本 kernel 未显式调用 `apply_index_n`（n 索引直接用 `n_block_idx * BLOCK_N + in_block_n_offset`）。[epilogue/transform.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/epilogue/transform.cuh) 提供的 `EpilogueHeadSplits` 用于 attention 场景把 Q/K/V 三段拼接的 N 轴索引跳过中间 head 段（要求三段都能被 `STORE_BLOCK_N` 整除），BF16 GEMM 走恒等的 `EpilogueIdentity`。 + +### 11.6 dtype 转换 + +```cpp +// bf16：一次指令完成 2 个 fp32 → bf16 的舍入与拼接 +__float22bfloat162_rn({shifted_accum[i*4+0], shifted_accum[i*4+1]}) // → nv_bfloat162（32-bit） +// fp32：直存 +make_float2(shifted_accum[i*4+0], shifted_accum[i*4+1]) // → st.shared.v2.f32 +``` + +`__float22bfloat162_rn` 与 SM100 的 `cast_into_bf16_and_pack`（[common/math.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/common/math.cuh) 第 72–76 行）是同一个东西——**一条指令完成两个 FP32 → BF16 的 round-to-nearest-even 舍入与拼接**。累加器永远是 FP32（`MMA_64xNx16_F32BF16BF16_SS` 的 F32 累加），BF16 只出现在输入与最终输出，不存在中间累加降精度。 + +--- + +## 12. 其他重要机制 + +### 12.1 PDL：把 prologue 藏进前驱 kernel 的尾巴 + +Programmatic Dependent Launch 允许后一个 kernel 在前一个 kernel **尚未执行完**时就被派发到 SM 上跑，只要它不碰前驱的输出。SM90 把 `cudaGridDependencySynchronize()`（PTX: `griddepcontrol.wait`）**刻意推迟**到 prologue 的最后（第 132 行）： + +``` +line 86 if (warp_idx == kNumMathThreads/32 && elect_one_sync()) prefetch_tma_descriptor(a/b/cd); // tensormap 进 L2 +line 113 if (warp_idx == kNumMathThreads/32+1 && elect_one_sync()) { …2×kNumStages 个 mbarrier init…; fence_barrier_init(); } +line 125 kNumTMAMulticast > 1 ? cluster_sync_with_relaxed_arrive() : __syncthreads(); +line 132 cudaGridDependencySynchronize(); ◄── 真正的依赖点 +line 136 …构造 scheduler… +line 148 TMA 分支：warpgroup_reg_dealloc<48>() → 第一条 TMA load +line 208 math 分支：warpgroup_reg_alloc<224/248>() → 第一条 wgmma +``` + +被提到依赖点之前的三件事（prefetch、mbarrier init、cluster sync）都**只读写本 kernel 自己的资源**（常量内存里的 descriptor、随 CTA 分配的私有 SMEM），与前驱的数据无因果关系；而第一条 `cp.async.bulk.tensor` 要读的 A/B 很可能是前驱 kernel 刚写出的，故留在依赖点之后。与 SM100 同构，只是 SM90 把**寄存器再配置（`setmaxnreg`）放在依赖点之后**（SM100 无此步）。 + +三个使用注意点与 SM100 完全一致（共用 `DeviceRuntime`）：① 默认关闭，需 `deep_gemm.set_pdl(True)`；② 不开 PDL 也正确（`cudaGridDependencySynchronize()` 立即返回）；③ 全仓库无 `cudaTriggerProgrammaticLaunchCompletion()`，只做「等待方」。 + +### 12.2 `cluster_sync_with_relaxed_arrive()` 与 `__syncthreads()` 的选择 + +```cpp +(kNumTMAMulticast > 1) ? comm::cluster_sync_with_relaxed_arrive() : __syncthreads(); // 第 125 行，仅 prologue +``` + +SM90 **只在 prologue 用一次** cluster sync（multicast 时），epilogue 无最终 cluster sync（§10.6）——对比 SM100 的三处。`cluster_sync_with_relaxed_arrive()`（[comm/barrier.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/comm/barrier.cuh) 第 14–19 行）= `cluster_arrive_relaxed()` + `cluster_wait()`，只做控制流汇合、不承诺内存序（比 `cute::cluster_sync` 快）。这里 relaxed 够用的原因：承担跨 CTA 可见性的是它上面第 121 行的 `fence_barrier_init()`（`fence.mbarrier_init.release.cluster`），cluster arrive 退化为纯汇合。单 CTA（`kNumTMAMulticast == 1`）时退化为 `__syncthreads()`。 + +### 12.3 NamedBarrier：epilogue 内的两级同步 + +epilogue 用 `cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 0)`（第 295/361 行）而非 `__syncthreads()`： + +```cpp +// 第一道（写完 smem_d 前）：把“上一块 TMA store 已读完 smem_d”从等待线程广播给全体 store 线程 +if (threadIdx.x < BLOCK_N/TMA_D_BLOCK_N) tma_store_wait<0>(); +NamedBarrier::sync(kNumWGMMAStoreThreads, 0); +… STSM / st.shared 写 smem_d … +// 第二道（发 TMA store 前）：确保全体 store 线程的写入对 TMA 可见 +cute::tma_store_fence(); +NamedBarrier::sync(kNumWGMMAStoreThreads, 0); +… TMA store … +``` + +`NamedBarrier`（`barrier.sync.aligned id, num_threads`，id = 0）**只绑 `kNumWGMMAStoreThreads` 个线程**（`BLOCK_M=128` 时 = 256，即两个 math warpgroup 全部）。不用 `__syncthreads()` 是因为后者会同步**全部**线程——包括正在异步跑自己那摊活的 TMA warpgroup，那会把 TMA 与 math 强行拉回同步，破坏流水。NamedBarrier 只在 math 线程内部建立两道屏障，TMA warpgroup 不受影响。 + +### 12.4 三层断言防御体系 + +与 SM100 共用 [common/exception.cuh](../third_party/DeepGEMM/deep_gemm/include/deep_gemm/common/exception.cuh) 的三个宏：`DG_STATIC_ASSERT`（编译期 `static_assert`，零成本）、`DG_DEVICE_ASSERT`（`printf` + `trap`）、`DG_TRAP_ONLY_DEVICE_ASSERT`（只 `trap`）。 + +本 kernel 的关键 static 断言： + +| 断言 | 行 | 检查内容 | +| --- | --- | --- | +| `BLOCK_M % WGMMA::M == 0 or BLOCK_M < WGMMA::M` | 62 | `BLOCK_M` 要么≥ 64 且整除 64，要么小于 64（§7.5） | +| `cd_dtype_t` 是 `float`/`bfloat16_t` | 65 | 输出 dtype 白名单 | +| `WGMMA_A_SIZE_PER_STAGE <= SMEM_A + SMEM_B*kNumStages` | 79 | WGMMA 越界读不冲出 A/B 区（§5.5，但用 fp8 size，偏松） | +| 三个 SMEM 区都是 1024 B 的倍数 | 95 | `__align__(1024)` + swizzle-128B 的前提 | +| `kNumTMAThreads >= 128` | 155 | TMA warpgroup 至少 4 warp（才能抽出第三个 warp 发 TMA） | +| `kNumTMAMulticast <= 2` | 164 | cluster 最多 2 CTA | +| `kGemmType` 与 `kMajorA` 组合合法 | 177 | m-grouped 必须 A K-major | +| `BLOCK_M >= 64 or kNumMathThreads == 128` | 228 | `BLOCK_M < 64` 只能单 math warpgroup | +| `BLOCK_M % WAVE_BLOCK_M == 0` | 224 | wave 划分整除 | +| `BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N/TMA_D_BLOCK_N <= 32` | 284 | TMA store 对齐且条数 ≤ 32 | +| `WGMMA::kNumAccum % 4 == 0` | 300 | STSM x2 向量化（bf16） | +| `kNumWGMMAStoreThreads >= BLOCK_N/TMA_D_BLOCK_N` | 365 | store 线程够发所有 TMA | + +与 SM100 一样，**真正编出的 SM90 cubin 里 `DG_DEVICE_ASSERT` 一条都没有**（唯一一处在 `#else` 非-sm90 分支，第 388 行 `false and "This kernel only support sm_90a"`）。所有配置合法性都在 JIT 编译阶段被 `static_assert` 拦住；运行期检查只剩 `make_gmma_desc` 里 MN-major 分支的 `DG_DEVICE_ASSERT(mn_idx % BLOCK_MN_ATOM == 0)`（而 `mn_idx` 在编译期已知，实际不会触发）。 + +### 12.5 编译期常量折叠的收益链 + +JIT 全量模板特化触发的死代码消除链（以 `compiled_dims = "nk"` 为例）： + +``` +SHAPE_K != 0 + └─► shape_k 成为编译期常量（Normal 时 current_shape_k = shape_k） + └─► num_total_k_blocks 可被常量传播（但外层 K 循环仍为运行期 for，不展开） +BLOCK_M / BLOCK_N / BLOCK_K 均为常量 + └─► 内层 local_idx（wave）、k（BLOCK_K/WGMMA::K）循环 `#pragma unroll` 完全展开 + └─► advance_gmma_desc_lo 的 offset/mn_idx/k_idx 全为立即数 → 折叠成一条 IADD3 + └─► WGMMA::wgmma → make_index_sequence 展成 N/2 个寄存器操作数 + └─► epilogue 的 i 循环（kNumAccum/4）、swizzle 地址计算全展开 + └─► TMA_D_BLOCK_N、BLOCK_N/TMA_D_BLOCK_N 为常量 → store 线程数确定 +``` + +被 `if constexpr` 彻底消除的分支维度：`kMajorA`（2）× `kMajorB`（2）× `kNumTMAMulticast`（2）× `kWithAccumulation`（2）× `kIsBatchedMM`（2）× `kGemmType`（7）× `cd_dtype_t`（2）× `kDoMergeStages`（2）。**一份 cubin 里只存在一条完全直线化的路径**。 + +与 SM100 的一个区别：SM90 **没有 tail-K 分支可消除**（§7.6 本就不生成），所以常量折叠链比 SM100 短一节；但描述符算术折叠、内层循环展开、STSM 展开的收益与 SM100 一致。同样地，`SHAPE_M` 默认不编译进来（推理场景 M 常变，避免 JIT 缓存爆炸），代价是 M 方向的 tail（`BLOCK_M < 64` 的无效行、`do_wgmma_store` 筛选）保留为运行期逻辑。 + +--- + +## 13. 不变式、限制与已知坑 + +### 13.1 必须成立的不变式 + +下表是「改动这个 kernel 或其 host 启发式时不能破坏」的硬约束。左列任一被打破，结果要么是编译失败，要么是静默的数值错误 / 死锁。与 SM100 相比，SM90 的不变式**少了 TMEM 那一整类**（`2×UMMA_N ≤ 512`、`32 ≤ kNumTmemCols ≤ 512`、base address 必须是第 0 列），**多了寄存器预算与描述符 16 B 整除**这两类——根源都是「累加器在寄存器」。 + +| 不变式 | 由谁保证 | 破坏后的表现 | +| --- | --- | --- | +| `BLOCK_K_ == 64` | host `block_k = 128 / element_size(BF16=2)` | 一个 swizzle atom 的 K 字节数必须 128 B；否则 `make_gmma_desc` 的 SBO/LBO 全错 | +| `BLOCK_M % WGMMA::M == 0 or BLOCK_M < WGMMA::M`（即 `%64==0` 或 `<64`） | host `block_m` 候选 `{16,32,64,128,256}` | `static_assert`（第 62 行） | +| `WGMMA::M == 64`、`WGMMA::K == 16` | `BF16MMASelector` / `BF16MMA` 硬编码 | wgmma 原子形状非法；`WAVE_BLOCK_M`、`kNumWGMMAStoreThreads` 推导失效 | +| 三个 SMEM 区各自 1024 B 对齐且为 1024 的倍数 | `SMEM_*_SIZE` 构造 + `constexpr_align` | `static_assert`（第 95 行）；swizzle-128B atom 跨边界 → 数据错乱 | +| `SMEM_A_SIZE_PER_STAGE % 16 == 0`（描述符 start_address 以 16 B 为单位） | 1024 对齐蕴含 | `a_desc_lo + stage_idx*(SMEM_A_SIZE_PER_STAGE/16)` 的整除截断出错（§9.3） | +| `WGMMA_A_SIZE_PER_STAGE <= SMEM_A + SMEM_B*kNumStages` | 第 79 行断言（**用 fp8 size，偏松 2×**） | WGMMA 按 M=64 读 A 越界冲出 A/B 区 → 非法 SMEM 访问（见 §13.3 第 1 条） | +| `kNumTMAMulticast ∈ {1,2}` | host cluster 过滤 + 第 164 行 | `static_assert`；>2 CTA 的 cluster multicast 不支持 | +| `kNumSMs % cluster_size == 0` | host 过滤（§2.2） | cluster 跨 grid 边界，launch 失败 | +| `ceil_div(n, block_n) % cluster_size == 0`（masked/psum 布局） | host 过滤 | cluster 内两 CTA 落到不同边界外，multicast 的 `arrive` 计数凑不齐 → 死锁 | +| 每个 barrier 的 `init` 计数与 arrive 次数严格相等（`full=1`，`empty=kNumTMAMulticast × kNumMathThreads/32`） | §10.3 的逐条推导 | 计数多 → 永久等待（死锁）；计数少 → parity 提前翻转（数据竞争） | +| `kNumTMAThreads >= 128` | host 恒 128 | `static_assert`（第 155 行）；抽不出第三个 warp 专发 TMA | +| 寄存器预算：`kNumMathThreads×kNumMathRegisters + 128×48 <= 65536` | `setmaxnreg` 取值 224/248（§3.2） | 超过 64K 寄存器/SM → `warpgroup_reg_alloc` 后 launch/执行失败 | +| m-grouped ⇒ A 必须 K-major | host `DG_HOST_ASSERT(major_a==K)` + 第 177 行 | `get_global_idx` 的 group 偏移公式失效 | +| k-grouped contiguous ⇒ A/B 都必须 MN-major | host `DG_HOST_ASSERT`（hpp 第 270 行） | K 方向的 group cumsum 索引算错 | +| `BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N/TMA_D_BLOCK_N <= 32` | 第 284 行 | TMA store 不对齐，或条数超过 32 个 store 线程所能发射的上限 | +| `kNumWGMMAStoreThreads >= BLOCK_N/TMA_D_BLOCK_N` | 第 365 行 | store 线程不够发所有 TMA 条 | +| `WGMMA::kNumAccum % 4 == 0`（bf16 路径） | 第 300 行 | STSM x2 向量化（一次 4 个累加器）失败 | +| `swap_ab == 0` | host `DG_HOST_ASSERT(layout.swap_ab==0)`（§2.3） | SM90 路径根本不支持 swap-AB | +| D 必须 N-major | `make_tma_cd_desc`（§2.6） | 输出布局错乱 | + +### 13.2 功能与平台限制 + +- **仅 SM90 / `sm_90a`**。`#if __CUDA_ARCH__ >= 900` 之外只有 `DG_DEVICE_ASSERT(false and "This kernel only support sm_90a")`（第 386–389 行）。编译 flag 必须是 `--gpu-architecture=sm_90a`（`wgmma` 与 TMA multicast 都需要 `a` 后缀特性）。SM100 走同仓库另一套 kernel（`sm100_*`），流水组织完全不同。 +- **输入 dtype 硬编码 BF16**。`smem_a`/`smem_b` 一律 cast 成 `cutlass::bfloat16_t*`，`BF16MMASelector` 选出 `MMA_64xNx16_F32BF16BF16_SS`，累加器恒 FP32。输出 `cd_dtype_t` 只有 `float` 与 `bfloat16_t` 两种，没有 FP16。FP8/FP4 走别的 kernel。 +- **不支持 swap-AB**。host 侧 `DG_HOST_ASSERT(layout.swap_ab == 0)` 直接拦死。SM100 支持 swap-AB（MoE 小 M 场景），SM90 没有这条路径——小 M 只能靠 `block_m ∈ {16,32}` + 承受 M=64 的算力浪费（§7.5）。 +- **cluster 最多 2 CTA，且协作只有 TMA multicast**。SM90 的 cluster 协作是把一份 A 或 B **复制**给 pair 内两个 CTA（省 L2/GMEM 读带宽，**不省 SMEM**，两 CTA 各算各的块），没有 SM100 的 2-CTA UMMA（**切分**操作数，两 SM tensor core 合算一块，既省带宽又省 SMEM）。所以 SM90 开 cluster 不会增加 stage 数。 +- **无 split-K**。K 方向由单个 CTA 串行跑完整个 `num_total_k_blocks`，K 很大而 M/N 很小的瘦长 shape 无法靠增加并行度填满 SM。`kWithAccumulation`（`cp.reduce.async.bulk.tensor…add`）提供「多次调用累加到同一块 D」的能力，算是 host 层面的手工 split-K，但调用之间没有 kernel 内同步。 +- **persistent 调度是静态的**。`next_block_idx = (++current_iter) * kNumSMs + blockIdx.x`，纯算术映射，没有原子操作也没有工作窃取。好处是零同步开销、每个 CTA 独立推算自己的块序列；代价是块间代价不均时无法再平衡——`MGroupedMasked` 下各 group 的 `masked_m` 差异很大时，末 wave 长尾直接暴露在关键路径上。 +- **compute 与 epilogue 串行、D 单缓冲、无重叠**（§9.6 / §11.4）。同一批 math 线程先跑完整个 k-loop（累加器占满寄存器），再做 epilogue；epilogue 期间 tensor core 完全空闲，D 靠 `tma_store_wait<0>()` 串行复用单缓冲。这是 SM90 相对 SM100（MMA warp / epilogue warp 分离 + TMEM 双缓冲重叠）最主要的结构性性能损失，**无法靠调参消除**。 +- **无 tail-K 专用分支**（§7.6）。靠 TMA 的 `CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE` 零填充处理 K 非整除 `BLOCK_K` 的尾巴：wgmma 内层恒发满 `BLOCK_K/WGMMA::K = 4` 条，被零填充的 K 部分算入 0，结果正确但末尾块有无效算力。好处是代码路径单一（少一整块 SM100 那样的 `issue_tail_k_block` 复杂度）。 +- **无 Tensor Core 利用率控制**。SM100 有 `kTensorCoreUtilControl` 旋钮（给功耗受限集群做可复现对比），SM90 路径没有对应模板参数，也没有 `clock64()` 自旋降速逻辑。 +- **`gridDim.x == num_sms`**，且 `deep_gemm.set_num_sms()` 同时影响 grid 大小与调度器的 `kNumSMs`；设成小于物理 SM 数会有 SM 闲置，两者必须一致否则块分配会漏。 + +### 13.3 已知坑与代码瑕疵 + +按「踩到的概率 × 排查成本」排序： + +1. **`WGMMA_A_SIZE_PER_STAGE` 用 fp8 的 `sizeof` 而非 bf16**（第 78 行）。`WGMMA::M * BLOCK_K * sizeof(__nv_fp8_e4m3)`——`sizeof(__nv_fp8_e4m3) == 1`，但本 kernel 是 BF16（应 ×2）。这是从 FP8 kernel 移植遗留的瑕疵，导致第 79 行的越界防护**松了 2×**：真正要保证的是「WGMMA 按 `M=64` 读 A 时不冲出 A/B 区」，但断言只按一半大小检查。之所以不出事：`SMEM_A_SIZE_PER_STAGE` 本身已是 `BLOCK_M*BLOCK_K*2`，且 B ring 通常有极大裕量（`kNumStages × SMEM_B_SIZE_PER_STAGE`），WGMMA 越界读到的垃圾行在 epilogue 里根本不会被读（只读 `BLOCK_M` 行）。但若把 `BLOCK_M` 调到 `< 64` 同时 B ring 很窄，理论上可能踩到未分配 SMEM。这是本 kernel 最值得警惕的潜伏 bug，改 `BLOCK_M` 候选或 SMEM 布局时要手工复核真实的 2× 边界。 + +2. **compute 与 epilogue 串行、无重叠**（§9.6 / §11.4）。累加器在寄存器 ⇒ 一个输出块的 k-loop 与 epilogue 由同一批 math 线程串行完成，epilogue 期间 tensor core 空闲，D 单缓冲靠 `tma_store_wait<0>` 串行复用。看 nsys 时间线会看到「wgmma 段」与「STSM+TMA store 段」交替而非重叠，这不是调度没调好，是架构约束（对照 SM100 的 TMEM 双缓冲）。想量化 tensor core 利用率时要意识到这段空窗是固有的。 + +3. **`warpgroup_wait<0>()` 每个 k_block 都阻塞**（§9.5，「MMA 半同步」）。因为累加器在寄存器 + wgmma 从 SMEM 异步读，必须等这一 stage 的 MMA 全部完成，才能 `empty_barrier_arrive` 释放 SMEM stage 给 TMA 覆写。`kDoMergeStages`（§7.4，stage≥10 且 NT-Normal-单 math warpgroup 时把 `BLOCK_K` 放大 `kNumStagesPerMerge` 倍）正是为了摊薄这个 `wait<0>` 的频率——一次 wait 覆盖更多 K。若关掉 stage 合并，小 `BLOCK_K` 会让 `wait<0>` 成为流水的主要气泡。 + +4. **`BLOCK_M < 64` 时算力浪费 + store warp 空转**（§7.5 / §3.4）。`WGMMA::M` 恒为 64，即使 `BLOCK_M=16/32`，wgmma 仍按 M=64 发射，累加器 64 行里只有 `BLOCK_M` 行有用；且 `do_wgmma_store = BLOCK_M>=64 or warp_idx < kNumWGMMAStoreThreads/32` 只让部分 math warp 参与 store，其余空转。换来的是避免小 M 块的 TMA L2 OOB + 单一代码路径。profiler 里「wgmma 指令数 / 有效 FLOP」比值异常时先确认 `BLOCK_M`。 + +5. **multicast 省带宽不省 SMEM**（§7.2）。SM90 的 `load_block_m/n` **不除 cluster**（§2.3），每个 CTA 存整份操作数。`is_tma_multicast_valid` + `SM90_TMA_LOAD_MULTICAST` 让 rank0 的一条 load 把操作数复制进 pair 内两个 CTA 的 SMEM，省的是 L2/GMEM 读带宽。若误以为它像 SM100 2-CTA UMMA 那样也省 SMEM，会错误预估 stage 数。 + +6. **`is_peer_cta_alive` 为假时 `empty_barrier_arrive` 的回退**（§10.3，第 237 行）。multicast 下 `target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster()`——当 peer CTA 已退出（persistent 尾部块数不均），两个 lane 都把 arrive 投给**本 CTA** 以保持 `empty_barriers`（init 为 `kNumTMAMulticast × kNumMathThreads/32`）的计数平衡。误改这个三元表达式、或把 `lane_idx < kNumTMAMulticast` 的门槛动一动，都会在尾部死锁。 + +7. **`a_desc_lo` 单值算术推进依赖 16 B 整除**（§9.3，第 245 行）。`a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16)`——GmmaDescriptor 的 start_address 以 16 B 为单位，所以这里整除 16。SM90 只用**一个** `a_desc_lo`/`b_desc_lo` 标量按 stage 加常量步长推进（对比 SM100 用 32 个 lane 各存一个 stage 的描述符低位）。这依赖 `SMEM_A_SIZE_PER_STAGE % 16 == 0`（由 1024 对齐蕴含，§13.1）。 + +8. **epilogue STSM「only 16 lanes' addresses are used」**（第 337 行注释）。`SM90_U32x2_STSM_N` 的 `stmatrix.x2.m8n8.b16` 只用一个 warp 里 **16 个 lane** 的地址（另 16 个被忽略），swizzle 地址计算（`col ^= row % (kSwizzleDMode/16)`、`kHasShortcut` 分支）必须与之匹配。读这段地址算术时容易误以为 32 lane 都在写 SMEM。 + +9. **TMA warpgroup 里大量 warp 空转**（§3.1）。TMA warpgroup 128 线程 / 4 warp 里，稳态只有**第三个 warp的 1 个 elected lane** 发 TMA；w8 prefetch（一次性）、w9 init barrier（一次性）、w11 恒空转。但与 SM100 不同的是：math warpgroup 的 128/256 线程稳态**都**在算或搬，所以 SM90 的「线程填充率」仍是有效健康指标（§1.4 末行），不像 SM100 那样 `sm__warps_active` 会把人引向错误的「低效」结论。 + +10. **`__shfl_sync(0xffffffff, …)` 要求整 warp 收敛**（第 82 / 213 / 219 / 220 行）。`warp_idx`、`math_wg_idx`、`a_desc_lo`、`b_desc_lo` 都靠整 warp 的 `__shfl_sync` 取统一值（注释：*encourage NVCC to use unified registers*）。它们都位于 warp-uniform 的位置，合法；但若将来有人在 math warpgroup 里引入依赖 `lane_idx` 的分支，这几处会立刻变成未定义行为。 + +--- + +## 14. 小结 + +这份 kernel 的核心思想可以压缩成三句话： + +1. **累加器住在寄存器，是塑造这个 kernel 一切形态的根本约束**。没有 TMEM ⇒ `wgmma.mma_async` 必须 128 线程协同发射（每线程「拿着」自己那份累加器）⇒ math warpgroup 要 224/248 的寄存器配额、逼出 `setmaxnreg` 再配置 ⇒ compute 与 epilogue 由同一批线程串行、D 只能单缓冲 ⇒ `BLOCK_M < 64` 也得按 M=64 发射。SM100 把这些约束统统甩给 Tensor Memory，于是能做到 MMA warp / epilogue warp 分离、TMEM 双缓冲重叠、单线程发射 UMMA。两代架构的分野，物理根源就在「累加器放哪」。 +2. **把同步外化成两类 mbarrier 的 parity 相位**（`full` / `empty`）。A/B 的 GMEM→SMEM 全异步、可超前 MMA 多达 `kNumStages` 个 k_block；但 MMA→释放 SMEM 这一段因 `warpgroup_wait<0>()` 而**半同步**——这是累加器在寄存器 + wgmma 从 SMEM 异步读带来的必然妥协（§9.5），也是 stage 合并存在的理由。 +3. **把能变成编译期常量的东西全部变成编译期常量**。JIT 全量特化 ⇒ 内层 wave/k 循环完全展开、描述符增量成即数、所有 `if constexpr` 分支塌缩成一条直线路径。SM90 比 SM100 少一节（没有 tail-K 分支可消除，因为它本就不生成），但收益同源——这才是 DeepGEMM 相对通用库的真正护城河。 + +理解「累加器在寄存器」这一条之后，其余所有细节——为什么要 128 线程发 wgmma、为什么要 `setmaxnreg`、为什么 compute/epilogue 串行、为什么 D 单缓冲、为什么 `BLOCK_M<64` 仍发 M=64、为什么 multicast 不省 SMEM、为什么 `WGMMA_A_SIZE_PER_STAGE` 的 fp8 瑕疵不出事——都是它的自然推论。把这份文档与 [SM100 姊妹篇](./deepgemm_sm100_bf16_gemm_design.md) 对照着读，两代 Hopper/Blackwell tensor core 编程模型的差异会一目了然：一个是把累加器塞进寄存器的 warpgroup MMA，一个是把累加器外置到 Tensor Memory 的单线程 UMMA。 diff --git a/docs/image-1.png b/docs/image-1.png new file mode 100644 index 000000000..132769040 Binary files /dev/null and b/docs/image-1.png differ diff --git a/docs/image-2.png b/docs/image-2.png new file mode 100644 index 000000000..66beaada8 Binary files /dev/null and b/docs/image-2.png differ diff --git a/docs/image-3.png b/docs/image-3.png new file mode 100644 index 000000000..d09453a94 Binary files /dev/null and b/docs/image-3.png differ diff --git a/docs/image-4.png b/docs/image-4.png new file mode 100644 index 000000000..a5e6bb590 Binary files /dev/null and b/docs/image-4.png differ diff --git a/docs/image.png b/docs/image.png new file mode 100644 index 000000000..7fc1037c1 Binary files /dev/null and b/docs/image.png differ diff --git "a/docs/kernel \344\274\230\345\214\226.md" "b/docs/kernel \344\274\230\345\214\226.md" new file mode 100644 index 000000000..40c294626 --- /dev/null +++ "b/docs/kernel \344\274\230\345\214\226.md" @@ -0,0 +1,319 @@ +# InfiniTrain Kernel优化 +> 摘要：在单卡A100(40G)上，对llama3.2-1B模型训练进行kernel优化。基线版本每轮迭代82.05ms(3122 tok/s)：Adam优化器，cast、Fill等kernel占据了时延大头。进行1）影子权重+cast融合，2）adam向量化，3）Fill-> memsetasync，4）slice H2D->传参，5）RMSNorm融合，6）Embedding稀疏化。经过6轮优化后，稳态wall时间来到59.27ms(4319 tok/s)，吞吐提升**38%**。 + +[TOC] + +## 1 实验环境 + +| 项 | 配置 | +|---|---| +| GPU | NVIDIA A100-SXM4-40GB（虚拟机 passthrough，Ampere / sm_80） | +| CUDA / 驱动 | 12.4 / 550.107.02 | +| 工具链 | gcc 13.4、cmake 3.30.9、nsys 2023.4.4 | +| 构建（全文统一） | `-DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_ARCHITECTURES=80 -DBUILD_TEST=OFF -DNVTX_MODE=ON -DUSE_CUDA=ON -DUSE_NCCL=ON` | +| 模型 | LLaMA3.2-1B（16 层，n_embd 2048，n_head 32，GQA n_kv_head 8，vocab 128256） | +| 精度 | BF16 autocast：Linear/Matmul 走 Tensor Core，master 权重与 Adam 保持 FP32 | +| 训练配置 | batch_size 4 × seq_len 64，total_batch_size 256 | +| Profiling | Nsight Systems（nsys）+ NVTX，逐 step / 逐阶段标注 | + +采集与汇总： + +```bash +nsys profile --trace=cuda,nvtx,osrt,cudnn,cublas --sample=none --cpuctxsw=none \ + --output=nsys/out/llama3_5iter_nvtx_release --export=sqlite --force-overwrite=true \ + ./build/llama3 --device cuda --dtype bfloat16 \ + --input_bin data/llama3/tiny_shakespeare_train.bin \ + --llmc_filepath data/llama3/llama3.2_1B_fp32.bin \ + --num_iteration 5 --batch_size 4 --sequence_length 64 --total_batch_size 256 +# 注意：--export=sqlite 生成的 .sqlite 会比 .nsys-rep 旧，stats 必须带 --force-export=true +nsys stats --report cuda_gpu_kern_sum --report cuda_api_sum --report nvtx_sum --format table \ + --force-overwrite=true --force-export=true \ + --output nsys/out/llama3_5iter_release_stats nsys/out/llama3_5iter_nvtx_release.nsys-rep +# 本章表格的来源：逐步窗口/GPU busy 用 analyze_steps.py，阶段归属与逐 kernel timeline 用 extract_timeline.py +python3 nsys/out/analyze_steps.py nsys/out/llama3_5iter_nvtx_release.sqlite +python3 nsys/out/extract_timeline.py nsys/out/llama3_5iter_nvtx_release.sqlite +``` + +## 2 性能现状 + +### 2.1 nsys分析结果 + +非 profiling 的稳态每轮 **82.05 ms**，单 A100 吞吐 **3122 tok/s**。nsys 采集下稳态每轮 94.80 ms——两者的差是 CUPTI 开销，见下方口径说明； + +| step | wall (ms) | Σ kernel (ms) | Σ memcpy (ms) | kernels | GPU busy (ms) | GPU util | +|---|---|---|---|---|---|---| +| Step_0 | 228.9 | 54.14 | 3.18 | 3478 | 57.3 | 25.0%（warm-up，host-bound） | +| Step_1 | 102.1 | 83.74 | 3.53 | 3573 | 87.3 | 85.5% | +| Step_2 | 96.4 | 77.84 | 3.11 | 3574 | 80.9 | 84.0% | +| Step_3 | 89.9 | 73.79 | 2.40 | 3574 | 76.2 | 84.7% | +| Step_4 | 90.8 | 73.34 | 2.56 | 3573 | 75.9 | 83.6% | + +Step_1 Top kernel（分母 = 该窗口 Σ kernel 83.74 ms）： + +| # | kernel | n | Σ (ms) | 占比 | +|---|---|---|---|---| +| 1 | `AdamAccumulateGradKernel` | 114 | 31.96 | 38.2% | +| 2 | `CastKernel`（f32→bf16） | 356 | 9.47 | 11.3% | +| 3 | `ampere_s16816gemm_bf16_128x128_nt` | 49 | 4.64 | 5.5% | +| 4 | `ampere_bf16_s16816gemm_128x256_f2f_tn` | 49 | 4.51 | 5.4% | +| 5 | `BinaryBackwardKernel`(Mul) | 194 | 3.67 | 4.4% | +| 6 | `FillKernel` | 678 | 3.18 | 3.8% | +| – | **全部 GEMM（BF16 Tensor Core）** | 339 | **16.85** | **20.0%** | +| – | **全部 Cast（autocast）** | 661 | **10.74** | **12.8%** | + + +nsys 下稳态 GPU util 为 84–85%；非 profiling 下约 **93%**。 + +> **测量口径**：本次 trace 由基线 commit `34a52a8` 的 **Release** 构建采集。需要区分两类量：**nsys 采集值**含 CUPTI 逐调用插桩开销（同一二进制按 step 配对实测：nsys 给基线叠加 +10.02 ms，按每步 3478 次 `cudaLaunchKernel` + 1515 次 `cudaMemcpyAsync` 折算 ≈ **2.0 μs/次 API**）；**非 profiling 运行值**才是真实性能（稳态 **82.05 ms** / 3122 tok/s）。全章收益一律以非 profiling wall 结算，两种口径不可直接相减。 +![nsys timeline](image.png) + +**关键结论**： + +1）BF16 训练引入了大量 Cast kernel 完成 bf16↔FP32 转换：**661 个/轮、10.74 ms、占 kernel 时间 12.8%**；其中 f32→bf16 下转 356 个就占 9.47 ms（以权重下转为主）。 + +2）Adam 优化器是所有 kernel 中耗时最多的：**114 次发射、31.96 ms、占 38.2%**，单一 kernel 家族接近全部 GPU 时间的四成；其中 lm_head 和 embedding 的参数更新，单个 kernel 就占了 5 ms 左右的时间。 + +3）整个训练过程只使用了一个 stream。 + +4）FillKernel占据的时间也比较长，可以尝试用异步memset替换。 + +5）CrossEntropyForward API占据了大量的host时间完成D2H的memcpy_async。当然并不是copy本身需要很长的时间，这个异步DMA会排在一个异步执行队列中，他要等前面的执行完成才能执行。并且CrossEntrypyForward后续mean计算放在了CPU中，他依赖于这次D2H拷贝的数据。也就是说CrossEntropy引入了一个同步点。 + +6）SliceForward Kernel中有5次异步H2D memory copy，其中有几次时间很长推测原因为队列耗尽， + +### 2.2 Timeline分析 + +参考timeline.md + + +### 2.3 优化项 +- **cast kernel泛滥** +- **FP32 Adam 降低内存带宽要求** +- **Fill 泛滥** +- **单 CUDA stream 串行** +- **Forward 末的 1024 B loss DtoH 读回**：GPU 侧仅 3.46 us，却让 host 阻塞 16.30 ms 等整条队列排空（四个构建都测到 12.2–16.3 ms）。改 pinned memory 异步读回、或只在需要打日志的轮次读，可直接消掉这段阻塞。 +- **每轮 1515 次微小 HtoD**（合计仅 76 KB、平均 50.4 B/次）：标量是逐个拷上 GPU 的，可合并为一次批量拷贝。 + +## 3 单GPU优化 +### 3.1 Cast Kernel优化 + +Section 2 定位到 Cast 是 BF16 训练的第二大开销（Step_1 661 次 / 10.7 ms / 12.8%），且分**两类正交来源**，需分别治理： + +#### 3.1.1 优化一：影子权重（Shadow Weights）——消除gemm前的cast kernel + +有两个方案：1）将cast kernel和后续的gemm kernel进行融合；2）备份一份bf16格式的权重，在forward阶段和backward阶段直接用备份权重。 + +考虑cuBLAS不支持进行定制融合，并且手写gemm很难达到cuBLAS的性能，所以采用方案2。 + +**方案**：为每个 fp32 master 权重维护一份 **bf16 影子副本**，autocast 需要权重 bf16 时直接复用，跳过 CastKernel。 + +- **存储 / 注册**：`Adam` 持有 `shadow_weights_`（每参数一个 bf16 Tensor）；`thread_local unordered_map> g_shadow_registry` 以 master 裸指针为 key 映射到影子（`optimizer.cc`）。 +- **零成本刷新**：Adam 更新与影子刷新**融合进同一 kernel** `AdamAccumulateGradShadowKernel`——写回 fp32 权重的同时，在adam kernel内写回bf16影子权重。 +- **autocast 接入**：`cast_arg` 在 `arg->To(target)` 之前先查 `GetShadow(arg.get())`，命中且 dtype 相符则直接返回影子（`autocast.h:122-130`）。 + + +#### 3.1.2 优化二：elementwise 混合输入融合 + +**方案**：把cast操作**融进 elementwise kernel**——用模板 `` 在**寄存器内 widen**（`common::cuda::Cast`），以 f32 计算，**输出 dtype 保持不变**。 + + +#### 3.1.3 优化结果 + +| 配置 | commit | peak used | step5–10 稳态 | 吞吐量 | 性能提升 | +|---|---|---|---|---|---| +| A 基线（无优化） | `34a52a8` | 23143 MB | **82.05 ms** | 3122 tok/s | 0% | +| B ＋cast 优化 | `fe5591a` | 26001 MB | **77.13 ms** | 3321 tok/s | +6.4% | + +专项指标（A→B）：Cast kernel 数 661 → **292**（−369，−56%）；nsys 下 GPU kernel Σ（Step_3–4 稳态窗口）73.56 → 68.65 ms（−6.7%，−4.91 ms）。 + + +### 3.2 AdamAccumulateGradKernel优化——向量化访存 + +`AdamAccumulateGradShadowKernel` 在稳态迭代中占据了最长时间：5 iter 共 545 次发射 / 154.05 ms，占全部 GPU kernel 时间的 **约 42%**。分析显存带宽是 Adam kernel 的性能瓶颈，尝试对它做**向量化访存**改造。 + +#### 3.2.1 瓶颈定位：ncu 证明已吃满 DRAM，而非发射受限 + +ncu（`llama3_adam_shadow_detailed.ncu-rep`）： + +| 指标 | 实测 | 含义 | +|---|---|---| +| DRAM Throughput | **85.6%**（1.33 TB/s，峰值 1.555 TB/s） | 已接近带宽上限 | +| Compute (SM) Throughput | 26.33% | 算力大量闲置 | +| Long Scoreboard stall | 48.5 cycles（占 warp cycles 87.0%） | warp 全在等全局访存 | +| Scheduler No Eligible | 73.44% | 无可发射 warp，非发射瓶颈 | + +每元素搬 **30 B**（grad/param/m/v 各读 4 B = 16 B，param/m/v 各写 4 B = 12 B，shadow 写 2 B），算术仅 ~5 FLOP → 算术强度 **0.17 FLOP/B**。 + +#### 3.2.2 方案：128-bit 向量化访存 + +- **载体**：进行float4访存： +```c++ +VecT grad_vec = *reinterpret_cast(&grad_data[base]); +VecT param_vec = *reinterpret_cast(¶m_data[base]); +VecT m_vec = *reinterpret_cast(&m_data[base]); +VecT v_vec = *reinterpret_cast(&v_data[base]); +``` +- **尾部**：`num_elements % VecSize` 的余数由 kernel 内标量循环处理（`:192` / `:252`），对任意长度均正确。 + +#### 3.2.3 派发守卫：对齐 + SM 数缩放的最小尺寸 + +向量化并非无条件更优，派发处设两道守卫（`:285-289` 无影子 / `:339-344` 有影子）： + +1. **对齐**：grad/param/m/v/shadow 五个指针都须按 `sizeof(T) * VecSize` 对齐，否则宽访存非法 → 回退标量。 +2. **最小尺寸**：`num_elements >= sms * threads_per_block * vec_size`，如果尺寸过小而使用向量访存会导致sm无法用满，并且一个线程所占的寄存器和指令膨胀也会影响单线程性能。 + +#### 3.2.4 优化结果 + +**LLaMA3.2-1B 5-iter 实测** kernel优化结果： + +| 张量规模 n | 发射数 | 优化前 | 优化后 | 加速 | reg/thread | 有效带宽 | 占 A100 峰值 | +|---|---|---|---|---|---|---|---| +| 262,668,288（embedding / lm_head） | 8 | 5958.80 us | 5738.98 us | **1.038x** | 18→46 | 1322→1373 GB/s | 85.0→**88.3%** | +| 16,777,216（MLP up/gate/down） | 232 | 381.56 us | 371.26 us | 1.028x | 18→46 | 1319→1356 GB/s | 84.8→87.2% | +| 6,291,456 | 78 | 145.70 us | 142.42 us | 1.023x | 18→46 | 1295→1325 GB/s | 83.3→85.2% | +| 4,194,304（q/o_proj） | 78 | 98.83 us | 96.65 us | 1.022x | 18→46 | 1273→1302 GB/s | 81.9→83.7% | +| 2,048（RMSNorm 权重） | 160 | 4.20 us | 4.24 us | 0.99x | 18→19 | — | 守卫路由至标量 | + +**端到端稳态时延：三配置受控对照** + +三个配置 checkout 到同一 commit 链的三个点，用**完全相同**的构建 flags 各自全新构建，在**无 profiling** 环境下每配置跑 10 轮 × 重复 3 次。统计法：每步先取 3 次运行的中位数，再对窗口求均值（对离群步稳健）。 + +| 配置 | commit | peak used | step5–10 稳态 | 吞吐量 | 性能提升 | +|---|---|---|---|---|---| +| A 基线（无优化） | `34a52a8` | 23143 MB | **82.05 ms** | 3122 tok/s | 0% | +| B ＋cast 优化 | `fe5591a` | 26001 MB | **77.13 ms** | 3321 tok/s | +6.4% | +| C ＋adam 向量化 | `c07805a` | 26001 MB | **75.69 ms** | 3378 tok/s | +8.2% | + +优化结果并不理想，原因是显存带宽已经接近理论上限，即使使用向量化访存，也无法提升太多性能。 + +### 3.3 FillKernel 优化 + +Section 2.1 显示 Fill 家族每轮 **710 次 / 3.31 ms / 4.0%**（`FillKernel` 678 次 / 3.18 ms + bf16 版 32 次 / 0.13 ms，Step_1 窗口），单次均值 **3.6–4.7 μs**——单次 GPU 执行时间几乎全部是 launch/setup，是典型的 **launch-overhead 主导型**开销，与带宽或算力无关。分两步优化： + +- **第一步：删除冗余 Fill**——即前置的 `Fill(0)` 后续 kernel 会覆盖全部输出元素，Fill 是死代码。 +- **第二步**：`Tensor::Fill(0.0)` 在 CUDA 后端特化为 `cudaMemsetAsync`，走 DMA 引擎，不占 SM，脱离 kernel launch 路径。 + +#### 3.3.1 删除 transform.cu 6 处 Fill + +| # | 函数 | 位置 | kernel 覆盖证据 | +|---|---|---|---| +| 1 | `TrilBackward` | transform.cu:96 | `TrilBackwardKernel` L62–76：每 idx 无条件写 `grad_output[idx]` 或 `T(0)`，grid 覆盖 `[0, rows*cols)` | +| 2 | `TriuBackward` | transform.cu:183 | `TriuBackwardKernel` L150–164：同上（else 分支写 `T(0)`） | +| 3 | `TransposeForward` | transform.cu:275 | `TransposeForwardKernel` L194–222：`output[idx] = input[in_flat_idx]`，`num_blocks = ceil(num_elements/256)` 全覆盖，无 atomicAdd | +| 4 | `MaskBackward` (rows) | transform.cu:441 | `MaskLeadsBackwardKernel` L410–415：每 i 无条件写 `mask ? T(0) : grad_output[i]`，`rows*inner = grad_output->NumElements()` | +| 5 | `MaskBackward` (tail) | transform.cu:455 | `MaskBackwardKernel` L402–407：同上，`batch_size*mask_size = grad_output->NumElements()` | +| 6 | `RepeatInterleaveBackward` | transform.cu:568 | `RepeatInterleaveBackwardKernel` L520–540：**gather-reduce 而非 scatter**，每 thread 独占一个 `grad_input[idx]`，寄存器内累加 `repeat` 个 grad_output 后单点写 | + +#### 3.3.2 Fill 特化为 `cudaMemsetAsync` + +`BinaryBwd[Mul]` 的 grad_a/grad_b 广播归约、`EmbeddingBwd`（1，atomicAdd scatter）等场景不能删除fill kernel，但可以将kernel更新为`cudaMemsetAsync`。 + +#### 3.3.3 优化结果 + + +**端到端稳态时延**（`build_new/llama3`，无 profiling，10 iter × 3 次，每步取中位数再对窗口求均值）： + +| 配置 | commit | peak used | step5–10 稳态 | 吞吐量 | 性能提升 | +|---|---|---|---|---|---| +| C ＋adam 向量化（基线） | `c07805a` | 26001 MB | **75.69 ms** | 3378 tok/s | 0% | +| D ＋P0 Fill 删除 | `3bbdd76` | — | **75.41 ms** | 3395 tok/s | +0.4% | +| **E ＋P1 memsetAsync** | `9c1947ff` | — | **74.90 ms** | 3418 tok/s | **+1.0%** | + +> **为什么 wall 收益远小于理论 GPU 收益**：单stream，并没有真正并行起来。优化结果不能排除是随机导致的。 + +### 3.4 slice Kernel + +`SliceForward` 每次有 5 次微小 H2D 拷贝，其中几次因命令队列耗尽而耗时很长。每 step **257 次 Slice**，每次拷 5 个长度 = `num_dims`（典型 3–4）的 int64 数组，合计 **1285 次 H2D + 257 对 `cudaMallocAsync`/`cudaFreeAsync`**。 + +![slice kernel](image-1.png) + +#### 3.4.1 H2D -> 参数传递 + +**方案**：把 5 个元数据数组打包进定长 POD 结构 `SliceMeta`，经 kernel parameter space 按值传入，彻底消除 `cudaMallocAsync` + 5×`cudaMemcpyAsync` + `cudaFreeAsync` 整条序列。 + +```cpp +constexpr int kMaxDims = 8; +struct SliceMeta { + int64_t new_dims[kMaxDims], starts[kMaxDims], steps[kMaxDims]; + int64_t in_strides[kMaxDims], out_strides[kMaxDims]; +}; +// SliceForwardKernel<<<...>>>(in, out, meta, num_dims, total_elements); // meta 按值 → constant cache +``` + +#### 3.4.3 优化结果 + +| 配置 | commit | peak used | step5–10 稳态 | 吞吐量 | 性能提升 | +|---|---|---|---|---|---| +| E ＋P0+P1（基线） | `9c1947ff` | — | **74.90 ms** | 3418 tok/s | 0% | +| F ＋Slice 参数化 | `c2b2de39` | — | **71.83 ms** | 3564 tok/s | **+4.3%** | + +消除了Slice kernel前全部的5次memcpy。 + +![slice优化](image-2.png) + +### 3.5 RMSNorm Kernel 融合 + +根据timeline 分析，一轮forward迭代会出现33次 RMSNorm调用（16 层 × 2 个 RMSNorm + 收尾 ln_f）。每次由 Pow → Mean → AddScalar → Rsqrt → Mul → Mul **六个独立 kernel** 组成，forward 侧合计 **198 个 kernel/step**。本节把整个 RMSNorm 融合为单 Function：forward、backward 各一个 kernel，减少重复的显存读写和kernel launch开销。 + +#### 3.5.1 方案：autograd::RMSNorm + 融合 kernel（参考 LayerNorm） + +- **Function 层**：新增 `autograd::RMSNorm`。Forward 经 `Dispatcher` 调 `RMSNormForward`，一次返回 `{output, rstd}`；`rstd` 标记不可微，连同 `input/weight` 存入反向上下文。Backward 调 `RMSNormBackward` 返回 `{grad_input, grad_weight}`。 +- **数学**：forward 分两步计算（`n = x·rstd` → `y = n·w`，不合并为 `x·(rstd·w)`），保持与 composite 相同的舍入行为；backward 用 `S1 = Σ g·w·x`、`K = (S1/H)·rstd`、`dx = rstd·(g·w − n·K)`、`dw = Σ g·n`（跨 token 块 `atomicAdd`）。 +- **grad_weight 清零**：权重梯度跨 token 块累加，必须在 launch 前于 host 侧 `Fill(0.0)` 清零（同 stream 顺序保证可见性）。 +- **autocast**：注册 `"RMSNorm" → kFP32`（与 fp32 残差流语义一致）； + +#### 3.5.2 优化结果 + +**数值等价性**：step1 loss **4.898438 → 4.896993**——融合改变了归约顺序（块内 BlockReduce vs composite 的 Mean kernel），不再 bit-exact，差 0.0014（相对 3×10⁻⁴）为预期量级。 + +**端到端稳态时延**（同会话 G/H 配对：`git stash` 暂存融合改动 → reconfigure + 重建基线 → 10 iter × 3 次 → 恢复改动重建并经 sha256 校验；每步取中位数再对窗口均值，与 3.2.4 同口径）： + +| 配置 | commit | peak used | step5–10 稳态 | 吞吐量 | 性能提升 | +|---|---|---|---|---|---| +| G 基线（无融合） | `384dfc5` | 26001 MB | **71.83 ms** | 3564 tok/s | 0% | +| H ＋RMSNorm 融合 | `5a50c0c` | 25999 MB | **66.41 ms** | 3855 tok/s | **+8.2%** | + + +### 3.6 Embedding 稀疏梯度（Sparse Row）优化 + +![adam kernel](image-3.png) +timeline 分析看出optimizer阶段，有两个adam kernel占据了很长的时间（一个kernel就占据了5.7ms）。这两个kernel分别在进行Embedding和lm_head的权重更新。绝大多数时间都是用来读写权重，梯度和优化器状态（权重 [128256, 2048] fp32 = **1.05 GB**）。但是对于embedding来说每次只有迭代遇到的token row需要更新，不需要每次都更新完整的权重和优化器状态。 同理在embedding的backward阶段也不需要全量计算梯度。 + +lm_head所有权重都参与了计算，无法进行优化。 + +#### 3.6.1 方案：持久 grad buffer + 行清单 + 稀疏 Adam + +- **持久 grad buffer 不变量**：`EmbeddingBackward` 的 `grad_weight` 改为跨步常驻，维持不变量「除行清单记录的命中行外恒零」。首次 backward 一次性 `Fill(0)`（走 3.3.5 的 memset 快路径）；此后 `ZeroGrad` 只清零上一步的 **~256 行（≈2 MB）** 并递增 `generation` 使去重 stamp 失效。返回的梯度是 buffer 的**零拷贝视图**（`make_shared(*buffer, 0, dims)`）。 +- **行去重 claim**：backward kernel 改 2D grid，`(blockIdx.y==0, threadIdx.x==0)` 的线程经 `stamp[vocab]`（int32，按 `generation` 区分批次）`atomicCAS` 争抢写入 `row_list`（capacity=vocab，天然无溢出）；命中数 `count` 常驻**设备端**，host 全程零同步（消除 2.3 类 D2H 风险）。 +- **稀疏 Adam**：新增 `AdamSparseRows(Shadow)Kernel`，grid-strided 遍历 `row_list[0..*count)`，只对命中行更新 param/m/v（及 shadow），行内复刻既有 128-bit 向量化（`dim%4==0` + 对齐守卫，否则标量回退）。 + +#### 3.6.3 优化结果 + +- **数值等价性**（本次 3×3 配对，10 iter）：step1 `4.896993`、step2 `4.544800` 在全部 6 次运行**逐位一致**（首步更新与 dense LazyAdam 数学等价）；step3 差 4.3×10⁻⁵、step4 差 1.2×10⁻³，至 step5 起差值（0.9–8.1×10⁻³）。差异来源：① LazyAdam 语义（未命中行冻结 vs dense 用陈旧动量继续更新）；② 原子累加顺序（同一二进制重跑亦抖动）。 + +**端到端稳态时延**： + +| 配置 | commit | peak used | step5–10 稳态 | 吞吐量 | 性能提升 | +|---|---|---|---|---|---| +| I 基线（含 3.1–3.5 全部优化） | `5a50c0c` | 25999 MB | **66.46 ms** | 3852 tok/s | 0% | +| J ＋Embedding 稀疏梯度 | `6a0b1785` | 26008 MB | **59.27 ms** | 4319 tok/s | **+12.1%** | + +GPU侧一个大Adam Kernel消失，EmbeddingBackward前的全量grad_weight清零也消失了。 +![embedding优化](image-4.png) + +## 4 优化结果汇总 + +十个配置端到端稳态（非 profiling；每步取 3 次运行中位数、对 step5–10 求均值；吞吐量 = 256 tok/step ÷ 稳态时延）： + +| 配置 | commit | peak used | step5–10 稳态 | 吞吐量 | 性能提升 | +|---|---|---|---|---|---| +| A 基线（无优化） | `34a52a8` | 23143 MB | **82.05 ms** | 3122 tok/s | 0% | +| B ＋cast 优化 | `fe5591a` | 26001 MB | **77.13 ms** | 3321 tok/s | +6.4%（vs A） | +| C ＋adam 向量化 | `c07805a` | 26001 MB | **75.69 ms** | 3378 tok/s | +8.2%（vs A） | +| D ＋P0 Fill 删除 | `3bbdd76` | — | **75.41 ms** | 3395 tok/s | +0.4%（vs C） | +| E ＋P1 memsetAsync | `9c1947ff` | — | **75.18 ms** | 3405 tok/s | +0.7%（vs C） | +| F ＋Slice 参数化 | `c2b2de39` | — | **71.83 ms** | 3564 tok/s | +4.3%（vs E） | +| H ＋RMSNorm 融合 | `5a50c0c` | 25999 MB | **66.41 ms** | 3855 tok/s | +8.2%（vs G） | +| J ＋Embedding 稀疏梯度 | `6a0b1785` | 26008 MB | **59.27 ms** | 4319 tok/s | +12.1%（vs I） | + +> **82.05 → 59.27 ms（时延 −27.8%）**、**3122 → 4319 tok/s（吞吐 +38.3%）**；所有收益一律以非 profiling wall 结算。 diff --git a/docs/linear_acceleration_roadmap.md b/docs/linear_acceleration_roadmap.md new file mode 100644 index 000000000..88bf0c134 --- /dev/null +++ b/docs/linear_acceleration_roadmap.md @@ -0,0 +1,167 @@ +# InfiniTrain Linear 提速路线图（arch-conditional，含 DeepGEMM SM90 可行性结论） + +> **文档定位**：本文件是**设计 / 路线图**，不是实现记录。遵循「本阶段不改动代码」的约束——每个 Tier 都是后续可**独立执行**的单元，落地前需另行确认。 +> +> **关联文档**：本文是 [`deepgemm_sm90_bf16_gemm_design.md`](./deepgemm_sm90_bf16_gemm_design.md) 与 [`deepgemm_sm100_bf16_gemm_design.md`](./deepgemm_sm100_bf16_gemm_design.md) 的**应用侧续篇**——前两篇讲「DeepGEMM 的 kernel 怎么工作」，本篇回答「能否 / 如何用它改进 InfiniTrain 的 Linear，以及在真实硬件上最划算的提速路径是什么」。 +> +> **行号约定**：所有源码坐标以撰写时的仓库状态为准，末尾附[源码坐标索引](#八源码坐标索引)便于快速跳转。 + +--- + +## 一、摘要（结论先行） + +- **直接替换不可行（sm_120 / RTX 5090）**：SM90 kernel 的核心指令 `wgmma.mma_async` 是 sm_90a 专属；DeepGEMM host 侧 `bf16_gemm_nt` 对 `arch_major == 12` 直接 `DG_HOST_UNREACHABLE`；整条路径还被 `#if DG_TENSORMAP_COMPATIBLE` 包裹。**三重硬阻断**。 +- **即使在 sm_90a（H100 / H800）也是低 ROI**：cuBLAS 的 BF16 NT GEMM 在 Hopper 上本就会调用 wgmma kernel；DeepGEMM 的增量优势主要在 **FP8 / grouped / MoE**，标准稠密 Linear 收益有限。 +- **架构冲突（与硬件无关）**：DeepGEMM 是 `torch + pybind11 + Python` 的运行期 JIT 扩展；InfiniTrain 是**零 torch 依赖**的纯 C++/CUDA 框架。接入需抽离 device 头 + JIT runtime 并重接裸指针（约 700 行胶水 + 维护 fork + 运行期 nvcc + cache 目录 + 首调延迟）。 +- **推荐做法**：走**分层提速路线**（Tier 0–2 两架构通用、零 / 极少 device 代码即可拿到大部分收益），并在 Dispatcher 的 `Gemm` 注册点引入 **arch-conditional 后端抽象**（Tier 3），把「sm_90a 可选 DeepGEMM、sm_120 走 cuBLASLt / CUTLASS」统一收敛到 L1 kernel 层，**框架层保持零 `#ifdef`**。 + +--- + +## 二、可行性结论：DeepGEMM SM90 能否替换 Linear + +### 2.1 关键事实（均已核对源码） + +| 维度 | 事实 | 出处 | +| --- | --- | --- | +| 目标硬件 | 构建架构 `75;80;90;120`；部署目标 RTX 5090 = **sm_120**（`arch_major = 12`） | [`CMakeLists.txt:112`](../CMakeLists.txt)；DeepGEMM `device_runtime->get_arch_major()`（[`gemm.hpp:430`](../third_party/DeepGEMM/csrc/apis/gemm.hpp) 调用） | +| SM90 路径门禁 | `if (arch_major==9) sm90_bf16_gemm(); else if (==10) sm100_...; else DG_HOST_UNREACHABLE("Unsupported architecture")` | [`gemm.hpp:430-437`](../third_party/DeepGEMM/csrc/apis/gemm.hpp) | +| 编译期门禁 | 整个 `bf16_gemm_nt` 及其 arch 分派都在 `#if DG_TENSORMAP_COMPATIBLE` 之内 | [`gemm.hpp:403`](../third_party/DeepGEMM/csrc/apis/gemm.hpp) | +| 指令级阻断 | `wgmma.mma_async.sync.aligned.m64nNk16` **仅 sm_90a**；kernel 需 `--gpu-architecture=sm_90a` 编译 | SM90 详设 §2.7 / §13.2 | +| 依赖冲突 | DeepGEMM = `find_package(Torch REQUIRED)` + `pybind11_add_module(_C)` + 「real compilation is done via JIT」；InfiniTrain 顶层 CUDA 库只链 `glog / CUDA::cudart / CUDA::cublas`，无 torch / DeepGEMM | [`DeepGEMM/CMakeLists.txt:1,17,18,28`](../third_party/DeepGEMM/CMakeLists.txt)；[`CMakeLists.txt:133-140`](../CMakeLists.txt) | +| Linear 现状 | 全部经 `Dispatcher::Call("Gemm")` → `cublasGemmEx(..., CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)` | [`linear.cu:113-133`](../infini_train/src/kernels/cuda/linear.cu)；[`gemm.cu:45-75`](../infini_train/src/kernels/cuda/common/gemm.cu) | +| 布局天然匹配 | `Linear::Forward` 恒传 `transpose=true`；weight=`[out,in]=[N,K]` K-major、input=`[M,K]` K-major、output=`[M,N]` N-major | [`autograd/linear.cc:16`](../infini_train/src/autograd/linear.cc)；[`linear.cu:56-61`](../infini_train/src/kernels/cuda/linear.cu) | + +### 2.2 结论 + +- **sm_120：不可行。** 即使绕过 host 门禁，`wgmma` 也无法为 sm_120 汇编 / 执行。DeepGEMM 自己在非 sm90/sm100 场景是**回退到 cuBLASLt**（`smxx_cublaslt.hpp`，`CUBLAS_COMPUTE_32F_FAST_TF32`）——这恰好反证了正解方向：**稠密 Linear 的正解是 cuBLASLt / CUTLASS，而非 DeepGEMM 的 wgmma kernel。** +- **sm_90a：技术可接、性价比低。** 布局上 InfiniTrain Linear 是 DeepGEMM 的**最佳路径**（`transpose=true` 恰好是原生 `bf16_gemm_nt` 的 K-major NT），但相对 cuBLAS-BF16 的增量收益小，还要背 torch/JIT 依赖。仅在 **FP8 / grouped-MoE** 等特殊形态才值得（见[附录四](#四deepgemm-sm90-条件性集成设计仅-sm_90a附录)）。 + +--- + +## 三、推荐路线：分层提速（Tier 0–4） + +> 设计原则：**收益从高到低、改动从小到大**排布；Tier 0–2 两架构通用且几乎不碰 device 代码，Tier 3 提供可移植的路由骨架，Tier 4 是 sm_120 的兜底。 + +### Tier 0 — BF16 端到端打通（两架构，最高 ROI，device 代码 ≈ 0） + +**现状（已核对）**： +- `linear.cu` 已支持 BF16：`ToCudaDataType` 映射 `kBFLOAT16 → CUDA_R_16BF`（[`gemm.cu:33-34`](../infini_train/src/kernels/cuda/common/gemm.cu)）；`BiasCopyKernel` 已对 FP32/BF16 双分派（[`linear.cu:75`](../infini_train/src/kernels/cuda/linear.cu)）；`LinearBackwardBias` 的 BF16 指针**已正确用** `static_cast`（[`linear.cu:316-318`](../infini_train/src/kernels/cuda/linear.cu)）。 + > ⚠️ **记忆订正**：历史记录中「`LinearBackwardBias` BF16 cast bug」在当前代码**已修复**，Tier 0 **无需再动此处**。 +- **autocast 命中 Linear 已确认**：`kOpCastPolicyMap` 中 `{"Linear", CastPolicy::kLowerPrecision}` 存在且大小写正确（[`autocast.h:47`](../infini_train/include/autocast.h)）；autograd 边界在 [`function.cc:181-183`](../infini_train/src/autograd/function.cc) 对每个输入调用 `Autocast(type_, t)`，`type_="LinearFunction"` 经 `GetBaseOpName` 剥后缀得 `"Linear"`（[`autocast.h:13-19,92`](../infini_train/include/autocast.h)）→ 命中 → cast 到 `autocast_dtype`。**原「待核对缺口 #2」核查通过。** +- **BF16 路径端到端可达**：示例程序用**双参** `AutocastGuard(device.type(), dtype)`，`dtype` 来自 `--dtype` 标志（GPT-2 仅接受 `float32/bfloat16`，[`gpt2/main.cc:91,263-270,472`](../example/gpt2/main.cc)）。故 `--dtype=bfloat16` 即可让 Linear 走 BF16 → cuBLAS BF16。 + +**待补缺口 / 风险点**： +1. **backward 三处 dtype promotion FIXME hack**：`output_dtype = (compute_dtype==kBFLOAT16) ? kFLOAT32 : compute_dtype`（[`linear.cu:172`](../infini_train/src/kernels/cuda/linear.cu) / `244` / `296`）——即「BF16 计算、FP32 输出」。需确认这是**数值稳定性权宜**还是可在 autocast/autograd 修正后去除。 +2. **CUDA 默认 autocast dtype 是 FP16，而非 BF16（潜在坑，本会话新发现）**：`kDeviceDefaultDtype = {kBFLOAT16(CPU), kFLOAT16(CUDA)}`（[`autocast.h:74-77`](../infini_train/include/autocast.h)）。 + - 单参 `AutocastGuard(device_type)`（[`autocast.h:163-164`](../infini_train/include/autocast.h)）在 CUDA 上会默认 **FP16**。 + - 但 Linear 的 bias 路径**只覆盖 FP32/BF16、不含 FP16**：`BiasCopyKernel` 分派表是 `DispatchCudaFunc`（[`linear.cu:75`](../infini_train/src/kernels/cuda/linear.cu)），`LinearBackwardBias` 的 `switch` 也只有 `kFLOAT32/kBFLOAT16` 两个 `DISPATCH_CASE`（[`linear.cu:308-321`](../infini_train/src/kernels/cuda/linear.cu)）。 + - **结论**：CUDA 的 FP16 默认值与 Linear 实际 dtype 覆盖**不一致**——一旦有人用单参 guard 触发 FP16 autocast，带 bias 的 Linear 会在分派处失配。**Tier 0 应显式统一走 BF16**（示例已用双参 guard 规避，但默认值本身是漂移隐患）。 +3. **策略表维护漂移的旁证**：`{"Layernorm", kFP32}`（[`autocast.h:70`](../infini_train/include/autocast.h)，小写 n）与 `LayerNorm` 模块 `kType="LayerNormFunction"`→`GetBaseOpName`→`"LayerNorm"`（大写 N）**大小写不匹配**，导致 BF16 下 LayerNorm **静默走 BF16 而非提升到 FP32**（已确认缺陷）。这与 Linear 无关，但说明 `kOpCastPolicyMap` 的**双份手工维护**（数组 `kLowerPrecisionOps`/`kFP32Ops` + map）存在漂移风险——建议 Tier 0 顺带由数组生成 map 或加启动期一致性 `CHECK`。 + +**改动点**：`autocast` 策略表校验 / 生成化 + CUDA 默认 dtype 与 Linear 覆盖对齐 + `linear.cu` backward 输出 dtype 逻辑复核；**无新增 kernel**。 + +**预期收益**：sm_120 上 cuBLAS BF16 走 Blackwell tensor core，端到端约 **1.3–1.5×**；sm_90a 上自动走 wgmma。 + +### Tier 1 — TF32 快路径（两架构，FP32 存储模式） + +**现状**：`gemm.cu` 硬编码 `const cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F;`（[`gemm.cu:63-65`](../infini_train/src/kernels/cuda/common/gemm.cu)，注释明说「always use CUBLAS_COMPUTE_32F」）；`GemmParams`（[`gemm.h:21-46`](../infini_train/src/kernels/common/gemm.h)）**无 compute 精度字段**；全仓库 InfiniTrain 侧无 TF32。 + +**改动点**：`GemmParams` 增 `compute_type`（或 `allow_tf32` bool）→ `Gemm()` 据此在 `CUBLAS_COMPUTE_32F` 与 `CUBLAS_COMPUTE_32F_FAST_TF32` 间分支；由 Linear/autocast 或全局开关驱动。**约 20 行。** + +**预期收益**：需要 FP32 存储但想要 tensor core 加速时约 **1.3×**；精度损失介于 FP32 与 BF16 之间。 + +### Tier 2 — cuBLASLt + bias epilogue 融合（两架构） + +**现状**：bias **未融合**——forward 用独立 `BiasCopyKernel` + `beta=1`（[`linear.cu:68-82`](../infini_train/src/kernels/cuda/linear.cu)），backward-bias 用 `ReduceColumnsKernel`（[`linear.cu:138-156,305-321`](../infini_train/src/kernels/cuda/linear.cu)）。相比 epilogue 融合多一次全量读写。 + +**改动点**：`Gemm()` 增 cuBLASLt 分支，用 `CUBLASLT_EPILOGUE_BIAS` 把 bias 并入 GEMM epilogue；可参考 DeepGEMM 的 `smxx_cublaslt.hpp`（arch-agnostic）。**约 300 行**（含 desc / heuristic 管理）。 + +**预期收益**：带 bias 的 Linear（如 GPT-2）省一次带宽往返，约 **+3–5%**。 + +### Tier 3 — Dispatcher 层 arch-conditional 后端抽象（可移植核心） + +**目标**：把「按架构选 GEMM 后端」收敛到 **L1 kernel 层**（`Gemm` 注册点，[`gemm.cu:87`](../infini_train/src/kernels/cuda/common/gemm.cu)），框架层**零 `#ifdef`**，符合 L0–L7 架构约束。 + +**设计**：在 `Gemm()`（或新增 `GemmBackend` 选择器）内按**运行期** compute capability 路由： +- `sm_120` → cuBLASLt（Tier 2）或 CUTLASS Sm120（Tier 4） +- `sm_90a` → cuBLAS BF16（默认）；特殊形态可选 DeepGEMM `bf16_gemm_nt`（[附录四](#四deepgemm-sm90-条件性集成设计仅-sm_90a附录)） + +**arch 获取**：复用 InfiniTrain 自己的 DeviceGuard / compute-capability 查询，**不引入** DeepGEMM 的 `device_runtime`。 + +**改动量**：约 50 行；优先级 **P1（与 Tier 0/1 并行）**——它本身不提速，但让后续后端可插拔。 + +### Tier 4 — CUTLASS Sm120 collective builder（sm_120 专用，兜底） + +**触发条件**：Tier 0–2 触顶且需要定制 shape / 融合 epilogue。 + +**改动点**：新增 CUTLASS 3.x Sm120 collective 的 GEMM 后端，注册进 Tier 3 抽象。**约 500 行。** CUTLASS 已提供经验证的 Sm120 builder，无需自研 kernel。 + +--- + +## 四、DeepGEMM SM90 条件性集成设计（仅 sm_90a，附录） + +> 若未来部署 H100/H800 且**确需** DeepGEMM（例如 FP8 / grouped / MoE），集成设计如下。**默认不推荐**用于标准稠密 Linear。 + +- **布局映射（天然匹配）**：InfiniTrain Linear 恒 `transpose=true`（[`autograd/linear.cc:16,29`](../infini_train/src/autograd/linear.cc)），三个 GEMM 对应 DeepGEMM： + - **forward**：`output[M,N] = input[M,K] @ weight[N,K]^T` → `bf16_gemm_nt(input, weight, output)`（A/B 均 K-major、D N-major，命中最佳路径，[`gemm.hpp:404-438`](../third_party/DeepGEMM/csrc/apis/gemm.hpp)） + - **dgrad**（`LinearBackwardInput`）：`bf16_gemm_nn` / `bf16_gemm_tn`（按 transpose 组合，内部 transpose 适配，[`gemm.hpp:440-454`](../third_party/DeepGEMM/csrc/apis/gemm.hpp)） + - **wgrad**（`LinearBackwardWeight`）：`bf16_gemm_tn`（reduce over bs） +- **bias**：DeepGEMM **无 bias 融合** → 仍需保留独立 kernel（此处无收益；融合要走 Tier 2 cuBLASLt）。 +- **dtype**：a/b 必须 BF16；d 可 BF16 或 FP32（[`gemm.hpp:421-423`](../third_party/DeepGEMM/csrc/apis/gemm.hpp)）→ 与现有 backward「BF16 计算、FP32 输出」的 promotion hack **兼容**。 +- **接入方式（二选一）**： + 1. 抽离 `deep_gemm/include` device 头 + `csrc/jit` runtime，重接到 InfiniTrain 裸指针 + 自建 `CUtensorMap`（避免 torch），**维护内部 fork**。 + 2. 直接依赖 libtorch/ATen（与 InfiniTrain 零 torch 架构**冲突，不推荐**）。 +- **成本 / 风险**：约 700 行胶水 + 运行期 nvcc / JIT cache + 首调延迟 + fork 维护；且 H100 上相对 cuBLAS-BF16 收益有限。**建议仅在 FP8 / grouped / MoE 时启用。** + +--- + +## 五、优先级与预期收益 + +| Tier | 适用硬件 | 改动量 | 预期收益 | 主要风险 | 优先级 | +| --- | --- | --- | --- | --- | --- | +| **0** BF16 端到端 | 两者 | ≈0（+校验 / hack 清理） | **1.3–1.5×** 端到端 | 收敛性（需 `compare_loss` 验证）；CUDA 默认 FP16 与 Linear 覆盖不一致 | **P0** | +| **1** TF32 快路径 | 两者 | ~20 行 | ~1.3×（FP32 模式） | 精度损失 | **P1** | +| **2** cuBLASLt + bias | 两者 | ~300 行 | +3–5%（带 bias） | 复杂度 | **P2** | +| **3** 后端抽象 | 两者 | ~50 行 | 使能可移植路由 | 抽象设计 | **P1（与 0/1 并行）** | +| **4** CUTLASS Sm120 | sm_120 | ~500 行 | shape 相关 | 集成工作量 | **P3（兜底）** | +| ✗ DeepGEMM SM90 | 仅 sm_90a | ~700 行 + fork | 低（cuBLAS 已用 wgmma） | 硬件 / 依赖 / JIT | **不推荐（除 FP8/MoE）** | + +> **注**：理论 GEMM 提速会被 host 侧瓶颈（CrossEntropy CPU 归约、StackForward HtoD）稀释，**端到端收益低于 device-time 估计**。 + +--- + +## 六、验证方法 + +- **数值**：每 Tier 后用 [`scripts/compare_loss.py`](../scripts/compare_loss.py) 对齐 loss 曲线（dtype 改动尤其需要）。 +- **性能**：用 **NVTX 标注 + nsys** 测 GEMM 段 device-time（**不要用 `PROFILE_MODE`**，其逐算子同步会使吞吐虚降约 2×）。 +- **后端确认**：nsys 里核对 cuBLAS 是否切到 **tensor-core 变体**（BF16 / TF32 kernel 名），而非 FP32 SIMT。 + +--- + +## 七、假设与边界 + +- 假设部署目标为 **RTX 5090（sm_120）**；若实际为 H100/H800，Tier 0–3 同样适用，[附录四](#四deepgemm-sm90-条件性集成设计仅-sm_90a附录)的 DeepGEMM 集成才进入可选范围。 +- 撰写时 shell **无 GPU**（`nvidia-smi` 无设备），硬件结论基于**构建配置**（[`CMakeLists.txt:112`](../CMakeLists.txt)）与既有实测记忆，未在本文成文时在线复测 compute capability。 +- **本阶段不产出代码改动**；任一 Tier 的落地为后续独立任务。 + +--- + +## 八、源码坐标索引 + +| 主题 | 文件 : 行 | 说明 | +| --- | --- | --- | +| 构建架构 | [`CMakeLists.txt:112`](../CMakeLists.txt) | `75;80;90;120` | +| InfiniTrain CUDA 库依赖 | [`CMakeLists.txt:133-140`](../CMakeLists.txt) | 只链 glog / cudart / cublas，无 torch | +| DeepGEMM 依赖 & JIT | [`DeepGEMM/CMakeLists.txt:1,17,18,28`](../third_party/DeepGEMM/CMakeLists.txt) | JIT / pybind11 / Torch / `_C` | +| DeepGEMM arch 门禁 | [`gemm.hpp:403,430-437`](../third_party/DeepGEMM/csrc/apis/gemm.hpp) | `#if DG_TENSORMAP_COMPATIBLE` + arch_major 分派 | +| `bf16_gemm_nt/nn/tn` | [`gemm.hpp:404-454`](../third_party/DeepGEMM/csrc/apis/gemm.hpp) | NT 最佳路径 + transpose 适配 | +| Linear transpose 恒 true | [`autograd/linear.cc:16,29`](../infini_train/src/autograd/linear.cc) | weight=`[N,K]` PyTorch 约定 | +| LinearForward / bias | [`linear.cu:56-133`](../infini_train/src/kernels/cuda/linear.cu) | BiasCopyKernel + Gemm | +| backward promotion hack | [`linear.cu:172,244,296`](../infini_train/src/kernels/cuda/linear.cu) | BF16→FP32 输出 | +| LinearBackwardBias（已修） | [`linear.cu:308-321`](../infini_train/src/kernels/cuda/linear.cu) | BF16 用 `nv_bfloat16*` | +| Gemm compute_type 硬编码 | [`gemm.cu:63-69`](../infini_train/src/kernels/cuda/common/gemm.cu) | `CUBLAS_COMPUTE_32F`（Tier 1 改动点） | +| `ToCudaDataType` | [`gemm.cu:29-41`](../infini_train/src/kernels/cuda/common/gemm.cu) | FP32 / BF16 / FP16 | +| `GemmParams`（无 compute 字段） | [`gemm.h:21-46`](../infini_train/src/kernels/common/gemm.h) | Tier 1 需扩字段 | +| autocast 策略表 | [`autocast.h:45-77`](../infini_train/include/autocast.h) | Linear@47 / Layernorm@70 / 默认 dtype@74-77 | +| autocast 边界 | [`function.cc:181-183`](../infini_train/src/autograd/function.cc) | autograd 边界按 op 名 cast | +| 示例 autocast 用法 | [`gpt2/main.cc:91,263-270,472`](../example/gpt2/main.cc) | `--dtype` → 双参 `AutocastGuard` | diff --git a/docs/timeline.md b/docs/timeline.md new file mode 100644 index 000000000..33e85d3a3 --- /dev/null +++ b/docs/timeline.md @@ -0,0 +1,210 @@ +# nsys Timeline 详细分析 + +> 本文是 `kernel 优化.md` 第 2.2 节的详细数据：阶段与 kernel 家族聚合表，以及逐 kernel 的 timeline 展开（以 Layer 0 为例）。测量口径与基线说明见该文 §2.1。 + +按 host 发射归属（kernel 经 correlationId 关联到其发射时刻所在的 NVTX 子区间，避开异步偏移），每轮 step 分三个阶段： + +| 阶段 | kernel 数 | Σ GPU 时间 | host 窗口 | 特征 | +|---|---|---|---|---| +| Forward | 1343 | 23.50 ms（含 339 Cast） | 62.84 ms | host 窗口是 GPU 时间的 2.7 倍，但其中 27.55 ms（44%）是 3 次 >1 ms 的阻塞调用（含一次 16.30 ms 的 loss DtoH 读回），扣除后真实发射工作约 35 ms，所以这不是单纯的 **host-bound** | +| Backward | 2116 | 28.28 ms（含 322 Cast） | 34.91 ms | host 与 GPU 接近平衡 | +| Optimizer | 115 | 32.29 ms（FP32 Adam，第一大头） | 0.86 ms | 115 个大 kernel 异步发完即返回，GPU 工作拖到后续窗口才排空 | +| **合计** | **3574** | **84.07 ms** | **98.61 ms** | 再加 ZeroGrad/DataUpload/LossReadback 三个不发 kernel 的窗口共 101.89 ms ≈ Step_1 wall 102.08 ms，nsys 下 host 全程无空闲 | + +另有三个不发射 kernel 的子区间：ZeroGrad 0.46 ms、DataUpload 0.03 ms、LossReadback 2.79 ms（host 窗口）。六个窗口之和 101.89 ms ≈ Step_1 wall 102.08 ms，说明 nsys 下整轮 host 全程无空闲。但「无空闲」不等于「都在发射」：Forward 那 62.84 ms 里 Σ CUDA API 占 51.24 ms，其中 27.55 ms 集中在 3 次 >1 ms 的阻塞调用上，真正逐次发射 3477 个 API 只花 23.69 ms（见结论 6）。全轮 API 合计 71.29 ms 已被 CUPTI 放大，非 profiling 下会明显缩短，故「Forward host-bound」只是 nsys 下的强结论；跨阶段看，host 有相当比例的时间其实是在等 GPU。 + +按 kernel 家族聚合： + +| 家族 | kernel 数 | Σ (ms) | 占比 | +|---|---|---|---| +| Adam（optimizer，fp32） | 115 | 32.29 | 38.4% | +| 其它 elementwise / 结构 | 1749 | 20.87 | 24.8% | +| GEMM（BF16 Tensor Core） | 339 | 16.85 | 20.0% | +| Cast（autocast 新增） | 661 | 10.74 | 12.8% | +| Fill | 710 | 3.31 | 3.9% | +| **合计** | **3574** | **84.07** | 100% | + +其中 Cast 再拆方向：f32→bf16 下转 356 个 / 9.47 ms（权重下转为主，平均 26.6 μs），bf16→f32 上转 305 个 / 1.28 ms（激活上转，平均 4.2 μs）——下转个数只多 17% 但耗时是上转的 **7.4 倍**，这就是 3.1 优先消除权重下转的依据。 + +#### 2.2.1 Forward Timeline +Forward 共 **1343 kernels / 23.502 ms**（含 339 个 autocast Cast）。 + +##### 2.2.1.1 迭代入口（整轮各 1 次） + +``` +[ 0] +0.547ms 8.54us EmbeddingFwd·f32 token embedding: (4,64) → (4,64,2048) +[ 1] +0.597ms 5.21us SliceFwd·f32 freqs_cis 取本轮窗口的 cos/sin +[ 2] +0.654ms 2.78us TriuFwd·f32 causal mask = Triu(ones(64,64),1) +``` + +之后进入 **16 层结构相同的 TransformerLayer**（单层 = RMSNorm(ln1) → Attention → 残差 → RMSNorm(ln2) → MLP → 残差），最后接 ln_f → lm_head → CrossEntropy。 + +##### 2.2.1.2 实测：Layer 0 的 Attention 半区（kernels [3..66]，含 cast） + +`+ms` 为相对 Step_1 起点的 host 发射时刻，`us` 为 GPU 执行时间： + +``` +── RMSNorm ln1（6 kernels，全 fp32）──────────────────────────────── + [ 3] +0.679ms 8.03us UnaryFwd[Pow]·f32 x² + [ 4] +0.695ms 6.59us Mean(Reduce)·f32 mean(x², -1) + [ 5] +0.710ms 2.46us UnaryFwd[AddScalar]·f32 + eps + [ 6] +0.724ms 2.53us UnaryFwd[Rsqrt]·f32 1/√(·) + [ 7] +0.740ms 22.05us BinaryFwd[Mul]·f32 x * rsqrt + [ 8] +0.760ms 21.50us BinaryFwd[Mul]·f32 norm * weight + +── QKV 投影（2 Cast + 1 bf16 GEMM）───────────────────────────────── + [ 9] +0.779ms 4.80us Cast[f32->bf16] ln1 输出 → bf16（Linear 输入） + [10] +0.789ms 41.50us Cast[f32->bf16] QKV master 权重 → bf16 ★ + [11] +0.834ms 28.67us ampere_bf16_...f2f_..._tn Linear 2048→3072 (=(H+2·KV)·D) + +── 拆分 q/k/v + 取 RoPE cos/sin/even/odd（7 Slice）───────────────── + [12] +0.869ms 16.48us SliceFwd·bf16 q = qkv[..., :2048] + [13] +0.903ms 6.72us SliceFwd·bf16 k = qkv[..., 2048:2560] + [14] +0.932ms 6.75us SliceFwd·bf16 v = qkv[..., 2560:3072] + [15] +0.969ms 6.24us SliceFwd·f32 cos = freqs_cis[...,0] + [16] +1.002ms 5.60us SliceFwd·f32 sin = freqs_cis[...,1] + [17] +1.033ms 11.23us SliceFwd·bf16 q_even = q[..., 0::2] + [18] +1.063ms 11.10us SliceFwd·bf16 q_odd = q[..., 1::2] + +── RoPE 作用于 q（fp32 算术：每个 Mul 前先 Cast[bf16->f32] 上转）──── + [19] +1.077ms 3.52us Cast[bf16->f32] q_even ↑f32 + [20] +1.086ms 15.90us BinaryFwd[Mul]·f32 q_even * cos + [21] +1.107ms 3.39us Cast[bf16->f32] q_odd ↑f32 + [22] +1.116ms 15.81us BinaryFwd[Mul]·f32 q_odd * sin + [23] +1.131ms 4.13us BinaryFwdNB[Sub]·f32 left = q_even·cos − q_odd·sin + [24] +1.153ms 3.42us Cast[bf16->f32] q_even ↑f32 + [25] +1.162ms 15.87us BinaryFwd[Mul]·f32 q_even * sin + [26] +1.176ms 3.36us Cast[bf16->f32] q_odd ↑f32 + [27] +1.185ms 15.84us BinaryFwd[Mul]·f32 q_odd * cos + [28] +1.200ms 4.16us BinaryFwdNB[Add]·f32 right = q_even·sin + q_odd·cos + [29] +1.224ms 8.38us StackFwd·f32 stack(left,right) → flatten + +── RoPE 作用于 k（2 Slice + 4×[Cast+Mul] + Sub + Add + Stack）────── + [30] +1.261ms 6.01us SliceFwd·bf16 k_even + [31] +1.288ms 5.95us SliceFwd·bf16 k_odd + [32] +1.301ms 2.78us Cast[bf16->f32] + [33] +1.310ms 7.71us BinaryFwd[Mul]·f32 k_even * cos + [34] +1.324ms 2.78us Cast[bf16->f32] + [35] +1.333ms 7.68us BinaryFwd[Mul]·f32 k_odd * sin + [36] +1.345ms 3.07us BinaryFwdNB[Sub]·f32 k left + [37] +1.361ms 2.78us Cast[bf16->f32] + [38] +1.369ms 7.68us BinaryFwd[Mul]·f32 k_even * sin + [39] +1.383ms 2.78us Cast[bf16->f32] + [40] +1.392ms 7.68us BinaryFwd[Mul]·f32 k_odd * cos + [41] +1.406ms 2.98us BinaryFwdNB[Add]·f32 k right + [42] +1.429ms 4.51us StackFwd·f32 k stack + +── GQA：K/V 8 头复制到 32 头（2 RepeatInterleave）────────────────── + [43] +1.461ms 9.60us RepeatInterleaveFwd·f32 k: repeat_interleave(n_rep=4) + [44] +1.487ms 9.50us RepeatInterleaveFwd·bf16 v: repeat_interleave(n_rep=4) + +── 转置到 (B,H,T,D) 并算 Kᵀ（4×[Fill + Transpose]）───────────────── + [45] +1.516ms 4.19us Fill·f32 [46] 19.23us TransposeFwd·f32 q + [47] +1.543ms 4.13us Fill·f32 [48] 18.46us TransposeFwd·f32 k + [49] +1.568ms 4.16us Fill·bf16 [50] 18.08us TransposeFwd·bf16 v + [51] +1.595ms 4.13us Fill·f32 [52] 19.97us TransposeFwd·f32 kᵀ + +── ★ Attention 核心（全 bf16 Tensor Core / WMMA）─────────────────── + [53] +1.628ms 4.96us Cast[f32->bf16] q → bf16 + [54] +1.639ms 4.80us Cast[f32->bf16] kᵀ → bf16 + [55] +1.681ms 7.29us cutlass_wmma_bf16_nn SCORE att = q · kᵀ (B·H=128 批) + [56] +1.697ms 4.96us UnaryFwd[MulScalar]·bf16 att *= 1/√D (=0.125) + [57] +1.719ms 2.82us Cast[f32->bf16] mask → bf16 + [58] +1.729ms 6.34us MaskFwd·bf16 masked_fill(causal, -inf) + [59] +1.749ms 26.17us SoftmaxFwd·bf16 softmax(att, -1) + [60] +1.773ms 7.17us cutlass_wmma_bf16_nn y = att · v + +── 输出整理 + 投影（[Fill+Transpose] + Cast + bf16 GEMM + Cast）───── + [61] +1.794ms 4.10us Fill·bf16 [62] 18.34us TransposeFwd·bf16 y→(B,T,H,D) + [63] +1.824ms 28.48us Cast[f32->bf16] out_proj master 权重 → bf16 ★ + [64] +1.844ms 23.93us ampere_bf16_...f2f_..._tn OutProj Linear 2048→2048 + [65] +1.872ms 4.99us Cast[bf16->f32] attn 输出 → fp32 + [66] +1.882ms 6.33us BinaryFwdNB[Add]·f32 x = x + attn_out（残差，fp32） +``` + +##### 2.2.1.3 实测：Layer 0 的 MLP(SwiGLU) 半区（kernels [67..85]） + +``` +── RMSNorm ln2（6 kernels，全 fp32）[67..72] ─────────────────────── + [67] Pow·f32 [68] Mean·f32 [69] AddScalar·f32 [70] Rsqrt·f32 [71] Mul·f32 [72] Mul·f32 + +── MLP SwiGLU（每个 Linear = 2 Cast + 1 bf16 GEMM）───────────────── + [73] +1.997ms 4.80us Cast[f32->bf16] c_fc 输入 → bf16 + [74] +2.008ms 102.29us Cast[f32->bf16] c_fc 权重(2048×8192) → bf16 ★★ + [75] +2.025ms 73.24us ampere_bf16_...f2f_..._tn c_fc Linear 2048→8192 (x1) + [76] +2.044ms 6.59us Cast[f32->bf16] c_fc2 输入 → bf16 + [77] +2.054ms 100.57us Cast[f32->bf16] c_fc2 权重(2048×8192) → bf16 ★★ + [78] +2.068ms 73.14us ampere_bf16_...f2f_..._tn c_fc2 Linear 2048→8192 (x2) + [79] +2.083ms 14.21us UnaryFwd[Sigmoid]·bf16 σ(x2) + [80] +2.098ms 12.57us BinaryFwdNB[Mul]·bf16 SiLU = x2 * σ(x2) + [81] +2.112ms 14.72us BinaryFwdNB[Mul]·bf16 x3 = x1 * SiLU(x2)（门控） + [82] +2.125ms 102.58us Cast[f32->bf16] c_proj 权重(8192×2048) → bf16 ★★ + [83] +2.150ms 90.78us ampere_bf16_...f2f_..._tn c_proj Linear 8192→2048 + [84] +2.171ms 5.34us Cast[bf16->f32] mlp 输出 → fp32 + [85] +2.181ms 6.72us BinaryFwdNB[Add]·f32 x = x + mlp_out（残差，fp32） +``` + +> **★★ 首要发现：权重 cast 比 bf16 GEMM 本身还贵。** + +##### 2.2.1.4 迭代收尾（整轮各 1 次） + +``` +[1331..1336] RMSNorm ln_f（6 kernels·f32） Pow 7.87 / Mean 6.59 / AddScalar 2.40 / Rsqrt 2.50 / Mul 21.76 / Mul 21.60 us —— 第 33 个 RMSNorm，+46.852ms 起 +[1337] +46.927ms 4.93us Cast f32→bf16 lm_head 输入 → bf16 +[1338] +46.937ms 1533.17us Cast f32→bf16 lm_head 权重 2048×128256 → bf16 + ★★ 全步最贵的单个 kernel（不只是最贵 Cast），是其 GEMM 698.65us 的 2.19 倍 +[1339] +46.955ms 698.65us ampere_bf16_...f2f_..._tn（128x256，grid 覆盖 vocab=128256） lm_head Linear 2048→128256 +[1340] +46.998ms 197.04us Cast bf16→f32：logits 上转（autocast 让 CE 走 fp32） +[1341] +47.010ms 268.13us CrossEntropyFwd·f32 损失 +[1342] +63.353ms 2.65us UnaryFwd MulScalar·f32 loss / grad_accum_steps 为 1 + +注：[1341] 与 [1342] 之间空了 16.34 ms——正是结论 6 里那次 1024 B loss DtoH 读回造成的 host 阻塞。 +Forward 的 GPU 工作在 +47.0 ms 就已全部发出，host 却要等到 +63.3 ms 才能发出最后一个标量 kernel。 +``` + +#### 2.2.2 Backward Timeline +Backward 是 Forward 的逆序 autograd，共 **2116 kernels / 28.276 ms**（host 发射）。结构： + +``` +CrossEntropyBwd·f32 损失反向（700 us，autocast 保持 fp32） +lm_head 反向 ampere_s16816gemm_bf16_128x256_nn（dW，单 kernel 686 us）+ dX +── 每层逆序（ln_f→…→layer0）────────────────────────────────────── + 残差 Add 反向 BinaryBwdNB[Add]·f32 + MLP 反向 c_proj/c_fc/c_fc2 各 2 bf16 GEMM（dX+dW）+ 权重 Cast[f32->bf16] + 门控：BinaryBwdNBVec[Mul]·bf16、UnaryBwd[Sigmoid]·bf16 + RMSNorm ln2 反向 UnaryBwd[Pow/Rsqrt/AddScalar]·f32、GenericReduceBwd·f32、BinaryBwd[Mul]·f32 + 残差 Add 反向 + Attention 反向 out_proj 2 GEMM；att·V 反向 cutlass_wmma_bf16_{nt,tn}； + SoftmaxBwd·bf16、MaskBwd、MulScalar 反向；Q·Kᵀ 反向 cutlass_wmma_bf16 + RoPE 反向 StackBwd/SliceBwd/BinaryBwd[Mul]/Sub/Add；RepeatInterleaveBwd + qkv 2 bf16 GEMM + Cast + RMSNorm ln1 反向 + ★ AccumulateGrad 每个产生梯度的张量：grad += rate·g（209 f32 + 16 bf16） +EmbeddingBwd·f32 词嵌入反向（scatter，1.313 ms） +``` + +Backward kernel Top（Σ GPU 时间，host 发射归属）： + +| Kernel | n | Σ | 角色 | +|---|---|---|---| +| `ampere_s16816gemm_bf16_128x128_..._32x5_nt` | 49 | 4.635 ms | Linear 反向 dX/dW（bf16 Tensor Core） | +| `BinaryBwd[Mul]·f32` | 194 | 3.674 ms | RoPE/RMSNorm 乘法反向（fp32） | +| `Fill·f32` | 630 | 2.982 ms | 梯度/输出缓冲清零 | +| `ampere_s16816gemm_bf16_128x128_..._64x3_nn` | 32 | 2.137 ms | Linear 反向 | +| `TransposeFwd·f32` | 80 | 1.551 ms | 反向中的转置 | +| `Cast[f32->bf16]` | 178 | 1.492 ms | 反向 Linear 输入/权重下转 | +| `EmbeddingBwd·f32` | 1 | 1.313 ms | 词嵌入 scatter 反向 | +| `ampere_s16816gemm_bf16_256x128_..._{nt,nn}` | 48 | 2.428 ms | Linear 反向 | +| `AccumulateGrad`（f32+bf16） | 225 | 1.267 ms | autograd 梯度累加 | +| `cutlass_wmma_bf16_..._{nt,tn}` | 64 | 0.506 ms | Attention batched 反向 | + +> `Fill` 在反向仍高达 **630 次**（2.982 ms）：每个需累加梯度的张量都要先清零 `.grad`，是「小 kernel 过多、launch 开销大」的主要来源。bf16 计算变快后，这类 launch/访存开销占比反而更突出。 + +#### 2.2.3 Optimizer Timeline + +``` +AdamAccumulateGrad·f32 ×115 Σ=32.292 ms avg=280.80 us +``` + +- `AdamAccumulateGradKernel` 是**融合的 Adam 更新**：单 kernel 内完成 `m=β1·m+(1-β1)g`、`v=β2·v+(1-β2)g²`、bias-correction、`param -= lr·m̂/(√v̂+eps)`（`accumulate_grad.cu`）。 +- **115 ≈ 参数张量数**：每层 7 个（ln1.w、ln2.w、qkv.w、proj.w、c_fc.w、c_fc2.w、c_proj.w）×16 = 112，加 embedding、ln_f、lm_head。 diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 7697ad213..5d05a5af3 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -36,6 +36,7 @@ #include "infini_train/include/nn/parallel/reduce_op_type.h" #include "infini_train/include/nn/parallel/tensor_parallel.h" #include "infini_train/include/optimizer.h" +#include "infini_train/include/utils/nvtx.h" #ifdef PROFILE_MODE #include "infini_train/include/profiler.h" #endif @@ -468,12 +469,20 @@ void Train(const nn::parallel::Rank &rank) { #ifdef PROFILE_MODE Profiler::Instance().SetTag("Step_" + std::to_string(step)); #endif +#ifdef NVTX_MODE + // Spans one whole training step. A scoped object is used here rather than + // a push/pop pair because the loop body has several exits. + const auto nvtx_step_name = "Step_" + std::to_string(step); + infini_train::utils::NvtxRange nvtx_step(nvtx_step_name.c_str()); +#endif const float current_lr = scheduler ? scheduler->learning_rate() : static_cast(FLAGS_learning_rate); float lossf = 0.0f; // model->Train(); if (pp_world_size == 1) { + INFINI_TRAIN_NVTX_PUSH("ZeroGrad"); optimizer->ZeroGrad(); + INFINI_TRAIN_NVTX_POP(); // if we are trying to overfit a single batch, we reset the loader here if (FLAGS_overfit_single_batch) { @@ -486,17 +495,20 @@ void Train(const nn::parallel::Rank &rank) { // (bs, seq_len), (bs, seq_len) auto [x, y] = next_train_batch(); + INFINI_TRAIN_NVTX_PUSH("DataUpload"); x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); + INFINI_TRAIN_NVTX_POP(); LOG(INFO) << "Rank " << rank.GlobalRank() << ": start forward"; - + INFINI_TRAIN_NVTX_PUSH("Forward"); // (bs, seq_len, vocab_size) auto logits = (*model)({x, y})[0]; LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish model forward, start loss forward"; auto loss = (*loss_fn)({logits, y})[0]; // FIXME(jym): verify gradient accumulation precision loss = loss / grad_accum_steps; + INFINI_TRAIN_NVTX_POP(); // disable autocast for the current step (backward is not under autocast) autocast_guard.Disable(); @@ -508,24 +520,32 @@ void Train(const nn::parallel::Rank &rank) { if (ddp_world_size > 1 && micro_step != grad_accum_steps - 1) { no_sync_guard = model->no_sync(); } + INFINI_TRAIN_NVTX_PUSH("Backward"); loss->Backward(); + INFINI_TRAIN_NVTX_POP(); // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA // between forward and backward. + INFINI_TRAIN_NVTX_PUSH("LossReadback"); auto loss_cpu = loss->To(Device()); lossf += static_cast(loss_cpu.DataPtr())[0]; + INFINI_TRAIN_NVTX_POP(); LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } + INFINI_TRAIN_NVTX_PUSH("Optimizer"); optimizer->Step(); if (scheduler) { scheduler->Step(); } + INFINI_TRAIN_NVTX_POP(); } else { auto [x, y] = next_train_batch(); x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); + INFINI_TRAIN_NVTX_PUSH("TrainStep"); lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype); + INFINI_TRAIN_NVTX_POP(); if (scheduler) { scheduler->Step(); } diff --git a/example/llama3/main.cc b/example/llama3/main.cc index a9f405fa5..142370fb4 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -38,6 +38,7 @@ #include "infini_train/include/nn/parallel/utils.h" #include "infini_train/include/optimizer.h" #include "infini_train/include/utils/global_module_hook_registry.h" +#include "infini_train/include/utils/nvtx.h" #include "infini_train/include/utils/precision_check_config.h" #include "infini_train/include/utils/precision_checker.h" #ifdef PROFILE_MODE @@ -380,6 +381,13 @@ void Train(const nn::parallel::Rank &rank) { start_step = resume_result.global_step; size_t consumed_train_samples = resume_result.consumed_train_samples; + // enable shadow weights (polymorphic; DistributedOptimizer degrades to a no-op with a warning). + // Note: this point is after ResumeFromCheckpoint, so the shadows are built directly from the + // already-restored FP32 params; no extra refresh is needed. + if (FLAGS_dtype == kDtypeBF16) { + optimizer->EnableShadowWeights(DataType::kBFLOAT16); + } + auto advance_train_iter = [&]() { ++train_iter; if (train_iter == train_loader.end()) { @@ -456,12 +464,20 @@ void Train(const nn::parallel::Rank &rank) { #ifdef PROFILE_MODE Profiler::Instance().SetTag("Step_" + std::to_string(step)); #endif +#ifdef NVTX_MODE + // Spans one whole training step. A scoped object is used here rather than + // a push/pop pair because the loop body has several exits. + const auto nvtx_step_name = "Step_" + std::to_string(step); + infini_train::utils::NvtxRange nvtx_step(nvtx_step_name.c_str()); +#endif const float current_lr = scheduler ? scheduler->learning_rate() : static_cast(FLAGS_learning_rate); float lossf = 0.0f; if (pp_world_size == 1) { // model->Train(); + INFINI_TRAIN_NVTX_PUSH("ZeroGrad"); optimizer->ZeroGrad(); + INFINI_TRAIN_NVTX_POP(); // if we are trying to overfit a single batch, we reset the loader here if (FLAGS_overfit_single_batch) { @@ -474,16 +490,20 @@ void Train(const nn::parallel::Rank &rank) { // (bs, seq_len), (bs, seq_len) auto [x, y] = next_train_batch(); + INFINI_TRAIN_NVTX_PUSH("DataUpload"); x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); + INFINI_TRAIN_NVTX_POP(); LOG(INFO) << "Rank " << rank.GlobalRank() << ": start forward"; + INFINI_TRAIN_NVTX_PUSH("Forward"); // (bs, seq_len, vocab_size) auto logits = (*model)({x, y})[0]; LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish model forward, start loss forward"; auto loss = (*loss_fn)({logits, y})[0]; // FIXME(jym): verify gradient accumulation precision loss = loss / grad_accum_steps; + INFINI_TRAIN_NVTX_POP(); // disable autocast for the current step (backward is not under autocast) autocast_guard.Disable(); @@ -495,24 +515,32 @@ void Train(const nn::parallel::Rank &rank) { if (ddp_world_size > 1 && micro_step != grad_accum_steps - 1) { no_sync_guard = model->no_sync(); } + INFINI_TRAIN_NVTX_PUSH("Backward"); loss->Backward(); + INFINI_TRAIN_NVTX_POP(); // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA // between forward and backward. + INFINI_TRAIN_NVTX_PUSH("LossReadback"); auto loss_cpu = loss->To(Device()); lossf += static_cast(loss_cpu.DataPtr())[0]; + INFINI_TRAIN_NVTX_POP(); LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward"; } + INFINI_TRAIN_NVTX_PUSH("Optimizer"); optimizer->Step(); if (scheduler) { scheduler->Step(); } + INFINI_TRAIN_NVTX_POP(); } else { auto [x, y] = next_train_batch(); x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); + INFINI_TRAIN_NVTX_PUSH("TrainStep"); lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype); + INFINI_TRAIN_NVTX_POP(); if (scheduler) { scheduler->Step(); } diff --git a/example/mnist/main.cc b/example/mnist/main.cc index 7744e0947..1e12bd1d3 100644 --- a/example/mnist/main.cc +++ b/example/mnist/main.cc @@ -47,14 +47,14 @@ int main(int argc, char *argv[]) { auto test_dataset = std::make_shared(FLAGS_dataset, false); DataLoader test_dataloader(test_dataset, FLAGS_bs); - auto network = MNIST(); + auto network = std::make_shared(); Device device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); Device cpu_device = Device(); - network.To(device); + network->To(device); - auto loss_fn = nn::CrossEntropyLoss(); - loss_fn.To(device); - auto optimizer = optimizers::SGD(network.Parameters(), FLAGS_lr); + auto loss_fn = std::make_shared(); + loss_fn->To(device); + auto optimizer = optimizers::SGD(network->Parameters(), FLAGS_lr); for (int epoch = 0; epoch < FLAGS_num_epoch; ++epoch) { int train_idx = 0; @@ -66,10 +66,10 @@ int main(int argc, char *argv[]) { auto new_image = std::make_shared(image->To(device)); auto new_label = std::make_shared(label->To(device)); - auto outputs = network.Forward({new_image}); + auto outputs = network->Forward({new_image}); optimizer.ZeroGrad(); - auto loss = loss_fn.Forward({outputs[0], new_label}); + auto loss = loss_fn->Forward({outputs[0], new_label}); loss[0]->Backward(); // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA @@ -104,9 +104,9 @@ int main(int argc, char *argv[]) { auto new_label = std::make_shared(label->To(device)); auto label_cpu = label->To(cpu_device); - auto outputs = network.Forward({new_image}); + auto outputs = network->Forward({new_image}); auto output_cpu = outputs[0]->To(cpu_device); - auto loss = loss_fn.Forward({outputs[0], new_label}); + auto loss = loss_fn->Forward({outputs[0], new_label}); auto loss_cpu = loss[0]->To(cpu_device); const int batch_size = output_cpu.Dims()[0]; diff --git a/infini_train/include/autocast.h b/infini_train/include/autocast.h index 201a0347b..d4e063d0c 100644 --- a/infini_train/include/autocast.h +++ b/infini_train/include/autocast.h @@ -11,6 +11,7 @@ #include "infini_train/include/tensor.h" namespace infini_train { +std::shared_ptr GetShadow(const Tensor *param); namespace { inline std::string_view GetBaseOpName(std::string_view op) { constexpr std::string_view function_suffix = "Function"; @@ -37,10 +38,10 @@ enum class CastPolicy : uint8_t { // Cast-policy maps and their associated operations. The op names should match the ones defined in the op registry. inline constexpr std::array kLowerPrecisionOps = {"Matmul", "Linear"}; inline constexpr std::array kFP32Ops - = {"Sin", "Cos", "Tan", "Asin", "Acos", "Atan", "Sinh", - "Cosh", "Tanh", "Asinh", "Acosh", "Atanh", "Exp", "Log", - "Sqrt", "Reciprocal", "Rsqrt", "Prod", "Pow", "CrossEntropy", "VocabParallelCrossEntropy", - "Layernorm"}; + = {"Sin", "Cos", "Tan", "Asin", "Acos", "Atan", "Sinh", + "Cosh", "Tanh", "Asinh", "Acosh", "Atanh", "Exp", "Log", + "Sqrt", "Reciprocal", "Rsqrt", "Prod", "Pow", "CrossEntropy", "VocabParallelCrossEntropy", + "LayerNorm", "RMSNorm"}; // Mapping from operation names to their cast policies. This is the primary construct that is used in autocasting. The // op names should match the ones defined in the op registry. @@ -69,7 +70,8 @@ inline const std::unordered_map kOpCastPolicyMap = {"Pow", CastPolicy::kFP32}, {"CrossEntropy", CastPolicy::kFP32}, {"VocabParallelCrossEntropy", CastPolicy::kFP32}, - {"Layernorm", CastPolicy::kFP32}, + {"LayerNorm", CastPolicy::kFP32}, + {"RMSNorm", CastPolicy::kFP32}, }; inline DataType GetDefaultAutocastDtype(Device::DeviceType device_type) { @@ -131,7 +133,12 @@ struct AutocastContext { if (is_floating_point(current_dtype)) { DataType target_dtype = get_target_dtype(); if (current_dtype != target_dtype) { - arg = std::make_shared(arg->To(target_dtype)); + auto shadow = GetShadow(arg.get()); + if (shadow && shadow->Dtype() == target_dtype) { + arg = shadow; + } else { + arg = std::make_shared(arg->To(target_dtype)); + } } } } diff --git a/infini_train/include/autograd/normalization.h b/infini_train/include/autograd/normalization.h index c4148cbaa..62fce3e68 100644 --- a/infini_train/include/autograd/normalization.h +++ b/infini_train/include/autograd/normalization.h @@ -21,6 +21,21 @@ class LayerNorm : public Function { const std::vector> &output_tensors) override; std::vector> Backward(const std::vector> &grad_outputs) override; +private: + const float eps_ = 1e-5f; +}; + +class RMSNorm : public Function { +public: + static constexpr char kType[] = "RMSNormFunction"; + + explicit RMSNorm(float eps) : Function(kType), eps_(eps) {} + + std::vector> Forward(const std::vector> &input_tensors) override; + void SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) override; + std::vector> Backward(const std::vector> &grad_outputs) override; + private: const float eps_ = 1e-5f; }; diff --git a/infini_train/include/autograd/sparse.h b/infini_train/include/autograd/sparse.h index ef1cc21bc..303b7cfbc 100644 --- a/infini_train/include/autograd/sparse.h +++ b/infini_train/include/autograd/sparse.h @@ -20,8 +20,5 @@ class Embedding : public Function { void SetupContext(const std::vector> &input_tensors, const std::vector> &output_tensors) override; std::vector> Backward(const std::vector> &grad_outputs) override; - -private: - std::vector weight_dims_; }; } // namespace infini_train::autograd diff --git a/infini_train/include/common/cuda/common_cuda.h b/infini_train/include/common/cuda/common_cuda.h index 862f41820..1fa71e0d0 100644 --- a/infini_train/include/common/cuda/common_cuda.h +++ b/infini_train/include/common/cuda/common_cuda.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -28,6 +29,17 @@ namespace infini_train::common::cuda { } \ } while (0) +// cuBLASLt shares cublasStatus_t with cuBLAS, so the same stringifier applies; only the +// log prefix differs, to tell the two libraries apart when a call fails. +#define CUBLASLT_CHECK(call) \ + do { \ + cublasStatus_t status = call; \ + if (status != CUBLAS_STATUS_SUCCESS) { \ + LOG(FATAL) << "CUBLASLT Error: " << cublasGetStatusString(status) << " at " << __FILE__ << ":" \ + << __LINE__; \ + } \ + } while (0) + #define CUDA_DRIVER_CHECK(call) \ do { \ CUresult status = call; \ diff --git a/infini_train/include/common/cuda/kernel_helper.cuh b/infini_train/include/common/cuda/kernel_helper.cuh index 6b532afb2..6852e275f 100644 --- a/infini_train/include/common/cuda/kernel_helper.cuh +++ b/infini_train/include/common/cuda/kernel_helper.cuh @@ -109,8 +109,10 @@ template __device__ __forceinline__ T Cos(const T &x) { } template __device__ __forceinline__ T Tanh(const T &x) { - if constexpr (std::is_same_v || std::is_same_v) { - return htanh(x); + if constexpr (std::is_same_v) { + return __float2half(tanhf(__half2float(x))); + } else if constexpr (std::is_same_v) { + return __float2bfloat16(tanhf(__bfloat162float(x))); } else if constexpr (std::is_same_v) { return tanhf(x); } else { diff --git a/infini_train/include/dispatcher.h b/infini_train/include/dispatcher.h index 6aa7e6e01..8f6c672a5 100644 --- a/infini_train/include/dispatcher.h +++ b/infini_train/include/dispatcher.h @@ -9,6 +9,7 @@ #include "infini_train/include/common/common.h" #include "infini_train/include/device.h" +#include "infini_train/include/utils/nvtx.h" #ifdef PROFILE_MODE #include "infini_train/include/profiler.h" #endif @@ -17,10 +18,18 @@ namespace infini_train { class KernelFunction { public: - template explicit KernelFunction(FuncT &&func) : func_ptr_(reinterpret_cast(func)) {} + // The operator name is stored per instance rather than read from a + // thread_local context: SetProfileContext() only records the outermost + // dispatch, so nested dispatches would overwrite each other's name. + template + explicit KernelFunction(std::string name, FuncT &&func) + : name_(std::move(name)), func_ptr_(reinterpret_cast(func)) {} // TODO(dcj): support auto-deduction of return type and parameter types template RetT Call(ArgsT... args) const { +#ifdef NVTX_MODE + utils::NvtxRange nvtx_range(name_.c_str()); +#endif #ifdef PROFILE_MODE const auto &ctx = GetProfileContext(); Profiler::Instance().StartRecord(ctx.name, ctx.device); @@ -45,6 +54,7 @@ class KernelFunction { } private: + std::string name_; void *func_ptr_ = nullptr; }; @@ -71,7 +81,7 @@ class Dispatcher { template void Register(const KeyT &key, FuncT &&kernel) { CHECK(!key_to_kernel_map_.contains(key)) << "Kernel already registered: " << key.second << " on device: " << static_cast(key.first); - key_to_kernel_map_.emplace(key, kernel); + key_to_kernel_map_.try_emplace(key, key.second, kernel); } template RetT Call(KeyT key, ArgsT... args) const { diff --git a/infini_train/include/nn/modules/normalization.h b/infini_train/include/nn/modules/normalization.h index 9cd5be886..61794af0f 100644 --- a/infini_train/include/nn/modules/normalization.h +++ b/infini_train/include/nn/modules/normalization.h @@ -36,6 +36,6 @@ class RMSNorm : public CloneableModule { std::vector> Forward(const std::vector> &x) override; private: - float eps_ = 1e-5f; + float eps_ = 1e-6f; }; } // namespace infini_train::nn diff --git a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h index d7cea198e..975289b07 100644 --- a/infini_train/include/nn/parallel/ddp/distributed_optimizer.h +++ b/infini_train/include/nn/parallel/ddp/distributed_optimizer.h @@ -46,6 +46,10 @@ class DistributedOptimizer final : public infini_train::Optimizer { virtual void set_learning_rate(float lr) override; virtual float learning_rate() const override; + // Shadow weights are not supported in the distributed setting: base_optimizer_ manages shard params, + // while autocast sees the full param in forward, so the registry never hits. Degrades to a no-op. + void EnableShadowWeights(DataType shadow_dtype) override; + private: using AddShardParam = std::function &, const std::shared_ptr &)>; diff --git a/infini_train/include/optimizer.h b/infini_train/include/optimizer.h index d85b1acea..cf11a1c01 100644 --- a/infini_train/include/optimizer.h +++ b/infini_train/include/optimizer.h @@ -8,6 +8,8 @@ #include #include +#include "infini_train/include/datatype.h" + namespace infini_train { class Tensor; } @@ -25,6 +27,8 @@ class Optimizer { Optimizer(const NamedParameterList &named_params, float learning_rate); + virtual ~Optimizer() = default; + virtual void ZeroGrad(bool set_to_none = true); virtual void Step() = 0; @@ -43,6 +47,16 @@ class Optimizer { void set_initial_learning_rate(float lr); + // ========== Shadow weights API (default no-op; Adam overrides) ========== + // Enable shadow weights: create a low-precision copy per param and register it in the global registry + // so autocast can reuse it directly, skipping the repeated CastKernel in forward. + // No-op by default; only Adam implements it. + virtual void EnableShadowWeights(DataType shadow_dtype = DataType::kBFLOAT16) {} + virtual void DisableShadowWeights() {} + // Re-sync all shadow weights from the current FP32 master weights (called after a checkpoint restore). + virtual void RefreshShadowWeights() {} + virtual bool ShadowWeightsEnabled() const { return false; } + protected: std::vector> params_; std::vector parameter_names_; @@ -80,6 +94,19 @@ class Adam : public Optimizer { static OptimizerCreatorNamed CreateNamed(float learning_rate = 1e-3, float beta1 = 0.9, float beta2 = 0.999, float eps = 1e-8); + // add shadow weights API + void EnableShadowWeights(DataType shadow_dtype) override; + + void DisableShadowWeights() override; + + void RefreshShadowWeights() override; + + bool ShadowWeightsEnabled() const override { return shadow_enable_; } + + std::shared_ptr GetShadow(const Tensor *param) const; + + ~Adam() override { DisableShadowWeights(); } + private: int64_t t_; const float beta1_; @@ -87,6 +114,10 @@ class Adam : public Optimizer { const float eps_; std::vector> m_; std::vector> v_; + // add shadow weights variable + bool shadow_enable_ = false; + DataType shadow_dtype_ = DataType::kBFLOAT16; + std::vector> shadow_weights_; }; } // namespace optimizers } // namespace infini_train diff --git a/infini_train/include/sparse_row_grad.h b/infini_train/include/sparse_row_grad.h new file mode 100644 index 000000000..57ab353a2 --- /dev/null +++ b/infini_train/include/sparse_row_grad.h @@ -0,0 +1,98 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train { + +// Bookkeeping shared by the sparse embedding backward and the optimizer, so that a vocab-sized +// gradient buffer can be kept alive across steps and only the rows actually hit by the batch need +// to be scattered into, cleared and updated. +// +// Invariant: grad_buffer is all-zero except for the rows listed in row_list[0..*count), which were +// claimed (deduplicated against stamp) since the last clear. generation is bumped at every clear +// so the stamp array can tag a row as "already claimed this cycle" without ever being reset. +struct SparseRowGradState { + const Tensor *weight = nullptr; + std::weak_ptr weight_guard; // detects a recycled pointer after the weight was freed + int64_t vocab = 0; + int64_t dim = 0; + int64_t capacity = 0; // row_list capacity; == vocab, so dedup alone guarantees no overflow + int32_t generation = 1; // stamp value identifying the current accumulation cycle + bool initialized = false; // one-time device zero-fill of grad_buffer/stamp/count is done + bool poisoned = false; // a dense (non-aliased) accumulation landed in a foreign storage + std::shared_ptr grad_buffer; // dense [vocab, dim], persistent storage for the grad + std::shared_ptr stamp; // kINT32 [vocab], per-row claim tags + std::shared_ptr row_list; // kINT32 [capacity], deduplicated rows since the last clear + std::shared_ptr count; // kINT32 [1], number of valid entries in row_list +}; + +class SparseRowGradRegistry { +public: + static SparseRowGradRegistry &Instance() { + static thread_local SparseRowGradRegistry instance; + return instance; + } + + SparseRowGradState *Lookup(const Tensor *weight) { + auto it = states_.find(weight); + if (it == states_.end()) { + return nullptr; + } + if (it->second->weight_guard.lock().get() != weight) { + states_.erase(it); // the address was recycled by a different tensor + return nullptr; + } + return it->second.get(); + } + + SparseRowGradState *GetOrCreate(const std::shared_ptr &weight) { + if (auto *state = Lookup(weight.get())) { + return state; + } + auto state = std::make_unique(); + state->weight = weight.get(); + state->weight_guard = weight; + state->vocab = weight->Dims()[0]; + state->dim = weight->Dims()[1]; + state->capacity = state->vocab; + const auto device = weight->GetDevice(); + state->grad_buffer = std::make_shared(weight->Dims(), weight->Dtype(), device); + state->stamp = std::make_shared(std::vector{state->vocab}, DataType::kINT32, device); + state->row_list = std::make_shared(std::vector{state->capacity}, DataType::kINT32, device); + state->count = std::make_shared(std::vector{1}, DataType::kINT32, device); + auto *raw = state.get(); + states_[weight.get()] = std::move(state); + return raw; + } + + void Erase(const Tensor *weight) { states_.erase(weight); } + +private: + SparseRowGradRegistry() = default; + std::unordered_map> states_; +}; + +// Free-function shim so call sites do not have to spell out the singleton. +inline SparseRowGradState *GetSparseRowGradState(const Tensor *weight) { + return SparseRowGradRegistry::Instance().Lookup(weight); +} + +// Zero the rows made dirty since the last clear, reset the row count and start a new claim +// generation. The caller must hold the device guard for the state's device. After this call the +// buffer is (again) all-zero, so nothing else needs to be touched to "zero the grad". +inline void ClearSparseRowGradRows(SparseRowGradState *state) { + const auto device = state->grad_buffer->GetDevice(); + auto clear_kernel = Dispatcher::Instance().GetKernel({device.type(), "SparseRowClearRows"}); + clear_kernel.Call(state->grad_buffer, state->row_list, state->count); + auto reset_kernel = Dispatcher::Instance().GetKernel({device.type(), "SparseRowResetCount"}); + reset_kernel.Call(state->count); + ++state->generation; +} + +} // namespace infini_train diff --git a/infini_train/include/utils/nvtx.h b/infini_train/include/utils/nvtx.h new file mode 100644 index 000000000..6d8de3cb5 --- /dev/null +++ b/infini_train/include/utils/nvtx.h @@ -0,0 +1,49 @@ +#pragma once + +// NVTX support for nsys timeline analysis. +// +// NVTX records host-side timestamps only and issues no CUDA call. This is what +// distinguishes it from PROFILE_MODE: the profiler wraps every dispatched +// kernel in SynchronizeStream + EventRecord/EventSynchronize, which serializes +// the pipeline and both inflates per-operator device time and hides every +// asynchronous-overlap and host-drain effect. NVTX leaves the execution shape +// intact, so ranges can be read directly against the CUDA kernel track. +// +// Everything below compiles away to nothing unless -DNVTX_MODE=1 is passed, +// which keeps CPU-only builds free of any CUDA include-path dependency. Callers +// should therefore use INFINI_TRAIN_NVTX_PUSH/POP unconditionally. + +#ifdef NVTX_MODE + +// The nvtx3 C API is header-only; nvtxRangePushA/Pop resolve through a runtime +// dlopen of the injection library, so no nvToolsExt library needs linking +// (libnvToolsExt.so was removed in CUDA 13). +#include + +namespace infini_train { +namespace utils { + +// RAII wrapper around the NVTX push/pop stack, for ranges whose extent matches +// a C++ scope. Preferred over the macros below: it keeps the push/pop balanced +// even if the enclosed code throws. +class NvtxRange { +public: + explicit NvtxRange(const char *name) { nvtxRangePushA(name); } + ~NvtxRange() { nvtxRangePop(); } + + NvtxRange(const NvtxRange &) = delete; + NvtxRange &operator=(const NvtxRange &) = delete; +}; + +} // namespace utils +} // namespace infini_train + +#define INFINI_TRAIN_NVTX_PUSH(name) ::nvtxRangePushA(name) +#define INFINI_TRAIN_NVTX_POP() ::nvtxRangePop() + +#else + +#define INFINI_TRAIN_NVTX_PUSH(name) ((void)0) +#define INFINI_TRAIN_NVTX_POP() ((void)0) + +#endif diff --git a/infini_train/src/autograd/accumulate.cc b/infini_train/src/autograd/accumulate.cc index 0c34819f6..a62d278e2 100644 --- a/infini_train/src/autograd/accumulate.cc +++ b/infini_train/src/autograd/accumulate.cc @@ -5,6 +5,7 @@ #include "infini_train/include/autograd/function_hook.h" #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/dispatcher.h" +#include "infini_train/include/sparse_row_grad.h" #include "infini_train/include/tensor.h" namespace infini_train::autograd { @@ -43,15 +44,33 @@ AccumulateGrad::Backward(const std::vector> &grad_output } auto grad = tensor_->grad(); + auto *sparse_state = GetSparseRowGradState(tensor_.get()); if (grad) { - if (overwrite) { - // If the tensor is marked to overrite its current grad on next grad update - // See notes in `infini_train::nn::parallel::Reducer::PrepareForBackward()` - // NOTE(zbl): must copy, cannot change grad buffer address - grad->CopyFrom(grad_output); + if (sparse_state && grad->DataPtr() == grad_output->DataPtr()) { + // Sparse embedding path: grad_output is a view of the persistent buffer the + // embedding backward already scatter-added into (learning_rate_ is 1.0f for + // AccumulateGrad nodes), so accumulating it again would double count. } else { - auto kernel = Dispatcher::Instance().GetKernel({device.type(), "AccumulateGrad"}); - kernel.Call(grad_output, learning_rate_, grad); + if (overwrite) { + // If the tensor is marked to overrite its current grad on next grad update + // See notes in `infini_train::nn::parallel::Reducer::PrepareForBackward()` + // NOTE(zbl): must copy, cannot change grad buffer address + grad->CopyFrom(grad_output); + } else { + auto kernel = Dispatcher::Instance().GetKernel({device.type(), "AccumulateGrad"}); + kernel.Call(grad_output, learning_rate_, grad); + } + if (sparse_state) { + // A dense grad from another producer (matmul on a tied weight, DDP bucket) + // landed in the storage; from here on the sparse optimizer shortcut is off. + sparse_state->poisoned = true; + if (grad->DataPtr() != sparse_state->grad_buffer->DataPtr()) { + // The accumulator belongs to that other producer: grad_output was the + // cumulative sparse buffer and the add/copy above merged it in full. + // Flush it so the next micro-batch contributes only its own delta. + ClearSparseRowGradRows(sparse_state); + } + } } } else { // FIXME(zbl): check whether need to do copying instead of slicing diff --git a/infini_train/src/autograd/normalization.cc b/infini_train/src/autograd/normalization.cc index eca830b06..f433ca452 100644 --- a/infini_train/src/autograd/normalization.cc +++ b/infini_train/src/autograd/normalization.cc @@ -51,4 +51,41 @@ std::vector> LayerNorm::Backward(const std::vector> RMSNorm::Forward(const std::vector> &input_tensors) { + CHECK_EQ(input_tensors.size(), 2); + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + + auto device = input->GetDevice().type(); + auto [output, rstd] = Dispatcher::Instance().Call, std::shared_ptr>>( + {device, "RMSNormForward"}, input, weight, eps_); + return {output, rstd}; +} + +void RMSNorm::SetupContext(const std::vector> &input_tensors, + const std::vector> &output_tensors) { + CHECK_EQ(output_tensors.size(), 2); + const auto &input = input_tensors[0]; + const auto &weight = input_tensors[1]; + const auto &rstd = output_tensors[1]; + ctx_.MarkNonDifferentiable({rstd}); + ctx_.SaveForBackward({input, weight, rstd}); +} + +std::vector> RMSNorm::Backward(const std::vector> &grad_outputs) { + auto saved_tensors = ctx_.GetSavedTensors(); + CHECK_EQ(saved_tensors.size(), 3); + const auto &input = saved_tensors[0]; + const auto &weight = saved_tensors[1]; + const auto &rstd = saved_tensors[2]; + CHECK_GE(grad_outputs.size(), 1); + const auto &grad_output = grad_outputs[0]; + + auto device = input->GetDevice().type(); + auto [grad_input, grad_weight] + = Dispatcher::Instance().Call, std::shared_ptr>>( + {device, "RMSNormBackward"}, input, weight, rstd, grad_output); + return {grad_input, grad_weight}; +} } // namespace infini_train::autograd diff --git a/infini_train/src/autograd/sparse.cc b/infini_train/src/autograd/sparse.cc index 5f515bdc1..f3a6e7396 100644 --- a/infini_train/src/autograd/sparse.cc +++ b/infini_train/src/autograd/sparse.cc @@ -19,19 +19,21 @@ void Embedding::SetupContext(const std::vector> &input_t const std::vector> &output_tensors) { const auto &input = input_tensors[0]; const auto &weight = input_tensors[1]; - weight_dims_ = weight->Dims(); - ctx_.SaveForBackward({input}); + // The weight itself (not only its dims) is needed by the backward kernel: the CUDA path routes + // the grad into a persistent per-weight sparse buffer keyed by the weight tensor. + ctx_.SaveForBackward({input, weight}); } std::vector> Embedding::Backward(const std::vector> &grad_outputs) { CHECK_EQ(grad_outputs.size(), 1); auto saved_tensors = ctx_.GetSavedTensors(); const auto &input = saved_tensors[0]; + const auto &weight = saved_tensors[1]; const auto &grad_output = grad_outputs[0]; auto device = input->GetDevice().type(); auto grad_weight = Dispatcher::Instance().Call>({device, "EmbeddingBackward"}, input, - weight_dims_, grad_output); + weight, grad_output); return {nullptr, grad_weight}; } } // namespace infini_train::autograd diff --git a/infini_train/src/kernels/cpu/embedding.cc b/infini_train/src/kernels/cpu/embedding.cc index a1da926f3..1986ed878 100644 --- a/infini_train/src/kernels/cpu/embedding.cc +++ b/infini_train/src/kernels/cpu/embedding.cc @@ -31,16 +31,16 @@ std::shared_ptr EmbeddingForward(const std::shared_ptr &input, c return output; } -std::shared_ptr EmbeddingBackward(const std::shared_ptr &input, const std::vector &weight_dims, +std::shared_ptr EmbeddingBackward(const std::shared_ptr &input, const std::shared_ptr &weight, const std::shared_ptr &grad_output) { CHECK(input->Dtype() == DataType::kINT64); - CHECK_EQ(weight_dims.size(), 2); - const int embedding_dim = weight_dims[1]; + CHECK_EQ(weight->Dims().size(), 2); + const int embedding_dim = weight->Dims()[1]; CHECK_EQ(input->Dims().size() + 1, grad_output->Dims().size()); for (int idx = 0; idx < input->Dims().size(); ++idx) { CHECK_EQ(input->Dims()[idx], grad_output->Dims()[idx]); } CHECK_EQ(*grad_output->Dims().rbegin(), embedding_dim); - auto grad_weight = std::make_shared(weight_dims, DataType::kFLOAT32); + auto grad_weight = std::make_shared(weight->Dims(), DataType::kFLOAT32); grad_weight->Fill(0.0); for (int i = 0; i < input->NumElements(); ++i) { diff --git a/infini_train/src/kernels/cpu/rmsnorm.cc b/infini_train/src/kernels/cpu/rmsnorm.cc new file mode 100644 index 000000000..39b8dfd76 --- /dev/null +++ b/infini_train/src/kernels/cpu/rmsnorm.cc @@ -0,0 +1,105 @@ +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::cpu { +std::tuple, std::shared_ptr> +RMSNormForward(const std::shared_ptr &input, const std::shared_ptr &weight, const float eps) { + /* + x: [..., embed_dim] + -> RMSNorm (w: [embed_dim]) + -> o: [..., embed_dim] + */ + // The composite path (Mean(-1)/Pow/Rsqrt/Mul) supports any rank, so the fused kernel keeps the + // same generality: one row per leading index, reducing over the last dimension. + CHECK_GE(input->Dims().size(), 2); + CHECK_EQ(input->Dims().back(), weight->Dims()[0]); + CHECK(input->Dtype() == DataType::kFLOAT32 && weight->Dtype() == DataType::kFLOAT32); + + auto input_c = input->IsContiguous() ? input : input->Contiguous(); + + const int embed_dim = static_cast(input_c->Dims().back()); + const int64_t rows = input_c->NumElements() / embed_dim; + + auto output = std::make_shared(input_c->Dims(), DataType::kFLOAT32); + auto rstd = std::make_shared(std::vector(input_c->Dims().begin(), input_c->Dims().end() - 1), + DataType::kFLOAT32); + + for (int64_t t = 0; t < rows; ++t) { + float sqsum = 0.0f; + for (int i = 0; i < embed_dim; ++i) { + float x = static_cast(input_c->DataPtr())[t * embed_dim + i]; + sqsum += x * x; + } + float s = 1.0f / sqrtf(sqsum / embed_dim + eps); + + for (int i = 0; i < embed_dim; ++i) { + float x = static_cast(input_c->DataPtr())[t * embed_dim + i]; + float n = x * s; // normalize + float o = n * static_cast(weight->DataPtr())[i]; // scale + static_cast(output->DataPtr())[t * embed_dim + i] = o; + } + // cache rstd for the backward pass later + static_cast(rstd->DataPtr())[t] = s; + } + + return {output, rstd}; +} + +std::tuple, std::shared_ptr> +RMSNormBackward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &rstd, const std::shared_ptr &grad_output) { + CHECK_GE(input->Dims().size(), 2); + CHECK_EQ(input->Dims().back(), weight->Dims()[0]); + CHECK_NE(rstd, nullptr); + CHECK(input->Dtype() == DataType::kFLOAT32 && weight->Dtype() == DataType::kFLOAT32 + && grad_output->Dtype() == DataType::kFLOAT32); + + auto input_c = input->IsContiguous() ? input : input->Contiguous(); + + const int embed_dim = static_cast(input_c->Dims().back()); + const int64_t rows = input_c->NumElements() / embed_dim; + + auto grad_input = std::make_shared(input_c->Dims(), DataType::kFLOAT32); + auto grad_weight = std::make_shared(weight->Dims(), DataType::kFLOAT32); + // grad_weight accumulates across rows; grad_input rows are fully overwritten below. + grad_weight->Fill(0.0); + + for (int64_t t = 0; t < rows; ++t) { + float rstd_t = static_cast(rstd->DataPtr())[t]; + + // S1 = sum_i(g_i * w_i * x_i); K = (S1 / H) * rstd + float S1 = 0.0f; + for (int i = 0; i < embed_dim; ++i) { + float x = static_cast(input_c->DataPtr())[t * embed_dim + i]; + float w = static_cast(weight->DataPtr())[i]; + float g = static_cast(grad_output->DataPtr())[t * embed_dim + i]; + S1 += g * w * x; + } + float K = (S1 / embed_dim) * rstd_t; + + for (int i = 0; i < embed_dim; ++i) { + float x = static_cast(input_c->DataPtr())[t * embed_dim + i]; + float w = static_cast(weight->DataPtr())[i]; + float g = static_cast(grad_output->DataPtr())[t * embed_dim + i]; + float n = x * rstd_t; + static_cast(grad_input->DataPtr())[t * embed_dim + i] = (g * w - K * n) * rstd_t; + static_cast(grad_weight->DataPtr())[i] += g * n; + } + } + return {grad_input, grad_weight}; +} +} // namespace infini_train::kernels::cpu + +#define REGISTER_CPU_RMSNORM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCPU, kernel_name, infini_train::kernels::cpu::kernel_name) + +REGISTER_CPU_RMSNORM_KERNEL(RMSNormForward) +REGISTER_CPU_RMSNORM_KERNEL(RMSNormBackward) + +#undef REGISTER_CPU_RMSNORM_KERNEL diff --git a/infini_train/src/kernels/cuda/accumulate_grad.cu b/infini_train/src/kernels/cuda/accumulate_grad.cu index 93409a7ef..ee39fd850 100644 --- a/infini_train/src/kernels/cuda/accumulate_grad.cu +++ b/infini_train/src/kernels/cuda/accumulate_grad.cu @@ -1,5 +1,10 @@ +#include +#include #include +#include +#include #include +#include #include "infini_train/include/common/cuda/kernel_helper.cuh" #include "infini_train/include/core/runtime/device_guard.h" @@ -10,6 +15,107 @@ #include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" namespace infini_train::kernels::cuda { +namespace { + +// Aligned vector type for vectorized loads/stores (up to 128-bit). The alignment matches the payload +// size so NVCC emits the widest legal access for it (16B -> .128, 8B -> .64, 4B -> .32). +template struct __align__(sizeof(T) * N) aligned_vector { T val[N]; }; + +// Elements per 128-bit access: float -> 4, bf16/half -> 8, double -> 2. +template constexpr int kVecSize = 16 / sizeof(T); + +// Vector width for kernels touching two dtypes (e.g. fp32 master weight + bf16 shadow). Sizing by the +// wider element keeps both payloads inside a single 128-bit access; the narrower one then uses a +// 64/32-bit access, still one instruction instead of VecSize scalar ones. +template constexpr int kMixedVecSize = 16 / (sizeof(T) > sizeof(U) ? sizeof(T) : sizeof(U)); + +// Vectorized access also needs the payload width to divide the pointer. Optimizer state is freshly +// allocated (hence 16B-aligned), but a param that is a view into a larger buffer may not be. +inline bool IsAlignedTo(const void *ptr, size_t bytes) { return (reinterpret_cast(ptr) % bytes) == 0; } + +// SM count for a device ordinal, resolved once and cached the way cuda_guard_impl.cc caches streams +// and handles. Returns 0 if the query fails or the ordinal is out of range. +inline int MultiProcessorCount(int index) { + constexpr int kMaxGpus = 8; + static std::array, kMaxGpus> cache{}; // 0 == not resolved yet + + if (index < 0 || index >= kMaxGpus) { + return 0; + } + int sms = cache[index].load(std::memory_order_relaxed); + if (sms != 0) { + return sms; + } + // Threads racing here query and store the same value, so no locking is needed. + if (cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, index) != cudaSuccess || sms <= 0) { + return 0; + } + cache[index].store(sms, std::memory_order_relaxed); + return sms; +} + +// Smallest tensor worth handing to the vectorized kernel. That kernel launches vec_size-times fewer +// threads and needs 46-57 registers/thread against the scalar kernel's 17-19, because every lane of +// all four payloads stays live from load to store; that caps it near half occupancy and gives its +// longer straight-line body a fixed cost of ~0.3us (vec_size 4) to ~1.3us (vec_size 8) on A100. +// Measured there it only overtakes the scalar kernel once the grid gives every SM at least one +// block, i.e. sms * threads_per_block vectors. Falls back to "one full vector" if the query failed. +inline size_t MinVectorizedElements(int device_index, int vec_size, int threads_per_block) { + const int sms = MultiProcessorCount(device_index); + if (sms <= 0) { + return static_cast(vec_size); + } + return static_cast(sms) * threads_per_block * vec_size; +} + +// Loop-invariant Adam scalars. beta1/beta2 are cast to T once per thread instead of once per element; +// the cast is deterministic, so the values fed to the math are unchanged. +template struct AdamParams { + T beta1; + T one_minus_beta1; + T beta2; + T one_minus_beta2; + float learning_rate; + float bias_correction_m; + float bias_correction_v; + float eps; +}; + +template +__device__ __forceinline__ AdamParams MakeAdamParams(float learning_rate, float beta1, float beta2, float eps, + float bias_correction_m, float bias_correction_v) { + return AdamParams{common::cuda::Cast(beta1), + common::cuda::Cast(1 - beta1), + common::cuda::Cast(beta2), + common::cuda::Cast(1 - beta2), + learning_rate, + bias_correction_m, + bias_correction_v, + eps}; +} + +// One Adam update step for a single element: m/v EMAs, bias correction, then the param update. +// Shared by the scalar and vectorized kernels so both paths stay bit-identical. +template +__device__ __forceinline__ void AdamUpdateElement(const T &grad, T ¶m, T &m, T &v, const AdamParams &p) { + m = common::cuda::Fma(p.beta1, m, p.one_minus_beta1 * grad); + v = common::cuda::Fma(p.beta2, v, p.one_minus_beta2 * grad * grad); + + const float m_hat = common::cuda::Cast(m) / p.bias_correction_m; + const float v_hat = common::cuda::Cast(v) / p.bias_correction_v; + const float rcp = __frcp_rn(__fsqrt_rn(v_hat) + p.eps); + + if constexpr (std::is_same_v) { + // Spell out the contraction of "param - (lr * m_hat) * rcp". Left implicit, nvcc fuses it into + // one FFMA in the scalar kernel but emits FMUL+FADD in the unrolled vectorized one, which costs + // 1 ULP and would make the result depend on which of the two paths the dispatch picked. + param = std::fma(-(p.learning_rate * m_hat), rcp, param); + } else { + param = common::cuda::Sub(param, common::cuda::Cast((p.learning_rate * m_hat) * rcp)); + } +} + +} // namespace template __global__ void AccumulateGradKernel(const T *grad_ptr, float rate, T *tensor_ptr, size_t num_elements) { @@ -43,18 +149,210 @@ template __global__ void AdamAccumulateGradKernel(const T *grad_data, T *param_data, size_t num_elements, T *m_data, T *v_data, float learning_rate, float beta1, float beta2, float eps, const float bias_correction_m, const float bias_correction_v) { + const AdamParams p = MakeAdamParams(learning_rate, beta1, beta2, eps, bias_correction_m, bias_correction_v); size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < num_elements) { - m_data[idx] = common::cuda::Fma(common::cuda::Cast(beta1), m_data[idx], - common::cuda::Cast(1 - beta1) * grad_data[idx]); - v_data[idx] = common::cuda::Fma(common::cuda::Cast(beta2), v_data[idx], - common::cuda::Cast(1 - beta2) * grad_data[idx] * grad_data[idx]); + AdamUpdateElement(grad_data[idx], param_data[idx], m_data[idx], v_data[idx], p); + } +} + +// Vectorized Adam update: each thread handles VecSize contiguous elements through wide accesses, +// cutting load/store instruction count (and hence Long Scoreboard stalls) by VecSize on this +// DRAM-bound kernel. Numerically identical to AdamAccumulateGradKernel via AdamUpdateElement. +// The trailing num_elements % VecSize elements are finished by a scalar pass in the same launch. +template +__global__ void AdamAccumulateGradKernelVectorized(const T *__restrict__ grad_data, T *__restrict__ param_data, + size_t num_elements, T *__restrict__ m_data, T *__restrict__ v_data, + float learning_rate, float beta1, float beta2, float eps, + const float bias_correction_m, const float bias_correction_v) { + using VecT = aligned_vector; + const AdamParams p = MakeAdamParams(learning_rate, beta1, beta2, eps, bias_correction_m, bias_correction_v); + const size_t num_vecs = num_elements / VecSize; + const size_t grid_stride = static_cast(gridDim.x) * blockDim.x; + + for (size_t vid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; vid < num_vecs; vid += grid_stride) { + const size_t base = vid * VecSize; + + // Vectorized loads + VecT grad_vec = *reinterpret_cast(&grad_data[base]); + VecT param_vec = *reinterpret_cast(¶m_data[base]); + VecT m_vec = *reinterpret_cast(&m_data[base]); + VecT v_vec = *reinterpret_cast(&v_data[base]); + +#pragma unroll + for (int i = 0; i < VecSize; ++i) { + AdamUpdateElement(grad_vec.val[i], param_vec.val[i], m_vec.val[i], v_vec.val[i], p); + } + + // Vectorized stores + *reinterpret_cast(¶m_data[base]) = param_vec; + *reinterpret_cast(&m_data[base]) = m_vec; + *reinterpret_cast(&v_data[base]) = v_vec; + } + + // Tail: numel % VecSize != 0 + const size_t tail_start = num_vecs * VecSize; + for (size_t idx = tail_start + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < num_elements; + idx += grid_stride) { + AdamUpdateElement(grad_data[idx], param_data[idx], m_data[idx], v_data[idx], p); + } +} + +template +__global__ void AdamAccumulateGradShadowKernel(const T *grad_data, T *param_data, TShadow *shadow_data, + size_t num_elements, T *m_data, T *v_data, float learning_rate, + float beta1, float beta2, float eps, const float bias_correction_m, + const float bias_correction_v) { + const AdamParams p = MakeAdamParams(learning_rate, beta1, beta2, eps, bias_correction_m, bias_correction_v); + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + + if (idx < num_elements) { + AdamUpdateElement(grad_data[idx], param_data[idx], m_data[idx], v_data[idx], p); + shadow_data[idx] = common::cuda::Cast(param_data[idx]); + } +} + +// Vectorized counterpart of AdamAccumulateGradShadowKernel. T and TShadow may differ in width, so +// VecSize is sized by the wider of the two (see kMixedVecSize): the wider dtype gets a 128-bit access +// and the narrower one a proportionally smaller but still single vector access. +template +__global__ void AdamAccumulateGradShadowKernelVectorized(const T *__restrict__ grad_data, T *__restrict__ param_data, + TShadow *__restrict__ shadow_data, size_t num_elements, + T *__restrict__ m_data, T *__restrict__ v_data, + float learning_rate, float beta1, float beta2, float eps, + const float bias_correction_m, const float bias_correction_v) { + using VecT = aligned_vector; + using VecTShadow = aligned_vector; + const AdamParams p = MakeAdamParams(learning_rate, beta1, beta2, eps, bias_correction_m, bias_correction_v); + const size_t num_vecs = num_elements / VecSize; + const size_t grid_stride = static_cast(gridDim.x) * blockDim.x; + + for (size_t vid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; vid < num_vecs; vid += grid_stride) { + const size_t base = vid * VecSize; + + // Vectorized loads + VecT grad_vec = *reinterpret_cast(&grad_data[base]); + VecT param_vec = *reinterpret_cast(¶m_data[base]); + VecT m_vec = *reinterpret_cast(&m_data[base]); + VecT v_vec = *reinterpret_cast(&v_data[base]); + + VecTShadow shadow_vec; +#pragma unroll + for (int i = 0; i < VecSize; ++i) { + AdamUpdateElement(grad_vec.val[i], param_vec.val[i], m_vec.val[i], v_vec.val[i], p); + shadow_vec.val[i] = common::cuda::Cast(param_vec.val[i]); + } + + // Vectorized stores + *reinterpret_cast(¶m_data[base]) = param_vec; + *reinterpret_cast(&m_data[base]) = m_vec; + *reinterpret_cast(&v_data[base]) = v_vec; + *reinterpret_cast(&shadow_data[base]) = shadow_vec; + } + + // Tail: numel % VecSize != 0 + const size_t tail_start = num_vecs * VecSize; + for (size_t idx = tail_start + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < num_elements; + idx += grid_stride) { + AdamUpdateElement(grad_data[idx], param_data[idx], m_data[idx], v_data[idx], p); + shadow_data[idx] = common::cuda::Cast(param_data[idx]); + } +} + +// Sparse-row Adam: walk only the rows claimed since the last clear (row_list[0, *count)) instead +// of the whole parameter. Rows that were never hit keep their param/m/v at their initialization, +// which is exactly what a dense step computes for a zero row with zero EMAs; rows hit in an +// earlier step but idle now are left untouched (LazyAdam semantics) rather than decaying in place. +// The row count lives on the device, so the host never synchronizes to size the grid. One block +// sweeps rows grid-strided; within a row the access pattern mirrors AdamAccumulateGradKernel. +template +__global__ void AdamSparseRowsKernel(const T *__restrict__ grad_data, T *__restrict__ param_data, + const int32_t *__restrict__ row_list, const int32_t *__restrict__ count, + size_t dim, T *__restrict__ m_data, T *__restrict__ v_data, float learning_rate, + float beta1, float beta2, float eps, const float bias_correction_m, + const float bias_correction_v) { + using VecT = aligned_vector; + const AdamParams p = MakeAdamParams(learning_rate, beta1, beta2, eps, bias_correction_m, bias_correction_v); + const int num_rows = *count; + const size_t num_vecs = dim / VecSize; + const size_t tail_start = num_vecs * VecSize; + + for (int r = blockIdx.x; r < num_rows; r += gridDim.x) { + const size_t base = static_cast(row_list[r]) * dim; + const T *row_grad = grad_data + base; + T *row_param = param_data + base; + T *row_m = m_data + base; + T *row_v = v_data + base; + + for (size_t vid = threadIdx.x; vid < num_vecs; vid += blockDim.x) { + const size_t idx = vid * VecSize; + VecT grad_vec = *reinterpret_cast(&row_grad[idx]); + VecT param_vec = *reinterpret_cast(&row_param[idx]); + VecT m_vec = *reinterpret_cast(&row_m[idx]); + VecT v_vec = *reinterpret_cast(&row_v[idx]); +#pragma unroll + for (int i = 0; i < VecSize; ++i) { + AdamUpdateElement(grad_vec.val[i], param_vec.val[i], m_vec.val[i], v_vec.val[i], p); + } + *reinterpret_cast(&row_param[idx]) = param_vec; + *reinterpret_cast(&row_m[idx]) = m_vec; + *reinterpret_cast(&row_v[idx]) = v_vec; + } + + // Tail: dim % VecSize != 0 (empty when it divides) + for (size_t idx = tail_start + threadIdx.x; idx < dim; idx += blockDim.x) { + AdamUpdateElement(row_grad[idx], row_param[idx], row_m[idx], row_v[idx], p); + } + } +} + +// Sparse-row counterpart of AdamAccumulateGradShadowKernel/VecSize conventions. +template +__global__ void AdamSparseRowsShadowKernel(const T *__restrict__ grad_data, T *__restrict__ param_data, + TShadow *__restrict__ shadow_data, const int32_t *__restrict__ row_list, + const int32_t *__restrict__ count, size_t dim, T *__restrict__ m_data, + T *__restrict__ v_data, float learning_rate, float beta1, float beta2, + float eps, const float bias_correction_m, const float bias_correction_v) { + using VecT = aligned_vector; + using VecTShadow = aligned_vector; + const AdamParams p = MakeAdamParams(learning_rate, beta1, beta2, eps, bias_correction_m, bias_correction_v); + const int num_rows = *count; + const size_t num_vecs = dim / VecSize; + const size_t tail_start = num_vecs * VecSize; + + for (int r = blockIdx.x; r < num_rows; r += gridDim.x) { + const size_t base = static_cast(row_list[r]) * dim; + const T *row_grad = grad_data + base; + T *row_param = param_data + base; + TShadow *row_shadow = shadow_data + base; + T *row_m = m_data + base; + T *row_v = v_data + base; + + for (size_t vid = threadIdx.x; vid < num_vecs; vid += blockDim.x) { + const size_t idx = vid * VecSize; + VecT grad_vec = *reinterpret_cast(&row_grad[idx]); + VecT param_vec = *reinterpret_cast(&row_param[idx]); + VecT m_vec = *reinterpret_cast(&row_m[idx]); + VecT v_vec = *reinterpret_cast(&row_v[idx]); - const float m_hat = common::cuda::Cast(m_data[idx]) / bias_correction_m; - const float v_hat = common::cuda::Cast(v_data[idx]) / bias_correction_v; + VecTShadow shadow_vec; +#pragma unroll + for (int i = 0; i < VecSize; ++i) { + AdamUpdateElement(grad_vec.val[i], param_vec.val[i], m_vec.val[i], v_vec.val[i], p); + shadow_vec.val[i] = common::cuda::Cast(param_vec.val[i]); + } - param_data[idx] = common::cuda::Sub( - param_data[idx], common::cuda::Cast(learning_rate * m_hat * __frcp_rn(__fsqrt_rn(v_hat) + eps))); + *reinterpret_cast(&row_param[idx]) = param_vec; + *reinterpret_cast(&row_m[idx]) = m_vec; + *reinterpret_cast(&row_v[idx]) = v_vec; + *reinterpret_cast(&row_shadow[idx]) = shadow_vec; + } + + // Tail: dim % VecSize != 0 (empty when it divides) + for (size_t idx = tail_start + threadIdx.x; idx < dim; idx += blockDim.x) { + AdamUpdateElement(row_grad[idx], row_param[idx], row_m[idx], row_v[idx], p); + row_shadow[idx] = common::cuda::Cast(row_param[idx]); + } } } @@ -67,7 +365,6 @@ void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_p const float bias_correction_v = 1.0f - std::pow(beta2, t); int threads_per_block = 256; - int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; auto device = grad->GetDevice(); const auto &cuda_stream = dynamic_cast( @@ -77,13 +374,197 @@ void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_p core::cuda::DispatchCudaFunc( grad->Dtype(), [=]() { - AdamAccumulateGradKernel<<>>( - static_cast(grad->DataPtr()), static_cast(param->DataPtr()), num_elements, - static_cast(m->DataPtr()), static_cast(v->DataPtr()), learning_rate, beta1, beta2, eps, - bias_correction_m, bias_correction_v); + const T *grad_ptr = static_cast(grad->DataPtr()); + T *param_ptr = static_cast(param->DataPtr()); + T *m_ptr = static_cast(m->DataPtr()); + T *v_ptr = static_cast(v->DataPtr()); + + constexpr int vec_size = kVecSize; + // Take the vectorized path only when every operand can legally serve a wide access and the + // tensor is big enough for the vec_size-times narrower grid to still fill the device. + // Anything else falls back to the scalar kernel, which produces identical values. + const size_t min_elements = MinVectorizedElements(device.index(), vec_size, threads_per_block); + const bool can_vectorize = num_elements >= min_elements && IsAlignedTo(grad_ptr, sizeof(T) * vec_size) + && IsAlignedTo(param_ptr, sizeof(T) * vec_size) + && IsAlignedTo(m_ptr, sizeof(T) * vec_size) + && IsAlignedTo(v_ptr, sizeof(T) * vec_size); + + if (can_vectorize) { + const size_t num_vecs = num_elements / vec_size; + int num_blocks = (num_vecs + threads_per_block - 1) / threads_per_block; + AdamAccumulateGradKernelVectorized<<>>( + grad_ptr, param_ptr, num_elements, m_ptr, v_ptr, learning_rate, beta1, beta2, eps, + bias_correction_m, bias_correction_v); + } else { + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + AdamAccumulateGradKernel<<>>( + grad_ptr, param_ptr, num_elements, m_ptr, v_ptr, learning_rate, beta1, beta2, eps, + bias_correction_m, bias_correction_v); + } }, "CUDA AdamAccumulateGrad"); } + +void AdamAccumulateGradShadow(const std::shared_ptr &grad, const std::shared_ptr ¶m, + const std::shared_ptr &shadow, const std::shared_ptr &m, + const std::shared_ptr &v, float learning_rate, float beta1, float beta2, + float eps, int64_t t) { + size_t num_elements = grad->NumElements(); + + const float bias_correction_m = 1.0f - std::pow(beta1, t); + const float bias_correction_v = 1.0f - std::pow(beta2, t); + + int threads_per_block = 256; + + auto device = grad->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + + core::cuda::DispatchCudaFunc( + grad->Dtype(), + [=]() { + core::cuda::DispatchCudaFunc( + shadow->Dtype(), + [=]() { + const T *grad_ptr = static_cast(grad->DataPtr()); + T *param_ptr = static_cast(param->DataPtr()); + TShadow *shadow_ptr = static_cast(shadow->DataPtr()); + T *m_ptr = static_cast(m->DataPtr()); + T *v_ptr = static_cast(v->DataPtr()); + + // T and TShadow may differ in width, so each operand is checked against its own + // payload width; the shadow check is what makes a narrower TShadow legal here. + constexpr int vec_size = kMixedVecSize; + const size_t min_elements = MinVectorizedElements(device.index(), vec_size, threads_per_block); + const bool can_vectorize + = num_elements >= min_elements && IsAlignedTo(grad_ptr, sizeof(T) * vec_size) + && IsAlignedTo(param_ptr, sizeof(T) * vec_size) && IsAlignedTo(m_ptr, sizeof(T) * vec_size) + && IsAlignedTo(v_ptr, sizeof(T) * vec_size) + && IsAlignedTo(shadow_ptr, sizeof(TShadow) * vec_size); + + if (can_vectorize) { + const size_t num_vecs = num_elements / vec_size; + int num_blocks = (num_vecs + threads_per_block - 1) / threads_per_block; + AdamAccumulateGradShadowKernelVectorized + <<>>( + grad_ptr, param_ptr, shadow_ptr, num_elements, m_ptr, v_ptr, learning_rate, beta1, + beta2, eps, bias_correction_m, bias_correction_v); + } else { + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + AdamAccumulateGradShadowKernel<<>>( + grad_ptr, param_ptr, shadow_ptr, num_elements, m_ptr, v_ptr, learning_rate, beta1, beta2, + eps, bias_correction_m, bias_correction_v); + } + }, + "CUDA AdamAccumulateGradShadow"); + }, + "CUDA AdamAccumulateGradShadow"); +} + +void AdamSparseRows(const std::shared_ptr &grad, const std::shared_ptr ¶m, + const std::shared_ptr &m, const std::shared_ptr &v, + const std::shared_ptr &row_list, const std::shared_ptr &count, float learning_rate, + float beta1, float beta2, float eps, int64_t t) { + const float bias_correction_m = 1.0f - std::pow(beta1, t); + const float bias_correction_v = 1.0f - std::pow(beta2, t); + + constexpr int threads_per_block = 256; + // Grid-strided over rows and sized for the common case (a few hundred rows = one row per + // block); idle blocks only read *count and exit. + constexpr int max_blocks = 512; + const size_t dim = static_cast(grad->Dims()[1]); + + auto device = grad->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + + core::cuda::DispatchCudaFunc( + grad->Dtype(), + [=]() { + const T *grad_ptr = static_cast(grad->DataPtr()); + T *param_ptr = static_cast(param->DataPtr()); + T *m_ptr = static_cast(m->DataPtr()); + T *v_ptr = static_cast(v->DataPtr()); + const int32_t *row_ptr = static_cast(row_list->DataPtr()); + const int32_t *count_ptr = static_cast(count->DataPtr()); + + // Rows start at multiples of dim, so a whole-row vector access needs dim % vec_size + // == 0 plus aligned bases; anything else falls back to the scalar form (identical + // values, VecSize == 1). + constexpr int vec_size = kVecSize; + const bool can_vectorize = (dim % vec_size == 0) && IsAlignedTo(grad_ptr, sizeof(T) * vec_size) + && IsAlignedTo(param_ptr, sizeof(T) * vec_size) + && IsAlignedTo(m_ptr, sizeof(T) * vec_size) + && IsAlignedTo(v_ptr, sizeof(T) * vec_size); + + if (can_vectorize) { + AdamSparseRowsKernel<<>>( + grad_ptr, param_ptr, row_ptr, count_ptr, dim, m_ptr, v_ptr, learning_rate, beta1, beta2, eps, + bias_correction_m, bias_correction_v); + } else { + AdamSparseRowsKernel<<>>( + grad_ptr, param_ptr, row_ptr, count_ptr, dim, m_ptr, v_ptr, learning_rate, beta1, beta2, eps, + bias_correction_m, bias_correction_v); + } + }, + "CUDA AdamSparseRows"); +} + +void AdamSparseRowsShadow(const std::shared_ptr &grad, const std::shared_ptr ¶m, + const std::shared_ptr &shadow, const std::shared_ptr &m, + const std::shared_ptr &v, const std::shared_ptr &row_list, + const std::shared_ptr &count, float learning_rate, float beta1, float beta2, + float eps, int64_t t) { + const float bias_correction_m = 1.0f - std::pow(beta1, t); + const float bias_correction_v = 1.0f - std::pow(beta2, t); + + constexpr int threads_per_block = 256; + constexpr int max_blocks = 512; + const size_t dim = static_cast(grad->Dims()[1]); + + auto device = grad->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + + core::cuda::DispatchCudaFunc( + grad->Dtype(), + [=]() { + core::cuda::DispatchCudaFunc( + shadow->Dtype(), + [=]() { + const T *grad_ptr = static_cast(grad->DataPtr()); + T *param_ptr = static_cast(param->DataPtr()); + TShadow *shadow_ptr = static_cast(shadow->DataPtr()); + T *m_ptr = static_cast(m->DataPtr()); + T *v_ptr = static_cast(v->DataPtr()); + const int32_t *row_ptr = static_cast(row_list->DataPtr()); + const int32_t *count_ptr = static_cast(count->DataPtr()); + + constexpr int vec_size = kMixedVecSize; + const bool can_vectorize = (dim % vec_size == 0) && IsAlignedTo(grad_ptr, sizeof(T) * vec_size) + && IsAlignedTo(param_ptr, sizeof(T) * vec_size) + && IsAlignedTo(m_ptr, sizeof(T) * vec_size) + && IsAlignedTo(v_ptr, sizeof(T) * vec_size) + && IsAlignedTo(shadow_ptr, sizeof(TShadow) * vec_size); + + if (can_vectorize) { + AdamSparseRowsShadowKernel + <<>>( + grad_ptr, param_ptr, shadow_ptr, row_ptr, count_ptr, dim, m_ptr, v_ptr, learning_rate, + beta1, beta2, eps, bias_correction_m, bias_correction_v); + } else { + AdamSparseRowsShadowKernel<<>>( + grad_ptr, param_ptr, shadow_ptr, row_ptr, count_ptr, dim, m_ptr, v_ptr, learning_rate, + beta1, beta2, eps, bias_correction_m, bias_correction_v); + } + }, + "CUDA AdamSparseRowsShadow"); + }, + "CUDA AdamSparseRowsShadow"); +} } // namespace infini_train::kernels::cuda #define REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(kernel_name) \ @@ -91,5 +572,7 @@ void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_p REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) - +REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGradShadow) +REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AdamSparseRows) +REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL(AdamSparseRowsShadow) #undef REGISTER_CUDA_ACCUMULATE_GRAD_KERNEL diff --git a/infini_train/src/kernels/cuda/cross_entropy.cu b/infini_train/src/kernels/cuda/cross_entropy.cu index 96f987fe5..d0af9dcea 100644 --- a/infini_train/src/kernels/cuda/cross_entropy.cu +++ b/infini_train/src/kernels/cuda/cross_entropy.cu @@ -99,20 +99,23 @@ std::shared_ptr CrossEntropyForward(const std::shared_ptr &input const Ttarget *target_ptr = static_cast(target->DataPtr()); const Tinput *input_ptr = static_cast(input->DataPtr()); Tinput *batched_loss_ptr = static_cast(batched_output->DataPtr()); - // FIXME(dcj): do reduce on GPU CrossEntropyForwardKernel <<>>(input_ptr, target_ptr, batched_loss_ptr, bs, num_classes); - auto loss_cpu = batched_output->To(Device()); - auto loss = std::make_shared(std::vector{}, input->Dtype(), Device()); - auto loss_cpu_typed_ptr = static_cast(loss_cpu.DataPtr()); - static_cast(loss->DataPtr())[0] - = std::accumulate(loss_cpu_typed_ptr, loss_cpu_typed_ptr + bs, 0.0f, - [](float acc, const Tinput &val) { return acc + common::cuda::Cast(val); }) - / bs; - - return std::make_shared(loss->To(input->GetDevice())); + // Reduce the per-sample losses to the scalar mean on-device. The previous host-side + // accumulate needed a D2H copy of batched_output into pageable memory, which is synchronous + // w.r.t. the host and stalled every forward step until the whole stream queue drained. + // Accumulate in fp32 (as the old host-side float accumulate did), then narrow once. + auto batched_f32 = input->Dtype() == DataType::kFLOAT32 + ? batched_output + : std::make_shared(batched_output->To(DataType::kFLOAT32)); + auto mean = Dispatcher::Instance().Call>({device.type(), "MeanForward"}, + batched_f32, int64_t{0}, false); + if (mean->Dtype() != input->Dtype()) { + mean = std::make_shared(mean->To(input->Dtype())); + } + return mean; }, "CUDA CrossEntropyForward"); } diff --git a/infini_train/src/kernels/cuda/elementwise.cu b/infini_train/src/kernels/cuda/elementwise.cu index 755f7433b..976ced53b 100644 --- a/infini_train/src/kernels/cuda/elementwise.cu +++ b/infini_train/src/kernels/cuda/elementwise.cu @@ -127,6 +127,35 @@ __global__ void BinaryForwardKernelNoBroadcast(T *__restrict__ output, Func fn, } } +// Mixed-input forward (broadcast path): operand a has dtype Ta, b has dtype Tb, output has dtype +// Tout (= PromoteDataTypes(Ta, Tb)). Both operands are widened to Tout *in registers*, so no +// intermediate Cast kernel is materialized. Numerically identical to promoting each input with +// Tensor::To(Tout) before the op, because Cast is exactly what To(Tout) applies per element. +template +__global__ void BinaryForwardKernelMixed(Tout *output, Func fn, BroadcastMeta meta, const Ta *a, const Tb *b, + size_t num_elements) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_elements) { + return; + } + + int64_t a_offset = CalcOffset(idx, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides); + int64_t b_offset = CalcOffset(idx, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides); + + output[idx] = fn(common::cuda::Cast(a[a_offset]), common::cuda::Cast(b[b_offset])); +} + +// Mixed-input forward fast path: no broadcast, contiguous tensors — skip CalcOffset entirely. +template +__global__ void BinaryForwardKernelNoBroadcastMixed(Tout *__restrict__ output, Func fn, const Ta *__restrict__ a, + const Tb *__restrict__ b, size_t num_elements) { + const size_t grid_stride = static_cast(gridDim.x) * blockDim.x; + for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < num_elements; + idx += grid_stride) { + output[idx] = fn(common::cuda::Cast(a[idx]), common::cuda::Cast(b[idx])); + } +} + // Fast path backward: no broadcast, contiguous — skip CalcOffset entirely template __global__ void BinaryBackwardKernelNoBroadcastFast(T *__restrict__ outA, T *__restrict__ outB, FuncA fn_a, FuncB fn_b, @@ -194,6 +223,22 @@ __global__ void BinaryBackwardKernelNoBroadcastVectorized(T *__restrict__ outA, } } +// Mixed-input backward fast path: no broadcast, contiguous. grad_out/outputs are Tout; saved operands +// a/b are Ta/Tb and widened to Tout in registers, so no Cast kernel is materialized. Bit-identical to +// promoting a/b with To(Tout) before the homogeneous backward. +template +__global__ void BinaryBackwardKernelNoBroadcastFastMixed(Tout *__restrict__ outA, Tout *__restrict__ outB, FuncA fn_a, + FuncB fn_b, size_t numel, const Tout *__restrict__ grad_out, + const Ta *__restrict__ inA, const Tb *__restrict__ inB) { + const size_t grid_stride = static_cast(gridDim.x) * blockDim.x; + for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < numel; idx += grid_stride) { + const Tout a = inA ? common::cuda::Cast(inA[idx]) : Tout(0); + const Tout b = inB ? common::cuda::Cast(inB[idx]) : Tout(0); + outA[idx] = Mul(grad_out[idx], fn_a(a, b)); + outB[idx] = Mul(grad_out[idx], fn_b(a, b)); + } +} + // Helper to choose optimal block size based on tensor size inline size_t ChooseBlockSize(size_t num_elements) { if (num_elements < 1024) { @@ -286,6 +331,90 @@ void LaunchForward(Func func, const std::shared_ptr &output, const Input } } +// Mixed-input forward launcher: reads input_a as Ta and input_b as Tb, computes in Tout and writes +// an output of dtype Tout. Used by BinaryForward to avoid materializing a Cast kernel when the two +// operands have different dtypes (e.g. bf16 activation * f32 RoPE cos/sin, f32 residual + bf16 block +// output). Requires both inputs to be contiguous; the broadcast metadata assumes contiguous strides. +template +void LaunchForwardMixed(Func func, const std::shared_ptr &output, const std::shared_ptr &input_a, + const std::shared_ptr &input_b) { + auto device = output->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + Tout *output_ptr = static_cast(output->DataPtr()); + + const auto &a_dims = input_a->Dims(); + const auto &b_dims = input_b->Dims(); + const auto &out_dims = output->Dims(); + const size_t num_elements = output->NumElements(); + + const Ta *a_ptr = static_cast(input_a->DataPtr()); + const Tb *b_ptr = static_cast(input_b->DataPtr()); + dim3 block_dims = ChooseBlockDims(num_elements); + + if (ShapesEqual(a_dims, out_dims) && ShapesEqual(b_dims, out_dims)) { + dim3 grid_dims(std::min(CEIL_DIV(num_elements, block_dims.x), static_cast(65535))); + BinaryForwardKernelNoBroadcastMixed + <<>>(output_ptr, func, a_ptr, b_ptr, num_elements); + } else { + BroadcastMeta meta = MakeBroadcastMeta(a_dims, b_dims, out_dims); + dim3 grid_dims(CEIL_DIV(num_elements, block_dims.x)); + BinaryForwardKernelMixed + <<>>(output_ptr, func, meta, a_ptr, b_ptr, num_elements); + } +} + +// Dispatch the mixed-input forward for the supported (Ta, Tb) combinations. The promoted output type +// for any two *different* floating dtypes is always float32 (see PromoteDataTypes), so Tout is fixed +// to float here. Returns false when the combination is not covered (e.g. non-floating, or a/b not +// contiguous), letting the caller fall back to the promote-via-To() path. +template +bool LaunchBinaryForwardMixed(Func func, const std::shared_ptr &output, const std::shared_ptr &a, + const std::shared_ptr &b, DataType a_dtype, DataType b_dtype, + DataType promoted_type) { + if (promoted_type != DataType::kFLOAT32 || !a->IsContiguous() || !b->IsContiguous()) { + return false; + } + switch (a_dtype) { + case DataType::kFLOAT32: + switch (b_dtype) { + case DataType::kBFLOAT16: + LaunchForwardMixed(func, output, a, b); + return true; + case DataType::kFLOAT16: + LaunchForwardMixed(func, output, a, b); + return true; + default: + return false; + } + case DataType::kBFLOAT16: + switch (b_dtype) { + case DataType::kFLOAT32: + LaunchForwardMixed(func, output, a, b); + return true; + case DataType::kFLOAT16: + LaunchForwardMixed(func, output, a, b); + return true; + default: + return false; + } + case DataType::kFLOAT16: + switch (b_dtype) { + case DataType::kFLOAT32: + LaunchForwardMixed(func, output, a, b); + return true; + case DataType::kBFLOAT16: + LaunchForwardMixed(func, output, a, b); + return true; + default: + return false; + } + default: + return false; + } +} + // Backward kernel for unary operators template __global__ void UnaryBackwardKernel(T *output, Func fn, size_t num_elements, size_t offset, const T *grad_output, @@ -551,6 +680,76 @@ __global__ void BinaryBackwardKernel(T *output_a, T *output_b, FuncA fn_a, FuncB } } +// Mixed-input variant of the float broadcast backward kernel above. Identical control flow and float +// accumulation; the only difference is that the saved operands are read as Ta/Tb and widened to Tout +// (= float) in registers, so BinaryBackward never materializes a bf16->f32 Cast kernel for them. +template +__global__ void BinaryBackwardKernelMixed(Tout *output_a, Tout *output_b, FuncA fn_a, FuncB fn_b, BroadcastMeta meta, + size_t num_elements, const Tout *grad_output, const Ta *input_a, + const Tb *input_b) { + extern __shared__ char shared_memory[]; + const int tid = threadIdx.x; + const int lane_id = tid % kLogicalWarpSize; + const int logical_warp_id = tid / kLogicalWarpSize; + + using WarpReduce = cub::WarpReduce; + auto *temp_storage = reinterpret_cast(shared_memory); + + size_t idx = blockIdx.x * blockDim.x + tid; + bool in_bounds = (idx < num_elements); + + int64_t a_offset = 0, b_offset = 0; + Tout a_val = Tout(0), b_val = Tout(0); + float grad_val = 0.0f; + + if (in_bounds) { + a_offset = CalcOffset(idx, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides); + b_offset = CalcOffset(idx, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides); + a_val = input_a ? common::cuda::Cast(input_a[a_offset]) : Tout(0); + b_val = input_b ? common::cuda::Cast(input_b[b_offset]) : Tout(0); + output_a[a_offset] = Mul(grad_output[idx], fn_a(a_val, b_val)); + grad_val = common::cuda::Cast(Mul(grad_output[idx], fn_b(a_val, b_val))); + } + + using WarpMask = decltype(__ballot_sync(~uint64_t{0}, true)); + const WarpMask full_mask = ~WarpMask{0}; + const WarpMask physical_active_mask = __ballot_sync(full_mask, in_bounds); + const int physical_lane = tid % warpSize; + const int logical_base = (physical_lane / kLogicalWarpSize) * kLogicalWarpSize; + const WarpMask logical_lane_mask = static_cast(uint64_t{0xffffffff} << logical_base); + const WarpMask active_mask = physical_active_mask & logical_lane_mask; + if (active_mask == 0) { + return; + } + + const unsigned logical_active_mask = static_cast(static_cast(active_mask) >> logical_base); + const int leader = __ffs(logical_active_mask) - 1; + // All lanes in a nonempty logical warp participate, including out-of-bounds lanes with zero gradients. + // Use the active mask only to select valid offsets so warp_uniform agrees across all lanes before Sum. + const int64_t common_offset = __shfl_sync(logical_lane_mask, b_offset, leader, kLogicalWarpSize); + + bool warp_uniform = true; + for (int i = 0; i < kLogicalWarpSize; ++i) { + if (!(logical_active_mask & (unsigned{1} << i))) { + continue; + } + const int64_t offset_i = __shfl_sync(logical_lane_mask, b_offset, i, kLogicalWarpSize); + if (offset_i != common_offset) { + warp_uniform = false; + break; + } + } + + if (warp_uniform) { + const float reduced = WarpReduce(temp_storage[logical_warp_id]).Sum(grad_val); + if (lane_id == leader) { + atomicAdd(&output_b[common_offset], common::cuda::Cast(reduced)); + } + } else if (in_bounds) { + atomicAdd(&output_b[b_offset], common::cuda::Cast(grad_val)); + } +} + // NOTE(dcj): Specialized BinaryBackwardKernel for low-precision types (__half / bfloat16) template __global__ void BinaryBackwardKernel(T *output_a, T *output_b, FuncA fn_a, FuncB fn_b, BroadcastMeta meta, @@ -759,6 +958,114 @@ void LaunchBackward(FuncA fun_a, FuncB fun_b, const std::shared_ptr &out } } +// Mixed-input backward launcher: grad_output and both gradient outputs are Tout (= float for every +// mixed floating combination), while the saved operands a/b keep their native Ta/Tb and are widened +// in registers. Mirrors the fast/broadcast split of the homogeneous LaunchBackward above. +template +void LaunchBackwardMixed(FuncA fun_a, FuncB fun_b, const std::shared_ptr &grad_a, + const std::shared_ptr &grad_b, const std::vector &a_dims, + const std::vector &b_dims, const std::shared_ptr &grad_output, + const std::shared_ptr &a, const std::shared_ptr &b) { + auto device = grad_a->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + + Tout *out_a_ptr = static_cast(grad_a->DataPtr()); + Tout *out_b_ptr = static_cast(grad_b->DataPtr()); + const Tout *grad_out_ptr = static_cast(grad_output->DataPtr()); + const Ta *a_ptr = a ? static_cast(a->DataPtr()) : nullptr; + const Tb *b_ptr = b ? static_cast(b->DataPtr()) : nullptr; + + const auto &out_dims = grad_output->Dims(); + const size_t num_elements = grad_output->NumElements(); + dim3 block_dims = ChooseBlockDims(num_elements); + + if (ShapesEqual(a_dims, b_dims) && ShapesEqual(a_dims, out_dims)) { + dim3 grid_dims(std::min(CEIL_DIV(num_elements, block_dims.x), static_cast(65535))); + BinaryBackwardKernelNoBroadcastFastMixed<<>>( + out_a_ptr, out_b_ptr, fun_a, fun_b, num_elements, grad_out_ptr, a_ptr, b_ptr); + return; + } + + BroadcastMeta meta = MakeBroadcastMeta(a_dims, b_dims, out_dims); + dim3 grid_dims(CEIL_DIV(num_elements, block_dims.x)); + const int block_threads = static_cast(block_dims.x); + const int num_warps = CEIL_DIV(block_threads, kLogicalWarpSize); + const size_t smem_size = num_warps * sizeof(cub::WarpReduce::TempStorage); + BinaryBackwardKernelMixed<<>>( + out_a_ptr, out_b_ptr, fun_a, fun_b, meta, num_elements, grad_out_ptr, a_ptr, b_ptr); +} + +// Dispatch the mixed-input backward for supported (Ta, Tb) combinations. Returns false (caller falls +// back to promote-via-To()) when: not actually mixed, operands are non-floating or non-contiguous. +// Fills grad_a/grad_b with zero only when a broadcast reduction (atomicAdd) will be used. +template +bool LaunchBinaryBackwardMixed(FuncA fn_a, FuncB fn_b, const std::shared_ptr &grad_a, + const std::shared_ptr &grad_b, const std::vector &a_dims, + const std::vector &b_dims, const std::shared_ptr &grad_output, + const std::shared_ptr &a, const std::shared_ptr &b, + bool needs_broadcast) { + const DataType a_dtype = a ? a->Dtype() : DataType::kFLOAT32; + const DataType b_dtype = b ? b->Dtype() : DataType::kFLOAT32; + auto is_float + = [](DataType d) { return d == DataType::kFLOAT32 || d == DataType::kBFLOAT16 || d == DataType::kFLOAT16; }; + const bool mixed = (a && a_dtype != DataType::kFLOAT32) || (b && b_dtype != DataType::kFLOAT32); + if (!mixed || !is_float(a_dtype) || !is_float(b_dtype)) { + return false; + } + if (!grad_output->IsContiguous() || (a && !a->IsContiguous()) || (b && !b->IsContiguous())) { + return false; + } + if (needs_broadcast) { + // grad_a needs no zero-init: a is never broadcast (one-way b->a only), so every kernel here + // writes output_a[a_offset] directly with full coverage (each element exactly once). Only + // grad_b is accumulated via atomicAdd and thus requires a zero start. + grad_b->Fill(0.0f); + } + switch (a_dtype) { + case DataType::kFLOAT32: + switch (b_dtype) { + case DataType::kBFLOAT16: + LaunchBackwardMixed(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, + b); + return true; + case DataType::kFLOAT16: + LaunchBackwardMixed(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, b); + return true; + default: + return false; + } + case DataType::kBFLOAT16: + switch (b_dtype) { + case DataType::kFLOAT32: + LaunchBackwardMixed(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, + b); + return true; + case DataType::kFLOAT16: + LaunchBackwardMixed(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, + b); + return true; + default: + return false; + } + case DataType::kFLOAT16: + switch (b_dtype) { + case DataType::kFLOAT32: + LaunchBackwardMixed(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, b); + return true; + case DataType::kBFLOAT16: + LaunchBackwardMixed(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, + b); + return true; + default: + return false; + } + default: + return false; + } +} + template std::shared_ptr UnaryForward(const std::shared_ptr &input, Func unary_fn) { auto dtype = input->Dtype(); auto output = std::make_shared(input->Dims(), dtype, input->GetDevice()); @@ -808,14 +1115,24 @@ std::shared_ptr BinaryForward(const std::shared_ptr &a, const st DataType promoted_type = PromoteDataTypes(a_dtype, b_dtype); - auto a_promoted = a_dtype == promoted_type ? a : std::make_shared(a->To(promoted_type)); - auto b_promoted = b_dtype == promoted_type ? b : std::make_shared(b->To(promoted_type)); // Currently a and b should have the same data type and only one-way broadcasting from b to a is assumed by // default CHECK(a->NumElements() >= b->NumElements() && a->NumElements() % b->NumElements() == 0); auto output = std::make_shared(a->Dims(), promoted_type, a->GetDevice()); + // Mixed-dtype fast path: widen operands to the promoted type *in registers* instead of + // materializing a Cast kernel. Bit-identical to the promote-via-To() path below (Cast is + // exactly the per-element conversion To(Tout) performs), but removes the bf16->f32 upcast kernels + // that dominate the timeline before RoPE Mul and the residual Add. Falls back to To() for + // combinations the mixed launcher does not cover. + if (a_dtype != b_dtype && LaunchBinaryForwardMixed(binary_fn, output, a, b, a_dtype, b_dtype, promoted_type)) { + return output; + } + + auto a_promoted = a_dtype == promoted_type ? a : std::make_shared(a->To(promoted_type)); + auto b_promoted = b_dtype == promoted_type ? b : std::make_shared(b->To(promoted_type)); + switch (promoted_type) { DISPATCH_CASE(WRAP(LaunchForward(binary_fn, output, a_promoted, b_promoted);), DataType::kFLOAT32) DISPATCH_CASE(WRAP(LaunchForward(binary_fn, output, a_promoted, b_promoted);), DataType::kBFLOAT16) @@ -849,6 +1166,22 @@ BinaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr CHECK(a_num_elements >= b_num_elements && a_num_elements % b_num_elements == 0); + auto grad_a = std::make_shared(a_dims, promoted_type, device); + auto grad_b = std::make_shared(b_dims, promoted_type, device); + + // Only Fill(0) when broadcast is needed (atomicAdd requires zero-init). + // The no-broadcast fast path writes every element directly. + const bool needs_broadcast = !ShapesEqual(a_dims, b_dims) || !ShapesEqual(a_dims, grad_output->Dims()); + + // Mixed-dtype fast path (mirror of BinaryForward): promoted compute/output type is float32, + // grad_output is already float32, and the saved operands are contiguous floats of possibly + // narrower dtype. Widens in registers instead of materializing Cast kernels, and must run BEFORE + // any promote-via-To() below so the bf16->f32 upcast is never launched. Uses the ORIGINAL a/b. + if (promoted_type == DataType::kFLOAT32 && dtype == DataType::kFLOAT32 + && LaunchBinaryBackwardMixed(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, b, needs_broadcast)) { + return {grad_a, grad_b}; + } + auto promote_if_needed = [&](std::shared_ptr &t, size_t expected_numel, DataType promoted_type) { if (t) { CHECK(expected_numel == t->NumElements()); @@ -863,17 +1196,11 @@ BinaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr grad_output_promoted = std::make_shared(grad_output_promoted->To(promoted_type)); } - auto grad_a = std::make_shared(a_dims, promoted_type, device); - auto grad_b = std::make_shared(b_dims, promoted_type, device); - - // Only Fill(0) when broadcast is needed (atomicAdd requires zero-init). - // The no-broadcast fast path writes every element directly. - const bool needs_broadcast = !ShapesEqual(a_dims, b_dims) || !ShapesEqual(a_dims, grad_output->Dims()); - switch (promoted_type) { DISPATCH_CASE(WRAP({ if (needs_broadcast) { - grad_a->Fill(0.0f); + // grad_a is written directly & fully covered (a never broadcast); only + // grad_b needs zero-init for atomicAdd accumulation. grad_b->Fill(0.0f); } LaunchBackward(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output_promoted, @@ -882,7 +1209,8 @@ BinaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr DataType::kFLOAT32) DISPATCH_CASE(WRAP({ if (needs_broadcast) { - grad_a->Fill(0.0f); + // grad_a is written directly & fully covered (a never broadcast); only + // grad_b needs zero-init for atomicAdd accumulation. grad_b->Fill(0.0f); } LaunchBackward(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output_promoted, diff --git a/infini_train/src/kernels/cuda/embedding.cu b/infini_train/src/kernels/cuda/embedding.cu index 89361e03a..4d83c10ce 100644 --- a/infini_train/src/kernels/cuda/embedding.cu +++ b/infini_train/src/kernels/cuda/embedding.cu @@ -3,6 +3,7 @@ #include "infini_train/include/common/cuda/common_cuda.h" #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/dispatcher.h" +#include "infini_train/include/sparse_row_grad.h" #include "infini_train/include/tensor.h" #include "infini_train/src/core/runtime/cuda/cuda_dispatch.h" @@ -81,38 +82,151 @@ __global__ void EmbeddingBackwardKernel(const int64_t *input_ptr, const T *grad_ } } -std::shared_ptr EmbeddingBackward(const std::shared_ptr &input, const std::vector &weight_dims, +// Sparse backward, part 1: scatter-add the per-token grad rows into the persistent gradient +// buffer and claim (deduplicate) every hit row in row_list, so later passes can touch exactly the +// rows that are dirty instead of the whole [vocab, dim] tensor. +// grid = (num_tokens, ceil(dim / blockDim.x)): block (t, c) covers the c-th chunk of token t. The +// first thread of every c == 0 block claims row tokens[t] with a CAS against the stamp array, so +// each row lands in row_list exactly once per accumulation cycle (generation); every thread then +// scatter-adds its element. Claim and scatter write disjoint memory and both complete before any +// later kernel (Adam, clear) observes them. +template +__global__ void EmbeddingBackwardSparseKernel(const int64_t *__restrict__ tokens, const T *__restrict__ grad_output, + T *__restrict__ grad_buffer, int32_t *__restrict__ stamp, + int32_t *__restrict__ row_list, int32_t *__restrict__ count, int gen, + int vocab_size, int embedding_dim) { + const int64_t token = tokens[blockIdx.x]; + if (token < 0 || token >= vocab_size) { + return; + } + const int row = static_cast(token); + + if (blockIdx.y == 0 && threadIdx.x == 0) { + int old = stamp[row]; + while (old != gen) { + const int prev = atomicCAS(&stamp[row], old, gen); + if (prev == old) { + row_list[atomicAdd(count, 1)] = row; + break; + } + old = prev; + } + } + + const int elem = blockIdx.y * blockDim.x + threadIdx.x; + if (elem < embedding_dim) { + atomicAdd(&grad_buffer[static_cast(row) * embedding_dim + elem], + grad_output[static_cast(blockIdx.x) * embedding_dim + elem]); + } +} + +// Sparse backward, part 2 (driven by the optimizer): zero exactly the rows claimed since the last +// clear. The row count lives on the device, so the host never needs to synchronize to size this. +template +__global__ void SparseRowClearRowsKernel(T *__restrict__ grad_buffer, const int32_t *__restrict__ row_list, + const int32_t *__restrict__ count, int embedding_dim) { + const int num_rows = *count; + for (int r = blockIdx.x; r < num_rows; r += gridDim.x) { + T *row_base = grad_buffer + static_cast(row_list[r]) * embedding_dim; + for (int e = threadIdx.x; e < embedding_dim; e += blockDim.x) { row_base[e] = static_cast(0.0f); } + } +} + +__global__ void SparseRowResetCountKernel(int32_t *count) { *count = 0; } + +void SparseRowClearRows(const std::shared_ptr &grad_buffer, const std::shared_ptr &row_list, + const std::shared_ptr &count) { + auto device = grad_buffer->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + const int embedding_dim = static_cast(grad_buffer->Dims()[1]); + + constexpr int kThreadsPerBlock = 256; + constexpr int kNumBlocks = 1024; // grid-strided; idle blocks only read *count and exit + core::cuda::DispatchCudaFunc( + grad_buffer->Dtype(), + [=]() { + SparseRowClearRowsKernel<<>>( + static_cast(grad_buffer->DataPtr()), static_cast(row_list->DataPtr()), + static_cast(count->DataPtr()), embedding_dim); + }, + "CUDA SparseRowClearRows"); +} + +void SparseRowResetCount(const std::shared_ptr &count) { + auto device = count->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + SparseRowResetCountKernel<<<1, 1, 0, cuda_stream>>>(static_cast(count->DataPtr())); +} + +std::shared_ptr EmbeddingBackward(const std::shared_ptr &input, const std::shared_ptr &weight, const std::shared_ptr &grad_output) { CHECK(input->Dtype() == DataType::kINT64); - CHECK_EQ(weight_dims.size(), 2); + CHECK_EQ(weight->Dims().size(), 2); auto device = input->GetDevice(); const auto &cuda_stream = dynamic_cast( infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) ->cuda_stream(); - const int vocab_size = weight_dims[0]; - const int embedding_dim = weight_dims[1]; + const int64_t vocab_size = weight->Dims()[0]; + const int64_t embedding_dim = weight->Dims()[1]; CHECK_EQ(input->Dims().size() + 1, grad_output->Dims().size()); for (int idx = 0; idx < input->Dims().size(); ++idx) { CHECK_EQ(input->Dims()[idx], grad_output->Dims()[idx]); } CHECK_EQ(*grad_output->Dims().rbegin(), embedding_dim); + const int64_t num_tokens = input->NumElements(); + + // Dtype mismatch (a grad produced outside autocast against a cast weight): the sparse buffer + // lives in the weight's dtype, so fall back to the legacy dense grad for that rare case. + if (grad_output->Dtype() != weight->Dtype()) { + auto grad_weight = std::make_shared(weight->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + const int threads_per_block = 256; + const int num_blocks = (num_tokens + threads_per_block - 1) / threads_per_block; + core::cuda::DispatchCudaFunc( + grad_output->Dtype(), + [=]() { + grad_weight->Fill(0.0); + EmbeddingBackwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), + static_cast(grad_weight->DataPtr()), static_cast(num_tokens), + static_cast(embedding_dim), static_cast(vocab_size)); + }, + "CUDA EmbeddingBackward"); + return grad_weight; + } - auto dtype = grad_output->Dtype(); - auto grad_weight = std::make_shared(weight_dims, dtype, grad_output->GetDevice()); - const int num_tokens = input->NumElements(); - const int threads_per_block = 256; - const int num_blocks = (num_tokens + threads_per_block - 1) / threads_per_block; + auto *state = SparseRowGradRegistry::Instance().GetOrCreate(weight); + if (!state->initialized) { + // One-time init: the "zero everywhere except rows claimed since the last clear" invariant + // has to start from an actually zero buffer. Everything afterwards only touches dirty rows. + state->grad_buffer->Fill(0.0); + cudaMemsetAsync(state->stamp->DataPtr(), 0, state->stamp->NumElements() * sizeof(int32_t), cuda_stream); + cudaMemsetAsync(state->count->DataPtr(), 0, sizeof(int32_t), cuda_stream); + state->initialized = true; + } - core::cuda::DispatchCudaFunc( - dtype, - [=]() { - grad_weight->Fill(0.0); - EmbeddingBackwardKernel<<>>( - static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), - static_cast(grad_weight->DataPtr()), num_tokens, embedding_dim, vocab_size); - }, - "CUDA EmbeddingBackward"); + if (num_tokens > 0) { + constexpr int kThreadsPerBlock = 256; + const dim3 grid(static_cast(num_tokens), + static_cast((embedding_dim + kThreadsPerBlock - 1) / kThreadsPerBlock)); + core::cuda::DispatchCudaFunc( + weight->Dtype(), + [=]() { + EmbeddingBackwardSparseKernel<<>>( + static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), + static_cast(state->grad_buffer->DataPtr()), static_cast(state->stamp->DataPtr()), + static_cast(state->row_list->DataPtr()), static_cast(state->count->DataPtr()), + state->generation, static_cast(vocab_size), static_cast(embedding_dim)); + }, + "CUDA EmbeddingBackward"); + } - return grad_weight; + // A view sharing the persistent buffer: AccumulateGrad notices that the grad storage is the + // scattered buffer and skips its own accumulation, so these rows are what ends up in + // param->grad() — with no per-call allocation or fill of a vocab-sized tensor. + return std::make_shared(*state->grad_buffer.get(), 0, state->grad_buffer->Dims()); } } // namespace infini_train::kernels::cuda @@ -121,5 +235,7 @@ std::shared_ptr EmbeddingBackward(const std::shared_ptr &input, REGISTER_CUDA_EMBEDDING_KERNEL(EmbeddingForward) REGISTER_CUDA_EMBEDDING_KERNEL(EmbeddingBackward) +REGISTER_CUDA_EMBEDDING_KERNEL(SparseRowClearRows) +REGISTER_CUDA_EMBEDDING_KERNEL(SparseRowResetCount) #undef REGISTER_CUDA_EMBEDDING_KERNEL diff --git a/infini_train/src/kernels/cuda/fill.cu b/infini_train/src/kernels/cuda/fill.cu index 3ddead5cf..d64ed9f4b 100644 --- a/infini_train/src/kernels/cuda/fill.cu +++ b/infini_train/src/kernels/cuda/fill.cu @@ -1,9 +1,13 @@ #include +#include +#include #include #include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/datatype.h" #include "infini_train/include/device.h" #include "infini_train/include/dispatcher.h" +#include "infini_train/include/scalar.h" #include "infini_train/include/tensor.h" #include "infini_train/src/core/runtime/cuda/cuda_dispatch.h" @@ -18,16 +22,65 @@ template __global__ void FillKernel(T *data, T value, size_t size) } } +// A Scalar is memset-compatible iff its stored bit pattern is all zeros, which for every +// dtype InfiniTrain supports (bool / intN / uintN / fp16 / bf16 / fp32 / fp64) coincides +// with the numeric value +0. Note that -0.0 has the sign bit set and is NOT all-zero bits, +// so we deliberately reject it here even though it compares == 0.0 numerically; callers +// that pass -0.0 fall through to FillKernel and get the correct bit pattern. +static bool IsZeroBitsScalar(const Scalar &s) { + switch (s.kind) { + case Scalar::Kind::kBool: + case Scalar::Kind::kUInt64: + return s.u == 0; + case Scalar::Kind::kInt64: + return s.i == 0; + case Scalar::Kind::kDouble: { + uint64_t bits = 0; + std::memcpy(&bits, &s.d, sizeof(bits)); + return bits == 0; + } + default: + return false; + } +} + // TODO(dcj): refactor Fill kernel with elementwise template void Fill(std::shared_ptr tensor, Scalar scalar) { - const int num_tokens = tensor->NumElements(); - const int threads_per_block = 256; - const int num_blocks = (num_tokens + threads_per_block - 1) / threads_per_block; + const size_t num_elements = tensor->NumElements(); + if (num_elements == 0) { + return; + } + auto device = tensor->GetDevice(); const auto &cuda_stream = dynamic_cast( infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) ->cuda_stream(); + // Fast path: Fill(0) on a contiguous tensor is a pure byte-zeroing operation. Route it to + // cudaMemsetAsync, which runs on the copy/DMA engine and does not consume an SM kernel + // launch slot. On llama3.2-1B this covers ~500 launches/step (ZeroGrad + scatter backward + // zero-init) that were previously 3-4 us FillKernels each - see docs/kernel 优化.md §3.3. + // + // Preconditions: + // 1. scalar is bit-exactly zero (see IsZeroBitsScalar; rejects -0.0 which has a nonzero + // sign bit and would produce a different bit pattern than memset). + // 2. tensor is contiguous (IsContiguous() is currently unconditionally true in Tensor; + // the check is kept for when strided views land). + // 3. dtype size is known (kDataTypeToSize covers every enum value). + // Any of these failing -> fall through to FillKernel, which handles the general case. + if (IsZeroBitsScalar(scalar) && tensor->IsContiguous()) { + auto size_it = kDataTypeToSize.find(tensor->Dtype()); + if (size_it != kDataTypeToSize.end()) { + const size_t bytes = num_elements * size_it->second; + cudaMemsetAsync(tensor->DataPtr(), 0, bytes, cuda_stream); + return; + } + } + + const int num_tokens = static_cast(num_elements); + const int threads_per_block = 256; + const int num_blocks = (num_tokens + threads_per_block - 1) / threads_per_block; + core::cuda::DispatchCudaFunc( tensor->Dtype(), [=]() { diff --git a/infini_train/src/kernels/cuda/rmsnorm.cu b/infini_train/src/kernels/cuda/rmsnorm.cu new file mode 100644 index 000000000..ece5914dc --- /dev/null +++ b/infini_train/src/kernels/cuda/rmsnorm.cu @@ -0,0 +1,184 @@ +#include + +#include "infini_train/include/common/cuda/common_cuda.h" +#include "infini_train/include/common/cuda/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/cuda/cuda_dispatch.h" +#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" + +namespace infini_train::kernels::cuda { + +template +__global__ void RMSNormForwardKernel(const T *__restrict__ input, const T *__restrict__ weight, float eps, + T *__restrict__ output, float *__restrict__ rstd_out, int embed_dim) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage_rstd; + __shared__ float shared_rstd; + + const int token_idx = blockIdx.x; + const T *x = input + token_idx * embed_dim; + T *y = output + token_idx * embed_dim; + + float sqsum = 0.0f; + + for (int i = threadIdx.x; i < embed_dim; i += BLOCK_SIZE) { + float val = common::cuda::Cast(x[i]); + sqsum += val * val; + } + + float total_sqsum = BlockReduce(temp_storage_rstd).Sum(sqsum); + + if (threadIdx.x == 0) { + float var = total_sqsum / embed_dim; + float rstd = rsqrtf(var + eps); + shared_rstd = rstd; + if (rstd_out) { + rstd_out[token_idx] = rstd; + } + } + __syncthreads(); + + for (int i = threadIdx.x; i < embed_dim; i += BLOCK_SIZE) { + // Keep the two multiplications separate: normalize by rstd first, then scale by weight. + // A fused x * (rstd * weight) would change the rounding order of the composite path. + float norm = common::cuda::Cast(x[i]) * shared_rstd; + y[i] = common::cuda::Cast(norm * common::cuda::Cast(weight[i])); + } +} + +std::tuple, std::shared_ptr> +RMSNormForward(const std::shared_ptr &input, const std::shared_ptr &weight, const float eps) { + /* + x: [..., embed_dim] + -> RMSNorm (w: [embed_dim]) + -> o: [..., embed_dim] + */ + // The composite path (Mean(-1)/Pow/Rsqrt/Mul) supports any rank, so the fused kernel keeps the + // same generality: one block per leading-index row, reducing over the last dimension. + CHECK_GE(input->Dims().size(), 2); + CHECK_EQ(input->Dims().back(), weight->Dims()[0]); + auto input_c = input->IsContiguous() ? input : input->Contiguous(); + + const int embed_dim = static_cast(input_c->Dims().back()); + const int64_t rows = input_c->NumElements() / embed_dim; + + auto dtype = input_c->Dtype(); + CHECK(dtype == weight->Dtype()); + + auto output = std::make_shared(input_c->Dims(), dtype, input_c->GetDevice()); + auto rstd = std::make_shared(std::vector(input_c->Dims().begin(), input_c->Dims().end() - 1), + DataType::kFLOAT32, input_c->GetDevice()); + + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = static_cast(rows); + + auto device = input_c->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + + core::cuda::DispatchCudaFunc( + dtype, + [=]() { + // Each token block writes its rstd exactly once; no Fill is needed. + RMSNormForwardKernel<<>>( + static_cast(input_c->DataPtr()), static_cast(weight->DataPtr()), eps, + static_cast(output->DataPtr()), static_cast(rstd->DataPtr()), embed_dim); + }, + "CUDA RMSNormForward"); + + return {output, rstd}; +} + +template +__global__ void RMSNormBackwardKernel(const T *__restrict__ input, const T *__restrict__ grad_output, + const T *__restrict__ weight, const float *__restrict__ rstd, + T *__restrict__ grad_input, T *__restrict__ grad_weight, int embed_dim, + size_t weight_num_elements) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + __shared__ float shared_K; + + const int token_idx = blockIdx.x; + const T *x = input + token_idx * embed_dim; + const T *g_out = grad_output + token_idx * embed_dim; + T *g_in = grad_input + token_idx * embed_dim; + const float rstd_val = rstd[token_idx]; + + // S1 = sum_i(g_i * w_i * x_i); K = (S1 / H) * rstd, shared by the whole row. + float S1 = 0.0f; + for (int i = threadIdx.x; i < embed_dim; i += BLOCK_SIZE) { + float xv = common::cuda::Cast(x[i]); + float wv = common::cuda::Cast(weight[i]); + float go = common::cuda::Cast(g_out[i]); + S1 += go * wv * xv; + } + S1 = BlockReduce(temp_storage).Sum(S1); + if (threadIdx.x == 0) { + shared_K = (S1 / embed_dim) * rstd_val; + } + __syncthreads(); + + for (int i = threadIdx.x; i < embed_dim; i += BLOCK_SIZE) { + float xv = common::cuda::Cast(x[i]); + float wv = common::cuda::Cast(weight[i]); + float go = common::cuda::Cast(g_out[i]); + float norm = xv * rstd_val; + g_in[i] = common::cuda::Cast((go * wv - shared_K * norm) * rstd_val); + common::cuda::fastAtomicAdd(grad_weight, i, weight_num_elements, common::cuda::Cast(go * norm), + true); + } +} + +std::tuple, std::shared_ptr> +RMSNormBackward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &rstd, const std::shared_ptr &grad_output) { + auto input_c = input->IsContiguous() ? input : input->Contiguous(); + + const int embed_dim = static_cast(input_c->Dims().back()); + const int64_t rows = input_c->NumElements() / embed_dim; + + auto dtype = input_c->Dtype(); + CHECK(dtype == weight->Dtype() && dtype == grad_output->Dtype() && rstd->Dtype() == DataType::kFLOAT32); + + auto grad_input = std::make_shared(input_c->Dims(), dtype, grad_output->GetDevice()); + auto grad_weight = std::make_shared(weight->Dims(), dtype, grad_output->GetDevice()); + + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = static_cast(rows); + + auto device = input_c->GetDevice(); + const auto &cuda_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->cuda_stream(); + core::cuda::DispatchCudaFunc( + dtype, + [=]() { + // grad_weight accumulates across token blocks via atomics and must start at zero; + // each token block fully overwrites its own grad_input slice, so no Fill is needed there. + grad_weight->Fill(0.0); + RMSNormBackwardKernel<<>>( + static_cast(input_c->DataPtr()), static_cast(grad_output->DataPtr()), + static_cast(weight->DataPtr()), static_cast(rstd->DataPtr()), + static_cast(grad_input->DataPtr()), static_cast(grad_weight->DataPtr()), embed_dim, + grad_weight->NumElements()); + }, + "CUDA RMSNormBackward"); + + return {grad_input, grad_weight}; +} +} // namespace infini_train::kernels::cuda + +#define REGISTER_CUDA_RMSNORM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name) + +REGISTER_CUDA_RMSNORM_KERNEL(RMSNormForward) +REGISTER_CUDA_RMSNORM_KERNEL(RMSNormBackward) + +#undef REGISTER_CUDA_RMSNORM_KERNEL diff --git a/infini_train/src/kernels/cuda/slice.cu b/infini_train/src/kernels/cuda/slice.cu index f71b4be61..f57d63fca 100644 --- a/infini_train/src/kernels/cuda/slice.cu +++ b/infini_train/src/kernels/cuda/slice.cu @@ -1,3 +1,4 @@ +#include #include #include @@ -12,11 +13,18 @@ #include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h" namespace infini_train::kernels::cuda { +constexpr int kMaxDims = 8; +struct SliceMeta { + int64_t new_dims[kMaxDims]; + int64_t starts[kMaxDims]; + int64_t steps[kMaxDims]; + int64_t in_strides[kMaxDims]; + int64_t out_strides[kMaxDims]; +}; template -__global__ void SliceForwardKernel(const T *input, T *output, const int64_t *new_dims, const int64_t *starts, - const int64_t *steps, const int64_t *in_strides, const int64_t *out_strides, - int num_dims, int64_t total_elements) { +__global__ void SliceForwardKernel(const T *input, T *output, const SliceMeta meta, int num_dims, + int64_t total_elements) { int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x; if (out_idx >= total_elements) { return; @@ -24,8 +32,8 @@ __global__ void SliceForwardKernel(const T *input, T *output, const int64_t *new int64_t in_index = 0; for (int i = 0; i < num_dims; ++i) { - int64_t idx = (out_idx / out_strides[i]) % new_dims[i]; - in_index += (starts[i] + idx * steps[i]) * in_strides[i]; + int64_t idx = (out_idx / meta.out_strides[i]) % meta.new_dims[i]; + in_index += (meta.starts[i] + idx * meta.steps[i]) * meta.in_strides[i]; } output[out_idx] = input[in_index]; @@ -38,6 +46,7 @@ std::shared_ptr SliceForward(const std::shared_ptr &input, const auto &dims = input->Dims(); CHECK_EQ(starts.size(), dims.size()); const int64_t num_dims = dims.size(); + CHECK_LE(num_dims, kMaxDims); std::vector new_dims; for (int i = 0; i < starts.size(); ++i) { @@ -65,28 +74,19 @@ std::shared_ptr SliceForward(const std::shared_ptr &input, const int64_t total_elements = stride; - int64_t *new_dims_dev, *starts_dev, *steps_dev, *input_strides_dev, *output_strides_dev; - auto device = input->GetDevice(); const auto &stream = dynamic_cast( infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) ->cuda_stream(); - cudaMallocAsync(&new_dims_dev, - (ends.size() + starts.size() + steps.size() + dims.size() + new_dims.size()) * sizeof(int64_t), - stream); - starts_dev = new_dims_dev + ends.size(); - steps_dev = starts_dev + starts.size(); - input_strides_dev = steps_dev + steps.size(); - output_strides_dev = input_strides_dev + dims.size(); - - cudaMemcpyAsync(new_dims_dev, new_dims.data(), ends.size() * sizeof(int64_t), cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(starts_dev, starts.data(), starts.size() * sizeof(int64_t), cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(steps_dev, steps.data(), steps.size() * sizeof(int64_t), cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(input_strides_dev, src_strides.data(), dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, - stream); - cudaMemcpyAsync(output_strides_dev, dst_strides.data(), new_dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, - stream); + // Metadata (5 arrays x num_dims int64) is passed by value through kernel parameter space + // (constant cache), so no device buffer / H2D memcpy is needed. + SliceMeta meta{}; + std::copy(new_dims.begin(), new_dims.end(), meta.new_dims); + std::copy(starts.begin(), starts.end(), meta.starts); + std::copy(steps.begin(), steps.end(), meta.steps); + std::copy(src_strides.begin(), src_strides.end(), meta.in_strides); + std::copy(dst_strides.begin(), dst_strides.end(), meta.out_strides); int threads_per_block = 256; int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; @@ -94,21 +94,18 @@ std::shared_ptr SliceForward(const std::shared_ptr &input, const core::cuda::DispatchCudaFunc( dtype, [=]() { - SliceForwardKernel<<>>( - static_cast(input->DataPtr()), static_cast(new_tensor->DataPtr()), new_dims_dev, - starts_dev, steps_dev, input_strides_dev, output_strides_dev, num_dims, total_elements); + SliceForwardKernel<<>>(static_cast(input->DataPtr()), + static_cast(new_tensor->DataPtr()), + meta, num_dims, total_elements); }, "CUDA SliceForward"); - cudaFreeAsync(new_dims_dev, stream); - return new_tensor; } template -__global__ void SliceBackwardKernel(const T *grad_output, T *grad_input, const int64_t *new_dims, const int64_t *starts, - const int64_t *steps, const int64_t *in_strides, const int64_t *out_strides, - int num_dims, int64_t total_elements) { +__global__ void SliceBackwardKernel(const T *grad_output, T *grad_input, const SliceMeta meta, int num_dims, + int64_t total_elements) { int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x; if (out_idx >= total_elements) { return; @@ -116,8 +113,8 @@ __global__ void SliceBackwardKernel(const T *grad_output, T *grad_input, const i int64_t in_index = 0; for (int i = 0; i < num_dims; ++i) { - int64_t idx = (out_idx / out_strides[i]) % new_dims[i]; - in_index += (starts[i] + idx * steps[i]) * in_strides[i]; + int64_t idx = (out_idx / meta.out_strides[i]) % meta.new_dims[i]; + in_index += (meta.starts[i] + idx * meta.steps[i]) * meta.in_strides[i]; } grad_input[in_index] = grad_output[out_idx]; } @@ -130,6 +127,7 @@ std::shared_ptr SliceBackward(const std::shared_ptr &grad_output auto &dims = input->Dims(); CHECK_EQ(starts.size(), dims.size()); const int64_t num_dims = dims.size(); + CHECK_LE(num_dims, kMaxDims); std::vector new_dims; for (int i = 0; i < starts.size(); ++i) { @@ -158,28 +156,18 @@ std::shared_ptr SliceBackward(const std::shared_ptr &grad_output int64_t total_elements = stride; - int dims_size = dims.size(); - int64_t *new_dims_dev, *starts_dev, *steps_dev, *input_strides_dev, *output_strides_dev; - auto device = input->GetDevice(); const auto &stream = dynamic_cast( infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) ->cuda_stream(); - cudaMallocAsync(&new_dims_dev, - (ends.size() + starts.size() + steps.size() + dims.size() + new_dims.size()) * sizeof(int64_t), - stream); - starts_dev = new_dims_dev + ends.size(); - steps_dev = starts_dev + starts.size(); - input_strides_dev = steps_dev + steps.size(); - output_strides_dev = input_strides_dev + dims.size(); - - cudaMemcpyAsync(new_dims_dev, new_dims.data(), ends.size() * sizeof(int64_t), cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(starts_dev, starts.data(), starts.size() * sizeof(int64_t), cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(steps_dev, steps.data(), steps.size() * sizeof(int64_t), cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(input_strides_dev, src_strides.data(), dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, - stream); - cudaMemcpyAsync(output_strides_dev, dst_strides.data(), new_dims.size() * sizeof(int64_t), cudaMemcpyHostToDevice, - stream); + + // Metadata is passed by value through kernel parameter space; no device buffer / H2D memcpy. + SliceMeta meta{}; + std::copy(new_dims.begin(), new_dims.end(), meta.new_dims); + std::copy(starts.begin(), starts.end(), meta.starts); + std::copy(steps.begin(), steps.end(), meta.steps); + std::copy(src_strides.begin(), src_strides.end(), meta.in_strides); + std::copy(dst_strides.begin(), dst_strides.end(), meta.out_strides); int threads_per_block = 256; int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; @@ -188,13 +176,11 @@ std::shared_ptr SliceBackward(const std::shared_ptr &grad_output grad_output_dtype, [=]() { SliceBackwardKernel<<>>( - static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), new_dims_dev, - starts_dev, steps_dev, input_strides_dev, output_strides_dev, num_dims, total_elements); + static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), meta, num_dims, + total_elements); }, "CUDA SliceBackward"); - cudaFreeAsync(new_dims_dev, stream); - return grad_input; } } // namespace infini_train::kernels::cuda diff --git a/infini_train/src/kernels/cuda/transform.cu b/infini_train/src/kernels/cuda/transform.cu index 81ff663c5..e7a43b0ba 100644 --- a/infini_train/src/kernels/cuda/transform.cu +++ b/infini_train/src/kernels/cuda/transform.cu @@ -93,7 +93,8 @@ std::shared_ptr TrilBackward(const std::shared_ptr &grad_output, core::cuda::DispatchCudaFunc( dtype, [=]() { - grad_input->Fill(0.0); + // No Fill(0) needed: TrilBackwardKernel writes every element of grad_input exactly once + // (in-region -> grad_output[idx]; out-of-region -> T(0)). Grid covers [0, rows*cols). TrilBackwardKernel<<>>( static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), rows, cols, diagonal); @@ -180,7 +181,8 @@ std::shared_ptr TriuBackward(const std::shared_ptr &grad_output, core::cuda::DispatchCudaFunc( dtype, [=]() { - grad_input->Fill(0.0); + // No Fill(0) needed: TriuBackwardKernel writes every element of grad_input exactly once + // (in-region -> grad_output[idx]; out-of-region -> T(0)). Grid covers [0, rows*cols). TriuBackwardKernel<<>>( static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), rows, cols, diagonal); @@ -272,7 +274,9 @@ std::shared_ptr TransposeForward(const std::shared_ptr &input, i core::cuda::DispatchCudaFunc( dtype, [=]() { - output->Fill(0.0); + // No Fill(0) needed: TransposeForwardKernel writes output[idx] = input[in_flat_idx] for every + // idx in [0, num_elements); num_blocks = ceil(num_elements / threads_per_block) covers all, + // each element written exactly once, no atomicAdd. Fill was pure dead launch. TransposeForwardKernel<<>>( static_cast(input->DataPtr()), static_cast(output->DataPtr()), in_dims_dev, in_strides_dev, out_strides_dev, ndim, dim0, dim1, num_elements); @@ -438,7 +442,9 @@ std::shared_ptr MaskBackward(const std::shared_ptr &grad_output, core::cuda::DispatchCudaFunc( dtype, [=]() { - grad_input->Fill(0.0); + // No Fill(0) needed: MaskLeadsBackwardKernel writes grad_input[i] for every i in + // [0, rows*inner) = [0, grad_output->NumElements()); mask-hit lanes are written as T(0) + // by the kernel itself, so no zero-init is required. MaskLeadsBackwardKernel<<>>( static_cast(grad_output->DataPtr()), static_cast(mask_casted->DataPtr()), static_cast(grad_input->DataPtr()), rows, inner); @@ -452,7 +458,9 @@ std::shared_ptr MaskBackward(const std::shared_ptr &grad_output, core::cuda::DispatchCudaFunc( dtype, [=]() { - grad_input->Fill(0.0); + // No Fill(0) needed: MaskBackwardKernel writes grad_input[i] for every i in + // [0, batch_size*mask_size) = [0, grad_output->NumElements()); mask-hit lanes are + // written as T(0) by the kernel itself, so no zero-init is required. MaskBackwardKernel<<>>( static_cast(grad_output->DataPtr()), static_cast(mask_casted->DataPtr()), static_cast(grad_input->DataPtr()), static_cast(batch_size), static_cast(mask_size)); @@ -565,7 +573,10 @@ std::shared_ptr RepeatInterleaveBackward(const std::shared_ptr & core::cuda::DispatchCudaFunc( grad_output->Dtype(), [=]() { - grad_input->Fill(0.0); + // No Fill(0) needed: RepeatInterleaveBackwardKernel is a gather-reduce (not scatter): + // each thread owns one grad_input[idx], sequentially sums `repeat` grad_output entries + // in registers, then writes grad_input[idx] = sum exactly once. No atomicAdd, so the + // buffer does not need to start at zero. RepeatInterleaveBackwardKernel<<>>( static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), outer, dim_size, inner, repeat); diff --git a/infini_train/src/nn/modules/normalization.cc b/infini_train/src/nn/modules/normalization.cc index 388b04de5..0edf34b19 100644 --- a/infini_train/src/nn/modules/normalization.cc +++ b/infini_train/src/nn/modules/normalization.cc @@ -5,7 +5,6 @@ #include "infini_train/include/autograd/normalization.h" #include "infini_train/include/device.h" -#include "infini_train/include/nn/functional.h" #include "infini_train/include/nn/init.h" #include "infini_train/include/tensor.h" @@ -39,8 +38,7 @@ RMSNorm::RMSNorm(int64_t dim, float eps, Device device) : CloneableModule(kType) } std::vector> RMSNorm::Forward(const std::vector> &x) { - // broadcasted Mul([4, 64, 2048] * [4, 64, 1]) - auto norm = x[0] * function::Rsqrt(function::Mean(function::Pow(x[0], 2), -1, true) + eps_); - return {norm * parameters_[kParamWeightName]}; + auto outputs = std::make_shared(eps_)->Apply({x[0], parameters_[kParamWeightName]}); + return {outputs[0]}; } } // namespace infini_train::nn diff --git a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc index 523bcf2d7..60e6060cc 100644 --- a/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc +++ b/infini_train/src/nn/parallel/ddp/distributed_optimizer.cc @@ -170,6 +170,12 @@ float DistributedOptimizer::learning_rate() const { return Optimizer::learning_rate(); } +void DistributedOptimizer::EnableShadowWeights(DataType /*shadow_dtype*/) { + LOG(WARNING) << "DistributedOptimizer: shadow weights are not supported in distributed mode " + << "(base optimizer manages sharded params while autocast sees full params); " + << "falling back to the regular cast path."; +} + void DistributedOptimizer::Step() { // 1. Ensure grads are synced FinishGradSync(); diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc index 39b999c77..99bf5c605 100644 --- a/infini_train/src/optimizer.cc +++ b/infini_train/src/optimizer.cc @@ -6,9 +6,19 @@ #include "infini_train/include/core/runtime/device_guard.h" #include "infini_train/include/device.h" #include "infini_train/include/dispatcher.h" +#include "infini_train/include/sparse_row_grad.h" #include "infini_train/include/tensor.h" namespace infini_train { +thread_local std::unordered_map> g_shadow_registry; + +// Free function for autocast.h to query shadow weights (decoupled from the concrete Optimizer type). +// Returns the shadow on a hit; nullptr on a miss (activations, or shadow disabled), letting autocast fall back to Cast. +std::shared_ptr GetShadow(const Tensor *param) { + auto it = g_shadow_registry.find(param); + return it != g_shadow_registry.end() ? it->second : nullptr; +} + Optimizer::Optimizer(const std::vector> ¶ms, float learning_rate) : params_(params), learning_rate_(learning_rate) {} @@ -26,7 +36,25 @@ Optimizer::Optimizer(const NamedParameterList &named_params, float learning_rate } void Optimizer::ZeroGrad(bool set_to_none) { - for (auto param : params_) { param->ZeroGrad(set_to_none); } + for (auto param : params_) { + auto *sparse_state = GetSparseRowGradState(param.get()); + if (sparse_state) { + auto device = param->GetDevice(); + core::DeviceGuard guard(device); + // Zero only the rows made dirty since the last clear: the persistent buffer is zero + // everywhere else by construction, so no full-size memset is needed. This also bumps + // the claim generation, re-arming the dedup for the next accumulation cycle. + ClearSparseRowGradRows(sparse_state); + // Non-poisoned: the cleared buffer *is* the accumulator, so nothing else has to be + // reset. Poisoned: the live storage is some dense foreign buffer and needs the dense + // reset path (grad_.reset() or a full fill). + if (set_to_none || sparse_state->poisoned) { + param->ZeroGrad(set_to_none); + } + continue; + } + param->ZeroGrad(set_to_none); + } } void Optimizer::set_learning_rate(float lr) { learning_rate_ = lr; } @@ -99,6 +127,55 @@ Adam::Adam(const NamedParameterList &named_params, float learning_rate, float be } } +void Adam::EnableShadowWeights(DataType shadow_dtype) { + shadow_enable_ = true; + shadow_dtype_ = shadow_dtype; + shadow_weights_.clear(); + shadow_weights_.reserve(params_.size()); + // init shadow weights form param + for (auto ¶m : params_) { + auto shadow_weight = std::make_shared(param->Dims(), shadow_dtype_, param->GetDevice()); + auto casted = param->To(shadow_dtype); + shadow_weight->CopyFrom(casted); + shadow_weights_.push_back(shadow_weight); + g_shadow_registry[param.get()] = shadow_weight; + } + LOG(INFO) << "Enable shadow weights for Adam optimizer, shadow dtype: " << static_cast(shadow_dtype_); +} + +void Adam::DisableShadowWeights() { + if (shadow_enable_ == false) { + return; + } + shadow_enable_ = false; + for (auto ¶m : params_) { g_shadow_registry.erase(param.get()); } + shadow_weights_.clear(); + LOG(INFO) << "Disable shadow weights for Adam optimizer"; +} + +void Adam::RefreshShadowWeights() { + if (!shadow_enable_) { + return; + } + // The param was updated in place by the checkpoint (LoadStateDict uses CopyFrom, so the pointer is unchanged): + // just re-cast the current FP32 master weights into the shadows; the registry mapping needs no changes. + for (size_t i = 0; i < params_.size(); ++i) { + auto casted = params_[i]->To(shadow_dtype_); + shadow_weights_[i]->CopyFrom(casted); + } + LOG(INFO) << "Refreshed " << shadow_weights_.size() << " shadow weights from current params"; +} + +std::shared_ptr Adam::GetShadow(const Tensor *param) const { + auto it = g_shadow_registry.find(param); + if (it != g_shadow_registry.end()) { + return it->second; + } else { + LOG(WARNING) << "Shadow weight not found for the given parameter."; + return nullptr; + } +} + void Adam::Step() { ++t_; @@ -114,8 +191,30 @@ void Adam::Step() { auto device = param->GetDevice(); core::DeviceGuard guard(device); - auto kernel = Dispatcher::Instance().GetKernel({device.type(), "AdamAccumulateGrad"}); - kernel.Call(grad, param, m, v, learning_rate_, beta1_, beta2_, eps_, t_); + auto *sparse_state = GetSparseRowGradState(param.get()); + if (sparse_state && !sparse_state->poisoned) { + // Sparse weight (embedding): only the rows hit since the last clear exist in the + // persistent buffer, so update just those rows of param/m/v/shadow. + if (shadow_enable_) { + auto shadow_weight = shadow_weights_[i]; + auto kernel = Dispatcher::Instance().GetKernel({device.type(), "AdamSparseRowsShadow"}); + kernel.Call(grad, param, shadow_weight, m, v, sparse_state->row_list, sparse_state->count, + learning_rate_, beta1_, beta2_, eps_, t_); + } else { + auto kernel = Dispatcher::Instance().GetKernel({device.type(), "AdamSparseRows"}); + kernel.Call(grad, param, m, v, sparse_state->row_list, sparse_state->count, learning_rate_, + beta1_, beta2_, eps_, t_); + } + continue; + } + if (shadow_enable_) { + auto shadow_weight = shadow_weights_[i]; + auto kernel = Dispatcher::Instance().GetKernel({device.type(), "AdamAccumulateGradShadow"}); + kernel.Call(grad, param, shadow_weight, m, v, learning_rate_, beta1_, beta2_, eps_, t_); + } else { + auto kernel = Dispatcher::Instance().GetKernel({device.type(), "AdamAccumulateGrad"}); + kernel.Call(grad, param, m, v, learning_rate_, beta1_, beta2_, eps_, t_); + } } } diff --git a/nsys/gap3.py b/nsys/gap3.py new file mode 100644 index 000000000..ba96dedb9 --- /dev/null +++ b/nsys/gap3.py @@ -0,0 +1,52 @@ +import sqlite3 +c = sqlite3.connect('llama3_nvtx.sqlite') +def q(s,*a): return c.execute(s,a).fetchall() +S = {i:v for i,v in c.execute("select id,value from StringIds")} +a,b = q("select start,end from NVTX_EVENTS where coalesce(text,(select value from StringIds where id=textId))='Step_2'")[0] +CK = {0:'HtoH',1:'HtoD',2:'DtoH',3:'HtoA',4:'AtoH',8:'DtoD'} + +ev = [] +for s,e,n in q("""select k.start,k.end,s.value from CUPTI_ACTIVITY_KIND_KERNEL k + join StringIds s on k.demangledName=s.id where k.start>=? and k.start=? and start=? and start merged[i][1]] +gaps.sort(key=lambda g: -(g[1]-g[0])) +idle = sum(g[1]-g[0] for g in gaps) +span = b-a +print(f"=== step3 窗口 {span/1e6:.3f} ms | GPU 活动段 {len(merged)} | gap {len(gaps)} 个 ===") +print(f" GPU 空闲合计 = {idle/1e6:.3f} ms = {100*idle/span:.1f}% 窗口\n") +print(f"=== 最大的 10 个 GPU 空隙 ===") +print(f" {'#':<3}{'起点(ms)':>9}{'空隙(ms)':>10} {'空隙前最后一个活动':<40}{'空隙后第一个活动'}") +for i,(gs,ge,pv,nx) in enumerate(gaps[:10],1): + print(f" {i:<3}{(gs-a)/1e6:>9.3f}{(ge-gs)/1e6:>10.4f} {pv[0]+':'+pv[1][:36]:<40}{nx[0]+':'+nx[1][:34]}") + +print(f"\n=== 最大 5 个 gap 期间 host 在做什么 ===") +for i,(gs,ge,pv,nx) in enumerate(gaps[:5],1): + print(f"\n --- gap#{i} 起点 +{(gs-a)/1e6:.3f} ms, 持续 {(ge-gs)/1e6:.4f} ms ---") + print(f" 前: {pv[0]}:{pv[1][:60]}") + print(f" 后: {nx[0]}:{nx[1][:60]}") + H = {} + for nid,s2,e2 in q("select nameId,start,end from CUPTI_ACTIVITY_KIND_RUNTIME where start>=? and start9.1f} us avg={t/n/1e3:.1f} us") + # 跨 gap 的长阻塞 API + lg = q("""select r.start,r.end,s.value from CUPTI_ACTIVITY_KIND_RUNTIME r join StringIds s on r.nameId=s.id + where r.start? order by (r.end-r.start) desc limit 3""",ge,gs) + for s2,e2,nm in lg: + if e2-s2 > (ge-gs)*0.5: + print(f" ★跨整个 gap 的阻塞调用: {nm} 时长 {(e2-s2)/1e6:.4f} ms (起 +{(s2-a)/1e6:.3f})") diff --git a/nsys/gpt2_nvtx.nsys-rep b/nsys/gpt2_nvtx.nsys-rep new file mode 100644 index 000000000..3c30f307a Binary files /dev/null and b/nsys/gpt2_nvtx.nsys-rep differ diff --git a/nsys/gpt2_nvtx.sqlite b/nsys/gpt2_nvtx.sqlite new file mode 100644 index 000000000..7d3ed258c Binary files /dev/null and b/nsys/gpt2_nvtx.sqlite differ diff --git a/nsys/llama3_nvtx.nsys-rep b/nsys/llama3_nvtx.nsys-rep new file mode 100644 index 000000000..a3c9c9fb7 Binary files /dev/null and b/nsys/llama3_nvtx.nsys-rep differ diff --git a/nsys/llama3_nvtx.sqlite b/nsys/llama3_nvtx.sqlite new file mode 100644 index 000000000..3764da446 Binary files /dev/null and b/nsys/llama3_nvtx.sqlite differ diff --git a/nsys/mem_3iter.txt b/nsys/mem_3iter.txt new file mode 100644 index 000000000..4fd19883a --- /dev/null +++ b/nsys/mem_3iter.txt @@ -0,0 +1,36 @@ +0 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +33 +4313 +24039 +3 +3 +3 +3 diff --git a/nsys/mem_probe.txt b/nsys/mem_probe.txt new file mode 100644 index 000000000..5d4803670 --- /dev/null +++ b/nsys/mem_probe.txt @@ -0,0 +1,683 @@ +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +1 +1595 +3675 +6235 +24009 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +3 +20 +4121 +24039 +3 +3 +3 +3 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 diff --git a/nsys/s1_full.py b/nsys/s1_full.py new file mode 100644 index 000000000..c88fbcfa8 --- /dev/null +++ b/nsys/s1_full.py @@ -0,0 +1,76 @@ +import sqlite3, re, math +c = sqlite3.connect('llama3_nvtx.sqlite') +def q(s,*a): return c.execute(s,a).fetchall() +def q1(s,*a): return c.execute(s,a).fetchone() +A,B = q1("""select start,end from NVTX_EVENTS + where coalesce(text,(select value from StringIds where id=textId))='Step_1'""") +W = B-A +SM, BW, MAXW, MAXB = 170, 1792128000000, 48, 24 +TPS = MAXW*32 # 每 SM 最大线程 + +def classify(n): + n = n.replace('void ','').strip() + if 'cutlass_' in n: + m = re.search(r'cutlass_(\d+)_(\w+)_(\w+)_(\d+x\d+)_(\d+x\d+)_(\w+)_align(\d+)', n) + if m: return f"GEMM cutlass_{m.group(3)}[{m.group(4)},{m.group(6)},align{m.group(7)}]" + if 'magma_sgemmEx' in n: return "GEMM magma_batched_sgemm" + n2 = n.replace('::','') + kname = n2.split('<')[0].split('::')[-1] + m = re.search(r'kernels::cuda::(\w+?)(Forward|Backward)\(', n) + if m and kname.startswith(('Binary','Unary','Generic')): + v = 'vec' if 'Vectorized' in kname else ('nb' if 'NoBroadcast' in kname else 'bcast') + return f"{m.group(1)}{m.group(2)} [{kname.replace('Kernel','')}/{v}]" + if 'GenericReduce' in kname: + m2 = re.search(r'::(\w+)Finalize', n) + return f"Reduce<{m2.group(1) if m2 else '?'}> [{kname.replace('Kernel','')}]" + return kname + +K = q("""select k.start,k.end,k.gridX*k.gridY*k.gridZ,k.blockX*k.blockY*k.blockZ, + k.registersPerThread,k.staticSharedMemory,k.correlationId,sm.value + from CUPTI_ACTIVITY_KIND_KERNEL k join StringIds sm on k.demangledName=sm.id + where k.start>=? and k.start5}{'Σ(ms)':>9}{'占比':>7}{'avg(us)':>9}{'SM填充':>8}") +for i,(nm,g) in enumerate(sorted(G.items(),key=lambda x:-x[1]['t']),1): + print(f" {i:<3}{nm[:47]:<48}{g['n']:>5}{g['t']/1e6:>9.4f}{100*g['t']/TOT:>6.2f}%" + f"{g['t']/g['n']/1e3:>9.2f}{100*g['fill']/g['t']:>7.1f}%") + +# ── 大类汇总 ── +CAT={'GEMM(cutlass+magma)':['GEMM'],'Optimizer(Adam)':['AdamAccumulate'],'GradAccum(SGD/累加)':['AccumulateGradKernel'], + 'Elementwise(逐元素)':['Mul','Add','Sub','Div','Sigmoid','Rsqrt','Pow','AddScalar','MulScalar'], + 'Reshape/搬运':['FillKernel','TransposeForward','Slice','Stack','RepeatInterleave','Triu'], + 'Reduce/Norm':['Reduce','GenericReduce'],'Embedding':['Embedding'],'Softmax':['Softmax'], + 'CrossEntropy':['CrossEntropy'],'Mask':['Mask']} +print(f"\n═══ ② 按功能大类汇总 ═══") +print(f" {'大类':<26}{'次数':>6}{'Σ(ms)':>10}{'占比':>8}{'SM填充率':>10}") +acc={} +for nm,g in G.items(): + hit='其他' + for cat,pats in CAT.items(): + if any(p in nm for p in pats): hit=cat; break + a=acc.setdefault(hit,[0,0,0.0]); a[0]+=g['n']; a[1]+=g['t']; a[2]+=g['fill'] +for cat,(n,t,f) in sorted(acc.items(),key=lambda x:-x[1][1]): + print(f" {cat:<26}{n:>6}{t/1e6:>10.4f}{100*t/TOT:>7.2f}%{100*f/t:>9.1f}%") +print(f" {'─'*26}{'─'*6}{'─'*10}{'─'*8}{'─'*10}") +print(f" {'合计':<26}{len(K):>6}{TOT/1e6:>10.4f}{100.0:>7.2f}%{100*sum(g['fill'] for g in G.values())/TOT:>9.1f}%") diff --git a/nsys/s1_kernel.py b/nsys/s1_kernel.py new file mode 100644 index 000000000..044bec8d8 --- /dev/null +++ b/nsys/s1_kernel.py @@ -0,0 +1,67 @@ +import sqlite3, re +c = sqlite3.connect('llama3_nvtx.sqlite') +def q(s,*a): return c.execute(s,a).fetchall() +def q1(s,*a): return c.execute(s,a).fetchone() +S = {i:v for i,v in c.execute("select id,value from StringIds")} +A,B = q1("""select start,end from NVTX_EVENTS + where coalesce(text,(select value from StringIds where id=textId))='Step_1'""") +W = B-A +st = q("""select coalesce(text,(select value from StringIds where id=textId)) t,start,end + from NVTX_EVENTS where t like 'Step_%' or t in ('Forward','Backward','Optimizer','ZeroGrad','LossReadback') + order by start""") +S0 = [x for x in st if x[0]=='Step_0'][0] +warm_end = S0[1] + +def clean(n): + n = n.replace('void ','').strip() + m = re.search(r'cutlass_(\d+)_(\w+?)_(\w+?)_(\d+x\d+x\d+)_(\w+?)_align', n) + if m: return f"cutlass_{m.group(2)}_{m.group(3)}_{m.group(4)}_{m.group(5)}" + m = re.search(r'cublas\w*', n, re.I) + if m: return m.group(0) + base = n.split('<')[0] + base = base.split('(')[0] + return base.split('::')[-1] + +K = q("""select k.start,k.end,k.gridX,k.gridY,k.gridZ,k.blockX,k.blockY,k.blockZ, + k.registersPerThread,k.staticSharedMemory,k.correlationId,sm.value + from CUPTI_ACTIVITY_KIND_KERNEL k join StringIds sm on k.demangledName=sm.id + where k.start>=? and k.start=? and start5} 个 kernel Σ {ownt[k]/1e6:>8.3f} ms ({100*ownt[k]/tot:>5.1f}% 设备时间)") + +print(f"\n════════ 各类 kernel 占用时间（按 Σ 设备时间排序）════════") +G={} +for s,e,gx,gy,gz,bx,by,bz,rp,shm,cid,nm in K: + g=G.setdefault(clean(nm),[0,0,gx*gy*gz,bx*by*bz,[]]) + g[0]+=1; g[1]+=e-s; g[4].append((e-s,gx*gy*gz,bx*by*bz,rp,shm)) +print(f" {'#':<3}{'kernel 类别':<40}{'次数':>6}{'Σ时间(ms)':>11}{'占比':>8}{'avg(us)':>9}{'blocks':>10}{'thr/blk':>8}") +for i,(nm,(n,t,gb,bb,_)) in enumerate(sorted(G.items(),key=lambda x:-x[1][1]),1): + print(f" {i:<3}{nm[:39]:<40}{n:>6}{t/1e6:>11.4f}{100*t/tot:>7.2f}%{t/n/1e3:>9.2f}{gb:>10}{bb:>8}") +print(f"\n 合计 {len(G)} 类，{len(K)} 个 kernel，Σ {tot/1e6:.3f} ms") diff --git a/nsys/s1_mem.py b/nsys/s1_mem.py new file mode 100644 index 000000000..124bb149f --- /dev/null +++ b/nsys/s1_mem.py @@ -0,0 +1,73 @@ +import sqlite3, math +c=sqlite3.connect('llama3_nvtx.sqlite') +def q(s,*a): return c.execute(s,a).fetchall() +def q1(s,*a): return c.execute(s,a).fetchone() +A,B=q1("select start,end from NVTX_EVENTS where coalesce(text,(select value from StringIds where id=textId))='Step_1'") +W=B-A; SM,BW,MAXW,MAXB=170,1792128000000,48,24; TPS=MAXW*32 +CK={1:'HtoD',2:'DtoH',8:'DtoD',9:'HtoH',0:'Unknown'} + +print("═══ ③ 各类 memory 操作占用时间与实测带宽 ═══") +print(f" {'操作':<10}{'次数':>6}{'Σ时间(ms)':>11}{'Σ字节':>14}{'实测带宽':>13}{'参考上限':>12}{'达成率':>9}{'avg(us)':>9}") +M=q("""select copyKind,count(*),sum(end-start),sum(bytes),min(bytes),max(bytes) + from CUPTI_ACTIVITY_KIND_MEMCPY where start>=? and start6}{t/1e6:>11.4f}{by:>14,}{bw/1e9:>10.2f}GB/s{refn:>12}{100*bw/ref:>8.2f}%{t/n/1e3:>9.2f}") + mtot+=t; mt+=by +Z=q1("select count(*),sum(end-start),sum(bytes) from CUPTI_ACTIVITY_KIND_MEMSET where start>=? and start6}{Z[1]/1e6:>11.4f}{Z[2]:>14,}{bw/1e9:>10.2f}GB/s{'显存1.79TB/s':>12}{100*bw/BW:>8.2f}%{Z[1]/Z[0]/1e3:>9.2f}") + mtot+=Z[1]; mt+=Z[2] +print(f" {'─'*10}{'─'*6}{'─'*11}{'─'*14}{'─'*13}") +print(f" {'合计':<10}{sum(r[1] for r in M)+Z[0]:>6}{mtot/1e6:>11.4f}{mt:>14,}") +print(f"\n memory 操作 Σ {mtot/1e6:.4f} ms = 窗口 {W/1e6:.3f} ms 的 {100*mtot/W:.2f}%，是 kernel 设备时间的 {100*mtot/72373800:.2f}%") + +print(f"\n═══ ④ memcpy 按字节数细分（揭示每一次传输的语义）═══") +print(f" {'方向':<6}{'字节':>10}{'次数':>6}{'Σ(ms)':>9}{'实测带宽':>13} 推断语义") +SEM={16:'StackForward 的 2 个指针数组 (sizeof(void*)*2)',32:'2 个 float 标量参数', + 24:'3 个 float 标量参数',40:'5 个 float 标量参数',96:'12 个 float 标量/Adam 超参', + 1024:'CrossEntropy 的 per-token loss (bs=256 × 4B)',4:'标量 loss 回读', + 2048:'输入 batch 数据',8192:'x/y 输入张量 (4*64*8B?)'} +for ck,by,n,t in q("""select copyKind,bytes,count(*),sum(end-start) from CUPTI_ACTIVITY_KIND_MEMCPY + where start>=? and start10,}{n:>6}{t/1e6:>9.4f}{by/(t/1e9)/1e9:>10.2f}GB/s {SEM.get(by,'?')}") + +print(f"\n═══ ⑤ GPU 利用率三层分解 ═══") +K=q("""select k.start,k.end,k.gridX*k.gridY*k.gridZ,k.blockX*k.blockY*k.blockZ,sm.value + from CUPTI_ACTIVITY_KIND_KERNEL k join StringIds sm on k.demangledName=sm.id + where k.start>=? and k.start4}{'Σ(ms)':>8}{'SM填充':>8}{'blocks':>9}{'容量':>7}") +for nm,(n,t,f,gb,bb,cap) in sorted(G.items(),key=lambda x:x[1][2]/x[1][1])[:9]: + print(f" {nm.replace('void infini_train::kernels::cuda::','').replace('::','')[:61]:<62}" + f"{n:>4}{t/1e6:>8.4f}{100*f/t:>7.1f}%{gb:>9}{cap:>7}") diff --git a/nsys/s1_rate.py b/nsys/s1_rate.py new file mode 100644 index 000000000..bbd078542 --- /dev/null +++ b/nsys/s1_rate.py @@ -0,0 +1,56 @@ +import sqlite3 +c=sqlite3.connect('llama3_nvtx.sqlite') +def q(s,*a): return c.execute(s,a).fetchall() +def q1(s,*a): return c.execute(s,a).fetchone() +A,B=q1("select start,end from NVTX_EVENTS where coalesce(text,(select value from StringIds where id=textId))='Step_1'") +BW=1792128000000; PEAK=170*128*2*2.407e9 +print("═══ ⑦ Adam kernel：grid 反推参数分组 + 显存带宽核算 ═══") +rows=q("""select k.gridX*k.gridY*k.gridZ gb,k.blockX bb,count(*) n,sum(k.end-k.start) t + from CUPTI_ACTIVITY_KIND_KERNEL k join StringIds sm on k.demangledName=sm.id + where k.start>=? and k.start13}{'次':>4}{'参数总量':>14}{'Σ(ms)':>9}{'avg(us)':>9}{'访存GB':>9}{'GB/s':>8}{'峰值%':>7} 推断") +for gb,bb,n,t in rows: + elem=gb*bb; p=elem*n; traf=p*7*4; bw=traf/(t/1e9) + print(f" {elem:>13,}{n:>4}{p:>14,}{t/1e6:>9.4f}{t/n/1e3:>9.1f}{traf/1e9:>9.3f}" + f"{bw/1e9:>8.0f}{100*bw/BW:>6.1f}% {NAM.get(elem,'?')}") + te+=p; tt+=t; tr+=traf +bw=tr/(tt/1e9) +print(f" {'─'*13}{'─'*4}{'─'*14}{'─'*9}{'─'*9}{'─'*9}{'─'*8}{'─'*7}") +print(f" {'合计':>13}{sum(r[2] for r in rows):>4}{te:>14,}{tt/1e6:>9.4f}{'':>9}{tr/1e9:>9.3f}{bw/1e9:>8.0f}{100*bw/BW:>6.1f}%") +print(f"\n ★ Adam 覆盖 {te:,} 参数 = {te/1e9:.3f}B ≈ llama3.2-1B 全量") +print(f" (16×6291456 qkv + 16×4194304 o + 48×16777216 MLP + 2×262668288 emb/head + 33×2048 norm") +print(f" = {16*6291456+16*4194304+48*16777216+2*262668288+33*2048:,}，本窗口见到 {te:,}，") +print(f" 差 {16*6291456+16*4194304+48*16777216+2*262668288+33*2048-te:,} 即溢出到 Step_2 的部分)") +print(f" ★ 访存 = 参数×7×4B（读 param/grad/m/v + 写 param/m/v）= {tr/1e9:.2f} GB / {tt/1e6:.3f} ms") +print(f" → {bw/1e9:.0f} GB/s = 显存峰值 {BW/1e9:.0f} GB/s 的 {100*bw/BW:.1f}%，各组一致 → 纯带宽受限，kernel 本身无优化空间") + +print(f"\n═══ ⑧ GEMM：kernel 数核对 + 算力达成率 ═══") +cut=q1("""select count(*),sum(k.end-k.start) from CUPTI_ACTIVITY_KIND_KERNEL k join StringIds sm on k.demangledName=sm.id + where k.start>=? and k.start=? and k.start9.1f} GFLOP ({100*L*f_layer/flops:>4.1f}%)") +print(f" lm_head = {f_head/1e9:>9.1f} GFLOP ({100*f_head/flops:>4.1f}%) ← 单个张量占 1/5") +print(f" attention = {L*f_attn/1e9:>9.3f} GFLOP ({100*L*f_attn/flops:>4.2f}%)") +print(f" 合计 = {flops/1e12:>9.4f} TFLOP") +print(f"\n GEMM 实测 Σ {sec*1e3:.4f} ms → 达成 {ach/1e12:.2f} TFLOPS") +print(f" RTX 5090 FP32 峰值 = 170×128×2×2.407GHz = {PEAK/1e12:.1f} TFLOPS") +print(f" ★ FP32 算力达成率 = {100*ach/PEAK:.1f}%") +print(f" ★ 全部 cutlass 名含 'simt'+'align1'：走 CUDA core FP32、未向量化，未用 Tensor Core") diff --git a/scripts/assets/prepare-infinitrain-assets.sh b/scripts/assets/prepare-infinitrain-assets.sh index 4df1abdd0..cbcc8950b 100755 --- a/scripts/assets/prepare-infinitrain-assets.sh +++ b/scripts/assets/prepare-infinitrain-assets.sh @@ -38,8 +38,23 @@ REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)" DATA_DIR="${DATA_DIR:-${REPO_ROOT}/data}" CACHE_DIR="${DATA_DIR}/.cache" PYTHON="${PYTHON:-python3}" +PIP_INDEX_URL="${PIP_INDEX_URL:-https://pypi.org/simple}" +PIP_EXTRA_INDEX_URL="${PIP_EXTRA_INDEX_URL:-}" FORCE="${FORCE:-0}" SKIP_LLAMA3_WEIGHTS="${SKIP_LLAMA3_WEIGHTS:-0}" +MODEL_SOURCE="${MODEL_SOURCE:-huggingface}" +MODEL_REPO_ID="${MODEL_REPO_ID:-meta-llama/Llama-3.2-1B}" + +if [[ "${MODEL_SOURCE}" == "modelscope" && -z "${MODEL_REPO_ID:-}" ]]; then + echo "ERROR: MODEL_SOURCE=modelscope requires MODEL_REPO_ID to point to a valid ModelScope repo." >&2 + echo "The official Meta-Llama-3.2-1B repository is not publicly available on ModelScope." >&2 + echo "Use Hugging Face with HF_TOKEN=hf_xxx, or set MODEL_REPO_ID to your own mirror/private repo." >&2 + exit 2 +fi + +if [[ "${MODEL_SOURCE}" == "modelscope" ]]; then + MODEL_REPO_ID="${MODEL_REPO_ID}" +fi GPT2_DIR="${DATA_DIR}/gpt2" LLAMA3_DIR="${DATA_DIR}/llama3" @@ -94,7 +109,7 @@ prepare_gpt2() { # InfiniTrain's GPT-2 LLMC loader currently accepts the FP32 v3 file. # These artifacts are the same llm.c starter-pack files used by TinyInfiniTrain. - local base="https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main" + local base="https://hf-mirror.com/datasets/karpathy/llmc-starter-pack/resolve/main" local files=( "gpt2_124M.bin" @@ -129,19 +144,34 @@ ensure_llama_python() { fi if ! "${py}" - <<'PY' >/dev/null 2>&1 +import os import numpy import huggingface_hub import socksio import transformers +if os.environ.get("MODEL_SOURCE", "huggingface") == "modelscope": + import modelscope PY then log "Installing LLaMA3 preparation dependencies into ${venv}" - "${py}" -m pip install --upgrade pip - "${py}" -m pip install \ + "${py}" -m pip install --upgrade pip setuptools wheel + local pip_args=( + --index-url "${PIP_INDEX_URL}" + ) + if [[ -n "${PIP_EXTRA_INDEX_URL}" ]]; then + pip_args+=(--extra-index-url "${PIP_EXTRA_INDEX_URL}") + fi + pip_args+=( + --trusted-host pypi.org + --trusted-host files.pythonhosted.org + --trusted-host mirrors.aliyun.com + ) + "${py}" -m pip install "${pip_args[@]}" \ "numpy>=1.24" \ "huggingface_hub>=0.24" \ "socksio>=1.0" \ - "transformers>=4.43" + "transformers>=4.43" \ + "modelscope>=1.18" fi } @@ -162,7 +192,10 @@ prepare_llama3() { # the token saved by `hf auth login`. TINY_SHAKESPEARE_TXT="${tiny_txt}" \ LLAMA3_OUTPUT_DIR="${LLAMA3_DIR}" \ - LLAMA3_CACHE_DIR="${CACHE_DIR}/llama3-hf" \ + LLAMA3_CACHE_DIR="${CACHE_DIR}/llama3-${MODEL_SOURCE}" \ + MODEL_SOURCE="${MODEL_SOURCE}" \ + MODEL_REPO_ID="${MODEL_REPO_ID}" \ + HF_TOKEN="${HF_TOKEN:-}" \ SKIP_LLAMA3_WEIGHTS="${SKIP_LLAMA3_WEIGHTS}" \ FORCE="${FORCE}" \ "${py}" "${SCRIPT_DIR}/prepare_llama3_assets.py" diff --git a/scripts/assets/prepare_llama3_assets.py b/scripts/assets/prepare_llama3_assets.py index 7487f9d71..9838f4ad3 100755 --- a/scripts/assets/prepare_llama3_assets.py +++ b/scripts/assets/prepare_llama3_assets.py @@ -7,10 +7,33 @@ from pathlib import Path import numpy as np -from huggingface_hub import get_token, snapshot_download from transformers import AutoTokenizer -MODEL_ID = "meta-llama/Llama-3.2-1B" +MODEL_SOURCE = os.environ.get("MODEL_SOURCE", "huggingface").lower() +MODEL_REPO_ID = os.environ.get("MODEL_REPO_ID") +if MODEL_REPO_ID is None: + MODEL_REPO_ID = "meta-llama/Llama-3.2-1B" +MODEL_ID = MODEL_REPO_ID + +if MODEL_SOURCE == "huggingface": + from huggingface_hub import get_token, snapshot_download + token = os.environ.get("HF_TOKEN") or get_token() + if not token: + raise SystemExit( + f"\nLLaMA3 preparation needs access to {MODEL_ID}.\n" + "1) Accept the model license on Hugging Face.\n" + "2) Run `hf auth login` or export HF_TOKEN=hf_xxx.\n" + ) +elif MODEL_SOURCE == "modelscope": + try: + from modelscope import snapshot_download as ms_snapshot_download + except ImportError as exc: + raise SystemExit( + "ModelScope support requires `pip install modelscope` in the asset-prep environment." + ) from exc + token = None +else: + raise SystemExit(f"Unsupported MODEL_SOURCE={MODEL_SOURCE!r}; expected 'huggingface' or 'modelscope'.") out_dir = Path(os.environ["LLAMA3_OUTPUT_DIR"]) cache_dir = Path(os.environ["LLAMA3_CACHE_DIR"]) @@ -21,14 +44,6 @@ out_dir.mkdir(parents=True, exist_ok=True) cache_dir.mkdir(parents=True, exist_ok=True) -token = os.environ.get("HF_TOKEN") or get_token() -if not token: - raise SystemExit( - f"\nLLaMA3 preparation needs access to {MODEL_ID}.\n" - "1) Accept the model license on Hugging Face.\n" - "2) Run `hf auth login` or export HF_TOKEN=hf_xxx.\n" - ) - allow_patterns = [ "config.json", "generation_config.json", @@ -45,13 +60,35 @@ "model.safetensors.index.json", ]) -print(f"[llama3] downloading/reusing Hugging Face files for {MODEL_ID}") -model_dir = Path(snapshot_download( - repo_id=MODEL_ID, - token=token, - cache_dir=str(cache_dir), - allow_patterns=allow_patterns, -)) +if MODEL_SOURCE == "huggingface": + print(f"[llama3] downloading/reusing Hugging Face files for {MODEL_ID}") + model_dir = Path(snapshot_download( + repo_id=MODEL_ID, + token=token, + cache_dir=str(cache_dir), + allow_patterns=allow_patterns, + )) +else: + print(f"[llama3] downloading/reusing ModelScope files for {MODEL_ID}") + try: + model_dir = Path(ms_snapshot_download( + model_id=MODEL_ID, + cache_dir=str(cache_dir), + allow_patterns=allow_patterns, + )) + except TypeError: + model_dir = Path(ms_snapshot_download( + MODEL_ID, + cache_dir=str(cache_dir), + allow_patterns=allow_patterns, + )) + except Exception as exc: # pragma: no cover - message is for users + raise SystemExit( + "\nModelScope download failed for the requested repo.\n" + f"repo_id={MODEL_ID}\n" + "This official Meta-Llama 3.2 1B repo is not publicly available on ModelScope.\n" + "Use Hugging Face with HF_TOKEN=hf_xxx, or set MODEL_REPO_ID to a valid mirror/private repo." + ) from exc # --------------------------------------------------------------------------- # TinyShakespeare -> InfiniTrain / llm.c LLaMA-3 data format @@ -83,7 +120,7 @@ def write_datafile(path: Path, toks): tokenizer = AutoTokenizer.from_pretrained( model_dir, local_files_only=True, - token=token, + token=token if MODEL_SOURCE == "huggingface" else None, use_fast=True, ) diff --git a/tests/autograd/test_autograd_normalization_backward.cc b/tests/autograd/test_autograd_normalization_backward.cc index 12241c2b6..7af79a26f 100644 --- a/tests/autograd/test_autograd_normalization_backward.cc +++ b/tests/autograd/test_autograd_normalization_backward.cc @@ -55,4 +55,28 @@ TEST_P(AutogradNormalizationBackwardTest, LayerNormBackwardZeroBias) { EXPECT_EQ(grad_inputs.size(), 3); } +TEST_P(AutogradNormalizationBackwardTest, RMSNormBackward) { + const std::vector input_dims{1, 2, 4}; + std::vector input_values{-1.5f, -0.5f, 0.5f, 1.5f, -3.0f, -1.0f, 1.0f, 3.0f}; + std::vector weight_values{1.0f, 0.5f, -1.0f, 2.0f}; + auto input = std::make_shared(input_values.data(), input_dims, DataType::kFLOAT32, GetDevice()); + auto weight + = std::make_shared(weight_values.data(), std::vector{4}, DataType::kFLOAT32, GetDevice()); + + auto rmsnorm_fn = std::make_shared(1e-5f); + auto result = rmsnorm_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 2); + test::ExpectTensorNear(result[1], {0.89442360f, 0.44721314f}, 1e-5f); + + std::vector grad_values{1.0f, 2.0f, 3.0f, 4.0f, 0.5f, -1.0f, 2.0f, -0.5f}; + auto grad = std::make_shared(grad_values.data(), input_dims, DataType::kFLOAT32, GetDevice()); + auto grad_inputs = rmsnorm_fn->Backward({grad}); + ASSERT_EQ(grad_inputs.size(), 2); + test::ExpectTensorNear( + grad_inputs[0], + {3.17518568f, 1.65467763f, -3.44352484f, 4.87462664f, -0.17888442f, -0.35777023f, -0.76026261f, -0.04472215f}, + 1e-5f); + test::ExpectTensorNear(grad_inputs[1], {-2.01245522f, -0.44721046f, 2.23606181f, 4.69572210f}, 1e-5f); +} + INFINI_TRAIN_REGISTER_TEST(AutogradNormalizationBackwardTest); diff --git a/tests/autograd/test_autograd_normalization_forward.cc b/tests/autograd/test_autograd_normalization_forward.cc index ebb876d8d..ba6791335 100644 --- a/tests/autograd/test_autograd_normalization_forward.cc +++ b/tests/autograd/test_autograd_normalization_forward.cc @@ -2,6 +2,7 @@ #include "gtest/gtest.h" +#include "infini_train/include/autocast.h" #include "infini_train/include/autograd/normalization.h" #include "infini_train/include/nn/parallel/global.h" #include "infini_train/include/tensor.h" @@ -61,4 +62,61 @@ TEST_P(AutogradNormalizationForwardTest, LayerNormThreeDim) { EXPECT_EQ(result[0]->Dims(), (std::vector{2, 1, 4})); } +TEST_P(AutogradNormalizationForwardTest, RMSNormForward) { + const std::vector input_dims{1, 2, 4}; + std::vector input_values{-1.5f, -0.5f, 0.5f, 1.5f, -3.0f, -1.0f, 1.0f, 3.0f}; + std::vector weight_values{1.0f, 0.5f, -1.0f, 2.0f}; + auto input = std::make_shared(input_values.data(), input_dims, DataType::kFLOAT32, GetDevice()); + auto weight + = std::make_shared(weight_values.data(), std::vector{4}, DataType::kFLOAT32, GetDevice()); + + auto rmsnorm_fn = std::make_shared(1e-5f); + auto result = rmsnorm_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 2); + EXPECT_FALSE(result[1]->requires_grad()); + EXPECT_EQ(result[1]->grad_fn(), nullptr); + test::ExpectTensorNear( + result[0], + {-1.34163547f, -0.22360590f, -0.44721180f, 2.68327093f, -1.34163940f, -0.22360657f, -0.44721314f, 2.68327880f}, + 1e-5f); + // rstd is the normalization statistic kept for the backward pass: 1/sqrt(mean(x^2) + eps) per row. + test::ExpectTensorNear(result[1], {0.89442360f, 0.44721314f}, 1e-5f); +} + +TEST_P(AutogradNormalizationForwardTest, RMSNormTwoDimInput) { + // The fused op must stay rank-agnostic: flattened [rows, embed_dim] inputs work like the + // composite path (Mean(-1)/Pow/Rsqrt/Mul), including the rstd shape (leading dims only). + auto input = std::make_shared(std::vector{8, 4}, DataType::kFLOAT32, GetDevice(), true); + input->Fill(2.0f); + auto weight = std::make_shared(std::vector{4}, DataType::kFLOAT32, GetDevice(), true); + weight->Fill(1.0f); + auto rmsnorm_fn = std::make_shared(1e-5f); + auto result = rmsnorm_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 2); + EXPECT_EQ(result[0]->Dims(), (std::vector{8, 4})); + EXPECT_EQ(result[1]->Dims(), (std::vector{8})); + test::ExpectTensorNear(result[0], 0.99999875f, 1e-5f); + test::ExpectTensorNear(result[1], 0.49999937f, 1e-5f); +} + +TEST_P(AutogradNormalizationForwardTest, RMSNormAutocastCastsInputToFP32) { + SKIP_CPU(); + // RMSNorm is registered as a kFP32 op: under a bf16 autocast context the bf16 activation is + // promoted to fp32 before the fused kernel runs, so the input/weight dtype CHECK holds and + // both outputs come back as fp32. + auto input = std::make_shared(std::vector{1, 2, 4}, DataType::kBFLOAT16, GetDevice()); + input->Fill(2.0f); + auto weight = std::make_shared(std::vector{4}, DataType::kFLOAT32, GetDevice()); + weight->Fill(1.0f); + + AutocastGuard guard(GetDevice().type(), DataType::kBFLOAT16); + auto rmsnorm_fn = std::make_shared(1e-5f); + auto result = rmsnorm_fn->Apply({input, weight}); + ASSERT_EQ(result.size(), 2); + EXPECT_EQ(result[0]->Dtype(), DataType::kFLOAT32); + EXPECT_EQ(result[1]->Dtype(), DataType::kFLOAT32); + test::ExpectTensorNear(result[0], 0.99999875f, 1e-5f); + test::ExpectTensorNear(result[1], 0.49999937f, 1e-5f); +} + INFINI_TRAIN_REGISTER_TEST(AutogradNormalizationForwardTest);