|
| 1 | +# NVIDIA Blackwell GPU Architecture Support Implementation Plan |
| 2 | + |
| 3 | +## Overview |
| 4 | + |
| 5 | +This document outlines the implementation plan for adding comprehensive NVIDIA Blackwell GPU architecture support to llama.cpp. The plan is structured in phases to ensure systematic development, testing, and validation of Blackwell-specific optimizations. |
| 6 | + |
| 7 | +## Current State Analysis |
| 8 | + |
| 9 | +- **Compute Capability**: Currently supports up to Ada Lovelace (8.9) |
| 10 | +- **Blackwell Support**: [PR #13360](https://github.com/ggml-org/llama.cpp/pull/13360) adds CUDA 12.8 + sm120 build support |
| 11 | +- **Missing Features**: Thread Block Clusters, L2 cache management, HBM3/HBM3e optimizations |
| 12 | +- **Flash Attention**: Multiple kernel variants but no Blackwell-specific optimizations |
| 13 | +- **Compatibility**: Basic functionality works via backward compatibility, but performance is sub-optimal |
| 14 | + |
| 15 | +## Architecture Constants Update |
| 16 | + |
| 17 | +**Critical Finding**: Blackwell GPUs use compute capability **12.0** (sm120), not 10.0 as initially assumed. |
| 18 | + |
| 19 | +## Flash Attention Analysis |
| 20 | + |
| 21 | +llama.cpp implements multiple Flash Attention kernel variants: |
| 22 | +- **MMA-based kernels** (`fattn-mma-f16.cuh`): Modern implementation for Turing+ |
| 23 | +- **Vector kernels** (`fattn-vec-f32/f16.cuh`): For smaller batches/specific dimensions |
| 24 | +- **WMMA kernels** (`fattn-wmma-f16.cu`): Legacy implementation for Volta |
| 25 | +- **Tile kernels** (`fattn-tile-f16/f32.cu`): For architectures without tensor cores |
| 26 | + |
| 27 | +Selection logic in `ggml_cuda_flash_attn_ext()` considers compute capability, batch size, head dimensions, and data types. |
| 28 | + |
| 29 | +## Phase 1: Foundation and Architecture Detection β‘ **ACCELERATED** |
| 30 | + |
| 31 | +### 1.1 Add Blackwell Constants and Detection **β
FOUNDATION COMPLETE** |
| 32 | + |
| 33 | +**Status**: Foundation provided by [PR #13360](https://github.com/ggml-org/llama.cpp/pull/13360) |
| 34 | +- β
CUDA 12.8 toolkit support |
| 35 | +- β
sm120 compilation target |
| 36 | +- β
Build system integration |
| 37 | + |
| 38 | +**Files to modify:** |
| 39 | +- `ggml/src/ggml-cuda/common.cuh` |
| 40 | +- `ggml/src/ggml-cuda/ggml-cuda.cu` |
| 41 | + |
| 42 | +**Updated Implementation:** |
| 43 | +```cpp |
| 44 | +// Add to common.cuh - CORRECTED for actual Blackwell compute capability |
| 45 | +#define GGML_CUDA_CC_BLACKWELL 1200 // B100/B200/RTX50 (12.0) - CORRECTED |
| 46 | +#define GGML_CUDA_CC_BLACKWELL_FUTURE 1300 // Future Blackwell variants |
| 47 | + |
| 48 | +#define GGML_CUDA_CC_IS_BLACKWELL(cc) (cc >= GGML_CUDA_CC_BLACKWELL && cc < GGML_CUDA_CC_BLACKWELL_FUTURE) |
| 49 | +#define GGML_CUDA_CC_SUPPORTS_CLUSTERS(cc) (cc >= GGML_CUDA_CC_BLACKWELL) |
| 50 | +``` |
| 51 | +
|
| 52 | +**Timeline:** ~~Week 1-2~~ **COMPLETE** β
|
| 53 | +
|
| 54 | +### 1.2 Enhanced Device Information Structure |
| 55 | +
|
| 56 | +**Files to modify:** |
| 57 | +- `ggml/src/ggml-cuda/ggml-cuda.cu` (cuda_device_info) |
| 58 | +
|
| 59 | +**New fields:** |
| 60 | +- `max_clusters_per_multiprocessor` |
| 61 | +- `max_blocks_per_cluster` |
| 62 | +- `l2_cache_size` |
| 63 | +- `hbm_bandwidth` |
| 64 | +
|
| 65 | +**Updated Timeline:** Week 1 β‘ (accelerated due to build foundation) |
| 66 | +
|
| 67 | +### 1.3 Blackwell Feature Detection **NEW** |
| 68 | +
|
| 69 | +**Files to create:** |
| 70 | +- `ggml/src/ggml-cuda/blackwell-detect.cu` |
| 71 | +
|
| 72 | +**Implementation:** |
| 73 | +```cpp |
| 74 | +bool ggml_cuda_supports_blackwell_features(int device_id) { |
| 75 | + cudaDeviceProp prop; |
| 76 | + cudaGetDeviceProperties(&prop, device_id); |
| 77 | + |
| 78 | + // Verify compute capability 12.0+ |
| 79 | + int cc = 100 * prop.major + 10 * prop.minor; |
| 80 | + if (!GGML_CUDA_CC_IS_BLACKWELL(cc)) return false; |
| 81 | + |
| 82 | + // Verify cluster support |
| 83 | + int max_cluster_size; |
| 84 | + cudaDeviceGetAttribute(&max_cluster_size, |
| 85 | + cudaDevAttrClusterLaunch, device_id); |
| 86 | + |
| 87 | + return max_cluster_size > 0; |
| 88 | +} |
| 89 | +``` |
| 90 | + |
| 91 | +**Timeline:** Week 1-2 |
| 92 | + |
| 93 | +## Phase 2: Thread Block Clusters Foundation |
| 94 | + |
| 95 | +### 2.1 Cluster Detection and Support Infrastructure |
| 96 | + |
| 97 | +**Files to create:** |
| 98 | +- `ggml/src/ggml-cuda/clusters.cuh` |
| 99 | +- `ggml/src/ggml-cuda/clusters.cu` |
| 100 | + |
| 101 | +**Key functions:** |
| 102 | +- `ggml_cuda_cluster_occupancy()` |
| 103 | +- `ggml_cuda_launch_kernel_clusters()` |
| 104 | +- `ggml_cuda_cluster_sync_init()` |
| 105 | + |
| 106 | +**Updated Timeline:** Week 2-3 β‘ (accelerated) |
| 107 | + |
| 108 | +### 2.2 L2 Cache Management |
| 109 | + |
| 110 | +**Files to modify:** |
| 111 | +- `ggml/src/ggml-cuda/ggml-cuda.cu` |
| 112 | + |
| 113 | +**Implementation:** |
| 114 | +- L2 cache persistence API wrappers |
| 115 | +- Cache allocation strategy for KV cache data |
| 116 | +- Stream-based cache management |
| 117 | + |
| 118 | +**Updated Timeline:** Week 3-4 β‘ (accelerated) |
| 119 | + |
| 120 | +## Phase 3: Flash Attention Blackwell Optimizations |
| 121 | + |
| 122 | +### 3.1 MMA Kernel Enhancements for Blackwell |
| 123 | + |
| 124 | +**Files to modify:** |
| 125 | +- `ggml/src/ggml-cuda/fattn-mma-f16.cuh` |
| 126 | +- `ggml/src/ggml-cuda/fattn-common.cuh` |
| 127 | + |
| 128 | +**Key optimizations:** |
| 129 | + |
| 130 | +#### 3.1.1 Enhanced Shared Memory Usage |
| 131 | +```cpp |
| 132 | +// Leverage 228 KB shared memory per SM vs 164 KB on Ada Lovelace |
| 133 | +template<int DKQ, int DV> |
| 134 | +struct fattn_blackwell_config : fattn_mma_f16_config<DKQ, DV> { |
| 135 | + static constexpr int cc_target = GGML_CUDA_CC_BLACKWELL; // 1200 |
| 136 | + static constexpr int smpb_blackwell = 228 * 1024; // 228 KB |
| 137 | + static constexpr int enhanced_batch_size = smpb_blackwell / (DKQ * sizeof(half)); |
| 138 | + |
| 139 | + // Increase tile sizes for better cache utilization |
| 140 | + static constexpr int nbatch_fa_blackwell = std::min(enhanced_batch_size, 128); |
| 141 | +}; |
| 142 | +``` |
| 143 | +
|
| 144 | +#### 3.1.2 Thread Block Cluster Integration |
| 145 | +```cpp |
| 146 | +template<int DKQ, int DV, int ncols1, int ncols2, int cluster_size> |
| 147 | +__cluster_dims__(cluster_size, 1, 1) |
| 148 | +__global__ void flash_attn_ext_f16_clustered(/* parameters */) { |
| 149 | + // Distributed shared memory across cluster |
| 150 | + extern __shared__ half2 cluster_shared_memory[]; |
| 151 | + |
| 152 | + // Cluster-wide synchronization |
| 153 | + cluster.sync(); |
| 154 | + |
| 155 | + // Enhanced memory access patterns |
| 156 | + // ... |
| 157 | +} |
| 158 | +``` |
| 159 | + |
| 160 | +#### 3.1.3 L2 Cache-Aware KV Access |
| 161 | +```cpp |
| 162 | +// Optimize KV cache access patterns for L2 persistence |
| 163 | +template<typename T> |
| 164 | +__device__ void prefetch_kv_to_l2(const T* kv_data, size_t size) { |
| 165 | + // Use cache hints for Blackwell L2 (126 MB vs 40 MB) |
| 166 | + __builtin_nontemporal_store(); // Blackwell-specific hints |
| 167 | +} |
| 168 | +``` |
| 169 | +
|
| 170 | +**Updated Timeline:** Week 4-7 β‘ (accelerated) |
| 171 | +
|
| 172 | +### 3.2 Kernel Selection Logic Updates |
| 173 | +
|
| 174 | +**Files to modify:** |
| 175 | +- `ggml/src/ggml-cuda/fattn.cu` |
| 176 | +
|
| 177 | +**Enhanced selection for Blackwell:** |
| 178 | +```cpp |
| 179 | +void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { |
| 180 | + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; |
| 181 | + |
| 182 | + if (GGML_CUDA_CC_IS_BLACKWELL(cc)) { |
| 183 | + // Prefer cluster-based kernels for larger problems |
| 184 | + if (can_use_clusters && problem_size_threshold_met) { |
| 185 | + ggml_cuda_flash_attn_ext_mma_f16_clusters(ctx, dst); |
| 186 | + return; |
| 187 | + } |
| 188 | + |
| 189 | + // Use enhanced MMA kernels with Blackwell optimizations |
| 190 | + ggml_cuda_flash_attn_ext_mma_f16_blackwell(ctx, dst); |
| 191 | + return; |
| 192 | + } |
| 193 | + |
| 194 | + // ... existing fallback logic ... |
| 195 | +} |
| 196 | +``` |
| 197 | + |
| 198 | +**Updated Timeline:** Week 6-7 β‘ (accelerated) |
| 199 | + |
| 200 | +### 3.3 Advanced Memory Access Optimizations |
| 201 | + |
| 202 | +#### 3.3.1 HBM3/HBM3e Bandwidth Optimization |
| 203 | +- Implement wider memory transactions (512-bit vs 256-bit) |
| 204 | +- Optimize memory coalescing patterns for higher bandwidth |
| 205 | +- Implement memory prefetching strategies |
| 206 | + |
| 207 | +#### 3.3.2 Async Copy Enhancements |
| 208 | +```cpp |
| 209 | +// Enhanced async copy for Blackwell |
| 210 | +template<int cluster_size> |
| 211 | +__device__ void async_copy_cluster_aware( |
| 212 | + void* dst, const void* src, size_t bytes, |
| 213 | + cuda::barrier<cuda::thread_scope_cluster>& barrier) { |
| 214 | + // Blackwell-optimized async copy with cluster coordination |
| 215 | +} |
| 216 | +``` |
| 217 | +
|
| 218 | +**Updated Timeline:** Week 8-9 β‘ (accelerated) |
| 219 | +
|
| 220 | +## Phase 4: Advanced Blackwell Features |
| 221 | +
|
| 222 | +### 4.1 Distributed Shared Memory Implementation |
| 223 | +
|
| 224 | +**Files to create:** |
| 225 | +- `ggml/src/ggml-cuda/distributed-shared-memory.cuh` |
| 226 | +
|
| 227 | +**Key features:** |
| 228 | +- Cross-block shared memory access |
| 229 | +- Cluster-wide data sharing for attention heads |
| 230 | +- Optimized memory layout for distributed access |
| 231 | +
|
| 232 | +**Updated Timeline:** Week 10-11 β‘ (accelerated) |
| 233 | +
|
| 234 | +### 4.2 Advanced Occupancy Management |
| 235 | +
|
| 236 | +**Files to modify:** |
| 237 | +- `ggml/src/ggml-cuda/fattn-common.cuh` |
| 238 | +
|
| 239 | +**Implementation:** |
| 240 | +- `cudaOccupancyMaxActiveClusters` integration |
| 241 | +- Dynamic cluster size selection |
| 242 | +- Load balancing across SMs |
| 243 | +
|
| 244 | +**Updated Timeline:** Week 11-12 β‘ (accelerated) |
| 245 | +
|
| 246 | +### 4.3 Multi-Head Attention Cluster Optimization |
| 247 | +
|
| 248 | +**New kernel variants:** |
| 249 | +- Cluster-aware multi-head processing |
| 250 | +- Cross-head data sharing via distributed shared memory |
| 251 | +- Optimized attention head grouping strategies |
| 252 | +
|
| 253 | +**Updated Timeline:** Week 12-13 β‘ (accelerated) |
| 254 | +
|
| 255 | +## Phase 5: General CUDA Kernel Optimizations |
| 256 | +
|
| 257 | +### 5.1 Matrix Operations Enhancement |
| 258 | +
|
| 259 | +**Files to modify:** |
| 260 | +- `ggml/src/ggml-cuda/gemm.cu` |
| 261 | +- `ggml/src/ggml-cuda/mul-mat.cu` |
| 262 | +
|
| 263 | +**Optimizations:** |
| 264 | +- Leverage 255 registers per thread with improved scheduling |
| 265 | +- Enhanced warp-level primitives for Blackwell |
| 266 | +- L2 cache persistence for weight matrices |
| 267 | +
|
| 268 | +**Updated Timeline:** Week 14-15 β‘ (accelerated) |
| 269 | +
|
| 270 | +### 5.2 Attention-Adjacent Operations |
| 271 | +
|
| 272 | +**Files to modify:** |
| 273 | +- `ggml/src/ggml-cuda/rope.cu` (Rotary Position Embedding) |
| 274 | +- `ggml/src/ggml-cuda/norm.cu` (Layer Normalization) |
| 275 | +
|
| 276 | +**Optimizations:** |
| 277 | +- Thread block cluster integration where beneficial |
| 278 | +- Enhanced shared memory usage |
| 279 | +- Optimized memory access patterns |
| 280 | +
|
| 281 | +**Updated Timeline:** Week 15-16 β‘ (accelerated) |
| 282 | +
|
| 283 | +## Phase 6: Performance Validation and Optimization |
| 284 | +
|
| 285 | +### 6.1 Benchmarking Infrastructure |
| 286 | +
|
| 287 | +**Files to create:** |
| 288 | +- `tools/blackwell-bench/` |
| 289 | +- Comprehensive benchmarking suite |
| 290 | +- Performance regression detection |
| 291 | +- A/B testing framework |
| 292 | +
|
| 293 | +**Updated Timeline:** Week 17-18 β‘ (accelerated) |
| 294 | +
|
| 295 | +### 6.2 Performance Tuning |
| 296 | +
|
| 297 | +**Focus areas:** |
| 298 | +- Kernel parameter auto-tuning |
| 299 | +- Dynamic optimization based on problem size |
| 300 | +- Memory allocation strategy optimization |
| 301 | +- Cache management tuning |
| 302 | +
|
| 303 | +**Updated Timeline:** Week 18-20 β‘ (accelerated) |
| 304 | +
|
| 305 | +### 6.3 Integration Testing |
| 306 | +
|
| 307 | +**Test scenarios:** |
| 308 | +- Various model architectures (Llama, Mistral, etc.) |
| 309 | +- Different sequence lengths and batch sizes |
| 310 | +- Mixed precision scenarios |
| 311 | +- Multi-GPU configurations |
| 312 | +
|
| 313 | +**Updated Timeline:** Week 20-21 β‘ (accelerated) |
| 314 | +
|
| 315 | +## Phase 7: Documentation and Integration |
| 316 | +
|
| 317 | +### 7.1 Documentation Updates |
| 318 | +
|
| 319 | +**Files to create/modify:** |
| 320 | +- `docs/backend/BLACKWELL.md` |
| 321 | +- Update existing CUDA documentation |
| 322 | +- Code documentation and examples |
| 323 | +
|
| 324 | +**Updated Timeline:** Week 22 β‘ (accelerated) |
| 325 | +
|
| 326 | +### 7.2 Build System Integration **β‘ FOUNDATION COMPLETE** |
| 327 | +
|
| 328 | +**Status**: Core build system complete via [PR #13360](https://github.com/ggml-org/llama.cpp/pull/13360) |
| 329 | +
|
| 330 | +**Remaining tasks:** |
| 331 | +- β
CUDA version detection (complete) |
| 332 | +- β
Blackwell-specific compilation flags (complete) |
| 333 | +- π Optional feature toggles for Blackwell optimizations |
| 334 | +
|
| 335 | +**Updated Timeline:** Week 22 β‘ (accelerated) |
| 336 | +
|
| 337 | +## Updated Success Metrics |
| 338 | +
|
| 339 | +### Performance Targets |
| 340 | +- **Flash Attention**: 20-40% improvement over Ada Lovelace |
| 341 | +- **Overall Inference**: 15-30% improvement in tokens/second |
| 342 | +- **Memory Efficiency**: Better utilization of 126 MB L2 cache |
| 343 | +- **Scalability**: Improved performance on larger context lengths |
| 344 | +
|
| 345 | +### Validation Criteria |
| 346 | +- All existing tests pass |
| 347 | +- No performance regression on older architectures |
| 348 | +- Blackwell-specific optimizations activate correctly for compute capability 12.0+ |
| 349 | +- Proper fallback behavior on non-Blackwell hardware |
| 350 | +
|
| 351 | +## Updated Risk Mitigation |
| 352 | +
|
| 353 | +### Technical Risks - REDUCED β‘ |
| 354 | +- β
**Build Infrastructure**: Resolved by PR #13360 |
| 355 | +- β
**Compute Capability Detection**: Corrected to 12.0 |
| 356 | +- π **Hardware Availability**: Still limited but build foundation ready |
| 357 | +- π **API Changes**: Version detection in place |
| 358 | +- π **Complexity**: Incremental implementation continues |
| 359 | +
|
| 360 | +### Timeline Risks - MITIGATED β‘ |
| 361 | +- β
**Foundation Delays**: Eliminated by PR #13360 |
| 362 | +- π **Scope Creep**: Strict phase gating maintained |
| 363 | +- π **Dependencies**: CUDA 12.8 foundation complete |
| 364 | +
|
| 365 | +## Updated Timeline Summary |
| 366 | +
|
| 367 | +**Original Timeline**: 24 weeks |
| 368 | +**Accelerated Timeline**: 22 weeks β‘ (2-week acceleration) |
| 369 | +
|
| 370 | +**Key Accelerations**: |
| 371 | +- Phase 1: Complete β Immediate start on Phase 2 |
| 372 | +- Phase 2-7: 1-2 week acceleration per phase |
| 373 | +- Build system risks eliminated |
| 374 | +
|
| 375 | +## Immediate Next Steps (Week 1) |
| 376 | +
|
| 377 | +1. **Implement Phase 1.2**: Enhanced device information structure |
| 378 | +2. **Begin Phase 1.3**: Blackwell feature detection |
| 379 | +3. **Start Phase 2.1**: Cluster infrastructure development |
| 380 | +4. **Update all compute capability constants**: 1000 β 1200 |
| 381 | +
|
| 382 | +## Conclusion |
| 383 | +
|
| 384 | +[PR #13360](https://github.com/ggml-org/llama.cpp/pull/13360) provides crucial foundation acceleration for our Blackwell implementation. The corrected compute capability (12.0) and completed build infrastructure allow us to begin advanced optimizations immediately. |
| 385 | +
|
| 386 | +**Key Benefits**: |
| 387 | +- β‘ **2-week timeline acceleration** |
| 388 | +- β
**Build foundation complete** |
| 389 | +- π― **Accurate architecture targeting** (cc 12.0) |
| 390 | +- π **Immediate development start** capability |
| 391 | +
|
| 392 | +The plan now reflects actual Blackwell specifications and leverages the completed foundation to achieve aggressive performance improvements while maintaining our systematic, phased approach. |
0 commit comments