From 0711c291e43fd9042c8a1f743f95fef34876944c Mon Sep 17 00:00:00 2001 From: jasl Date: Sat, 25 Apr 2026 04:47:03 +0800 Subject: [PATCH 01/14] Upgrade CUTLASS to 4.4.2 --- third-party/cutlass | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/third-party/cutlass b/third-party/cutlass index f3fde58372..da5e086dab 160000 --- a/third-party/cutlass +++ b/third-party/cutlass @@ -1 +1 @@ -Subproject commit f3fde58372d33e9a5650ba7b80fc48b3b49d40c8 +Subproject commit da5e086dab31d63815acafdac9a9c5893b1c69e2 From 2206a1d7bc6d4df81737d74efb6bacf5888f1fbb Mon Sep 17 00:00:00 2001 From: jasl Date: Sat, 25 Apr 2026 05:50:22 +0800 Subject: [PATCH 02/14] Add SM120 reference fallbacks for V4 blockers DeepGEMM currently rejects SM12x devices in the DeepSeek V4 forward path because the HC prenorm GEMM and paged MQA logits families only dispatch SM90/SM100 implementations. This change adds an explicit compatibility path for compute capability 12.x while keeping the existing high-performance SM90/SM100 paths unchanged. Normalize CUDA major 12 devices, including SM120 and SM121, to the sm_120f JIT target and sm120 include suffix so a GB10/DGX Spark can build a shared SM12x kernel family. Add an SM12x tf32_hc_prenorm_gemm reference fallback using ATen CUDA matmul plus the required square-sum output, preserving the existing split-output ABI by storing the full reference result in split zero and zeroing the remaining splits. Allow SM12x paged MQA metadata and add an FP8-only SM120 paged MQA logits reference runtime. The new CUDA kernel avoids SM90 WGMMA and SM100 tcgen05/TMA assumptions; it directly dequantizes FP8 q/kv values, applies the KV scale, accumulates in FP32, applies ReLU and per-head weights, and stores invalid token positions as -inf. FP4 paged MQA, non-paged MQA, MegaMoE, and general FP8/FP4 GEMM remain gated until they have separate SM12x implementations. Add focused SM120 regression coverage for HC prenorm GEMM and FP8 paged MQA logits against PyTorch references, including the fused paged KV-cache layout used by the API. Also install the requested project-scoped C++ readability skills under .codex/skills so future CUDA/C++ work in this checkout can use the same local guidance. Verification on 10.0.0.116 (/home/jasl/tmp/DeepGEMM): ./develop.sh completed with CUDA 13.1, then PYTHONPATH=. pytest -q tests/test_sm120_reference.py -q returned '.. [100%]'. A separate JIT command check confirmed the SM12x path compiles with --gpu-architecture=sm_120f. --- .codex/skills/cpp/SKILL.md | 111 ++++ .codex/skills/cpp/references/guidelines.md | 152 +++++ .codex/skills/cpp/references/structure.md | 89 +++ .codex/skills/readable-cpp/SKILL.md | 596 ++++++++++++++++++ csrc/apis/attention.hpp | 12 +- csrc/apis/hyperconnection.hpp | 27 + csrc/jit/device_runtime.hpp | 5 + .../impls/smxx_fp8_fp4_paged_mqa_logits.hpp | 127 ++++ .../impls/sm120_fp8_paged_mqa_logits.cuh | 75 +++ tests/test_sm120_reference.py | 126 ++++ 10 files changed, 1316 insertions(+), 4 deletions(-) create mode 100644 .codex/skills/cpp/SKILL.md create mode 100644 .codex/skills/cpp/references/guidelines.md create mode 100644 .codex/skills/cpp/references/structure.md create mode 100644 .codex/skills/readable-cpp/SKILL.md create mode 100644 deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh create mode 100644 tests/test_sm120_reference.py diff --git a/.codex/skills/cpp/SKILL.md b/.codex/skills/cpp/SKILL.md new file mode 100644 index 0000000000..2c61ec1e86 --- /dev/null +++ b/.codex/skills/cpp/SKILL.md @@ -0,0 +1,111 @@ +--- +name: cpp +description: Comprehensive C/C++ programming reference covering everything from C11-C23 and C++11-C++23, system programming, CUDA GPU computing, debugging tools, Rust interop, and advanced topics. Use for: C/C++ questions, C/C++ interview preparation, modern language features, RAII/memory management, templates/generics, CUDA programming, system programming, debugging/profiling, performance optimization, cross-platform development, build systems, assembly, shell scripting, and any C/C++/CUDA development tasks. +--- + +# C/C++ Comprehensive Cheat Sheets (/cpp) + +Complete C/C++ development reference combining local documentation with live examples from cppcheatsheet.com, covering everything from basic syntax to advanced GPU programming and system-level development. + +## How It Works + +Help users write functional, correct C/C++ code and answer C/C++ questions by fetching proven patterns and examples from cppcheatsheet.com. + +When a user asks a C/C++ question or wants to write C/C++ code: + +1. Look up the relevant topic(s) in [Structure](references/structure.md) to find the matching URL(s) +2. **Always fetch** the URL(s) using WebFetch to get real examples and patterns from the site +3. Use the fetched content to: + - **Write code**: Apply the patterns to produce functional, correct code that solves the user's task + - **Answer questions**: Provide thorough explanations backed by the examples and information from the site +4. Follow the [Guidelines](references/guidelines.md) for code quality + +## Key Principle + +**Functionality first, cleanliness second.** The code must work correctly and handle the task properly. Fetching from cppcheatsheet.com ensures solutions use battle-tested patterns rather than guessing. The site contains rich examples covering edge cases, common pitfalls, and practical usage that go beyond basic documentation. + +## Coverage Areas + +### C/C++ Interview Preparation +**Interview Cheatsheet:** Curated C and C++ interview questions grouped by topic (memory management, move semantics, templates, STL, concurrency, debugging), each deep-linked to the section of the cheat sheet that answers it. Use for interview prep or quick topic review. + +### Modern C Programming (C11-C23) +**Core Language:** Syntax, types, memory management, preprocessor macros +**GNU Extensions:** Compiler-specific features and optimizations +**Build Systems:** Makefiles, compilation, linking +**Assembly:** X86 assembly integration and inline assembly + +### Modern C++ Programming (C++11-C++23) +**Core Features:** RAII, templates, STL containers, iterators, algorithms +**Modern Standards:** Move semantics, constexpr, lambdas, concepts, coroutines, modules +**Memory Management:** Smart pointers, resource management, optimization (RVO) +**Build Systems:** CMake, package management, cross-platform builds + +### System Programming +**Process Management:** POSIX processes, signals, process communication +**File Systems:** File I/O, directory operations, filesystem monitoring +**Networking:** Sockets, protocols, network programming patterns +**Threading:** Multithreading, synchronization, parallel programming +**IPC:** Inter-process communication, shared memory, message queues + +### CUDA Programming +**GPU Computing:** CUDA kernels, memory hierarchy, performance optimization +**Advanced CUDA:** libcu++, Thrust library, cooperative groups +**Multi-GPU:** GPU-GPU communication, hardware topology, NVSHMEM +**Async Programming:** CUDA pipelines, memory visibility, asynchronous execution + +### Debugging & Profiling +**Debug Tools:** GDB debugging, Valgrind memory analysis, sanitizers, binary inspection (nm, readelf, objdump, otool) +**Performance:** Perf profiling, tracing, performance optimization +**GPU Debugging:** Nsight Systems, CUDA debugging and profiling + +### System Tools & Automation +**Shell Scripting:** Bash programming, system administration +**System Tools:** OS utilities, networking tools, hardware inspection +**Service Management:** Systemd, process management, system monitoring + +### Cross-Language Development +**Rust Interop:** Rust for C++ developers, FFI, memory safety comparison +**Language Bridging:** C/C++ integration, foreign function interfaces + +### Advanced Topics +**Blog Content:** Deep-dive articles, RDMA networking, GPU-initiated communication +**Low-Level Programming:** Hardware interfaces, performance tuning +**Architecture:** System design, scalable applications + +## References + +For detailed information, I can access: +- **[Structure](references/structure.md)** - Complete topic-to-URL reference map +- **[Guidelines](references/guidelines.md)** - Code quality and best practices + +## Examples + +### C/C++ Interview Prep +- "What should I review for a C++ interview?" → Fetch https://cppcheatsheet.com/notes/interview/index.html and walk the reader through the topic groups +- "Common C++ interview questions on smart pointers" → Fetch https://cppcheatsheet.com/notes/interview/index.html and then drill into https://cppcheatsheet.com/notes/cpp/cpp_smartpointers.html for detailed answers + +### C/C++ Core +- "How do smart pointers work?" → Fetch https://cppcheatsheet.com/notes/cpp/cpp_smartpointers.html and explain with the site's examples +- "How does RAII work in C++?" → Fetch https://cppcheatsheet.com/notes/cpp/cpp_raii.html and explain with practical examples +- "How to use STL containers?" → Fetch https://cppcheatsheet.com/notes/cpp/cpp_container.html and explain with practical examples +- "Template metaprogramming techniques" → Fetch https://cppcheatsheet.com/notes/cpp/cpp_template.html and explain with practical examples + +### System Programming +- "POSIX socket server in C" → Fetch https://cppcheatsheet.com/notes/os/os_socket.html, use the patterns to write a working server +- "Multithreading with std::thread" → Fetch https://cppcheatsheet.com/notes/os/os_thread.html and explain with practical examples +- "Signal handling and process management" → Fetch https://cppcheatsheet.com/notes/os/os_signal.html and explain with practical examples + +### CUDA & GPU Programming +- "Write a CUDA kernel" → Fetch https://cppcheatsheet.com/notes/cuda/cuda_basics.html, use the patterns to write working GPU code +- "CUDA memory hierarchy and optimization" → Fetch https://cppcheatsheet.com/notes/cuda/cuda_memory_visibility.html and explain with practical examples +- "Multi-GPU communication with NCCL" → Fetch https://cppcheatsheet.com/notes/cuda/cuda_nccl.html and explain with practical examples + +### Debugging & Tools +- "Debug a segfault with GDB" → Fetch https://cppcheatsheet.com/notes/debug/gdb.html and explain with practical examples +- "Valgrind memory leak detection" → Fetch https://cppcheatsheet.com/notes/debug/valgrind.html and explain with practical examples +- "Performance profiling with perf" → Fetch https://cppcheatsheet.com/notes/debug/perf.html and explain with practical examples + +### Build & Development +- "CMake cross-platform build systems" → Fetch https://cppcheatsheet.com/notes/cpp/cpp_cmake.html and explain with practical examples +- "Makefile patterns and best practices" → Fetch https://cppcheatsheet.com/notes/c/make.html and explain with practical examples \ No newline at end of file diff --git a/.codex/skills/cpp/references/guidelines.md b/.codex/skills/cpp/references/guidelines.md new file mode 100644 index 0000000000..9af3ec1a39 --- /dev/null +++ b/.codex/skills/cpp/references/guidelines.md @@ -0,0 +1,152 @@ +# C/C++ Comprehensive Development Guidelines + +Always fetch relevant examples from cppcheatsheet.com first to ensure correctness, then apply these guidelines when writing code. + +## Modern C Programming (C11-C23) + +### Code Quality & Standards +- Follow C11/C17/C23 standards for modern C development +- Use static analysis tools: clang-static-analyzer, cppcheck, scan-build +- Implement proper error handling with return codes and errno +- Use const correctness and restrict pointers for optimization +- Leverage compound literals and designated initializers (C99+) + +### Memory Management +- Always pair malloc/free, avoid memory leaks +- Use valgrind and AddressSanitizer for memory debugging +- Consider memory alignment for performance-critical code +- Use restrict keyword for pointer aliasing optimization + +### Build Systems & Portability +- Write portable Makefiles with proper dependency tracking +- Use feature test macros for cross-platform compatibility +- Leverage GNU extensions judiciously with fallbacks + +## Modern C++ Programming (C++11-C++23) + +### Resource Management (RAII) +- Always use RAII for resource management +- Prefer smart pointers: unique_ptr, shared_ptr, weak_ptr +- Use custom deleters for non-standard resources +- Implement move semantics for expensive-to-copy objects + +### Modern Language Features +- Use auto for type deduction, but maintain readability +- Leverage constexpr for compile-time computation +- Use structured bindings (C++17+) for multiple return values +- Implement concepts (C++20+) for better template constraints +- Use coroutines (C++20+) for asynchronous programming +- Adopt modules (C++20+) for better build times and encapsulation + +### Template Programming +- Use SFINAE and concepts for template constraints +- Prefer if constexpr over template specialization where applicable +- Use fold expressions (C++17+) for variadic templates +- Implement type traits for generic programming + +### Performance Optimization +- Profile before optimizing (use perf, vtune, or similar) +- Understand memory hierarchy and cache-friendly algorithms +- Use move semantics and perfect forwarding +- Leverage Return Value Optimization (RVO) and copy elision +- Consider noexcept specifications for optimization opportunities + +## System Programming + +### POSIX Best Practices +- Always check return values from system calls +- Handle EINTR properly in system call loops +- Use signal-safe functions in signal handlers +- Implement proper cleanup with RAII wrappers + +### Multithreading & Concurrency +- Use std::thread and C++11+ concurrency primitives +- Avoid data races with proper synchronization +- Prefer lock-free algorithms when appropriate +- Use thread-local storage for per-thread data +- Consider memory ordering for atomic operations + +### Network Programming +- Handle partial reads/writes in socket programming +- Use non-blocking I/O with proper error handling +- Implement timeout mechanisms for robust networking +- Consider endianness for network protocols + +## CUDA Programming + +### Performance Guidelines +- Optimize memory access patterns (coalesced access) +- Use shared memory effectively for data reuse +- Minimize divergence in warp execution +- Consider occupancy vs. resource usage tradeoffs + +### Memory Management +- Use CUDA unified memory judiciously +- Understand memory hierarchy: global, shared, constant, texture +- Implement proper error checking with cudaGetLastError() +- Use CUDA streams for overlapping computation and communication + +### Advanced CUDA +- Leverage cooperative groups for flexible thread cooperation +- Use libcu++ for standard library features on GPU +- Implement multi-GPU communication with NVSHMEM or NCCL +- Consider GPU-initiated communication for advanced patterns + +## Debugging & Profiling + +### Debug Tools Usage +- Use GDB with proper debug symbols (-g flag) +- Leverage Valgrind for memory error detection +- Use sanitizers: AddressSanitizer, ThreadSanitizer, UBSan +- Implement proper logging with different verbosity levels + +### Performance Analysis +- Use perf for CPU profiling and cache analysis +- Implement custom trace points for application-specific profiling +- Use Nsight Systems/Compute for CUDA application analysis +- Profile both CPU and GPU components in heterogeneous applications + +## Build Systems & DevOps + +### CMake Best Practices +- Use modern CMake (3.15+) with target-based design +- Properly handle dependencies with find_package or FetchContent +- Implement proper install rules and export configurations +- Use generator expressions for configuration-specific settings + +### Cross-Platform Development +- Write portable code with proper feature detection +- Use standard library features over platform-specific ones +- Handle different compiler behaviors (GCC, Clang, MSVC) +- Test on multiple platforms and architectures + +## Security & Safety + +### Memory Safety +- Use static analysis tools to catch buffer overflows +- Implement bounds checking in debug builds +- Consider using safer alternatives (std::array vs C arrays) +- Use string handling functions that prevent overflows + +### Code Review & Quality +- Implement automated testing with appropriate coverage +- Use continuous integration for multi-platform testing +- Enforce coding standards with clang-format and linters +- Document public APIs with clear contracts and examples + +## Cross-Language Integration + +### C/C++ Interoperability +- Use extern "C" linkage for C++ code called from C +- Handle exceptions at C++ boundaries when interfacing with C +- Use proper name mangling considerations + +### Rust Integration +- Understand ownership model differences between C++ and Rust +- Use proper FFI patterns for safe interoperation +- Handle error propagation between language boundaries +- Consider memory safety implications at interfaces + +## Related Documentation + +This skill is based on the comprehensive C/C++ reference available at https://cppcheatsheet.com/ which includes working code snippets, performance benchmarks, real-world patterns, and integration guides. The reference is continuously updated with the latest C/C++ features and best practices. \ No newline at end of file diff --git a/.codex/skills/cpp/references/structure.md b/.codex/skills/cpp/references/structure.md new file mode 100644 index 0000000000..be32cc2faa --- /dev/null +++ b/.codex/skills/cpp/references/structure.md @@ -0,0 +1,89 @@ +# C/C++ Topics Reference Map + +## C/C++ Interview Cheatsheet +- [C/C++ Interview Cheatsheet](https://cppcheatsheet.com/notes/interview/index.html) - Curated interview questions, each deep-linked to the section that answers it + +## Modern C Programming +- [C Basics](https://cppcheatsheet.com/notes/c/c_basic.html) - Core C language features from C11 to C23 +- [Memory](https://cppcheatsheet.com/notes/c/c_memory.html) - Memory management and allocation +- [Preprocessor & GNU Extensions](https://cppcheatsheet.com/notes/c/c_macro.html) - Macros and GNU extensions +- [Makefile](https://cppcheatsheet.com/notes/c/make.html) - Build system configuration +- [X86 Assembly](https://cppcheatsheet.com/notes/c/asm.html) - Assembly language integration + +## Modern C++ Programming +- [C++ Basics](https://cppcheatsheet.com/notes/cpp/cpp_basic.html) - Modern C++ features from C++11 to C++23 +- [Resource Management](https://cppcheatsheet.com/notes/cpp/cpp_raii.html) - RAII and resource management +- [String](https://cppcheatsheet.com/notes/cpp/cpp_string.html) - String handling and operations +- [Container](https://cppcheatsheet.com/notes/cpp/cpp_container.html) - STL containers and data structures +- [Iterator](https://cppcheatsheet.com/notes/cpp/cpp_iterator.html) - Iterator patterns and usage +- [Template](https://cppcheatsheet.com/notes/cpp/cpp_template.html) - Generic programming and metaprogramming +- [Casting](https://cppcheatsheet.com/notes/cpp/cpp_casting.html) - Type casting and conversions +- [Constexpr](https://cppcheatsheet.com/notes/cpp/cpp_constexpr.html) - Compile-time computation +- [Lambda](https://cppcheatsheet.com/notes/cpp/cpp_lambda.html) - Lambda expressions and functional programming +- [Concepts](https://cppcheatsheet.com/notes/cpp/cpp_concepts.html) - C++20 concepts and constraints +- [Requires](https://cppcheatsheet.com/notes/cpp/cpp_requires.html) - Requirements and constraints +- [Time](https://cppcheatsheet.com/notes/cpp/cpp_time.html) - Time and chrono utilities +- [Smart Pointers](https://cppcheatsheet.com/notes/cpp/cpp_smartpointers.html) - Modern memory management +- [Return Value Optimization](https://cppcheatsheet.com/notes/cpp/cpp_rvo.html) - RVO and move semantics +- [Algorithm](https://cppcheatsheet.com/notes/cpp/cpp_algorithm.html) - STL algorithms and functional patterns +- [Coroutine](https://cppcheatsheet.com/notes/cpp/cpp_coroutine.html) - C++20 coroutines +- [Modules](https://cppcheatsheet.com/notes/cpp/cpp_modules.html) - C++20 modules system +- [CMake](https://cppcheatsheet.com/notes/cpp/cpp_cmake.html) - Build system and project configuration + +## System Programming +- [Process](https://cppcheatsheet.com/notes/os/os_process.html) - Process management and control +- [File](https://cppcheatsheet.com/notes/os/os_file.html) - File I/O and filesystem operations +- [Signal](https://cppcheatsheet.com/notes/os/os_signal.html) - Signal handling and IPC +- [Socket](https://cppcheatsheet.com/notes/os/os_socket.html) - Network programming and sockets +- [Thread](https://cppcheatsheet.com/notes/os/os_thread.html) - Threading and synchronization +- [IPC](https://cppcheatsheet.com/notes/os/os_ipc.html) - Inter-process communication + +## CUDA Programming +- [CUDA Basics](https://cppcheatsheet.com/notes/cuda/cuda_basics.html) - GPU programming fundamentals +- [CUDA C++ (libcu++)](https://cppcheatsheet.com/notes/cuda/cuda_cpp.html) - Modern CUDA C++ features +- [Thrust](https://cppcheatsheet.com/notes/cuda/cuda_thrust.html) - High-level parallel algorithms +- [Cooperative Groups](https://cppcheatsheet.com/notes/cuda/cuda_coop_groups.html) - Thread cooperation patterns +- [Memory Visibility](https://cppcheatsheet.com/notes/cuda/cuda_memory_visibility.html) - Memory models and synchronization +- [CUDA Graph](https://cppcheatsheet.com/notes/cuda/cuda_graph.html) - Graph-based execution and optimization +- [CUDA Quantization](https://cppcheatsheet.com/notes/cuda/cuda_quantization.html) - Quantization techniques on GPU +- [Pipelines](https://cppcheatsheet.com/notes/cuda/cuda_pipelines.html) - Asynchronous execution pipelines +- [NCCL](https://cppcheatsheet.com/notes/cuda/cuda_nccl.html) - Multi-GPU collective communication +- [GPU-GPU Communication](https://cppcheatsheet.com/notes/cuda/cuda_ipc.html) - Multi-GPU programming +- [Hardware Topology](https://cppcheatsheet.com/notes/cuda/cuda_hwloc.html) - Hardware locality and topology + +## Blog +- [Building NVSHMEM from Scratch](https://cppcheatsheet.com/notes/blog/nvshmem.html) - GPU-initiated networking deep dive + +## Bash & System Tools +- [Bash](https://cppcheatsheet.com/notes/tools/bash.html) - Shell scripting and command-line tools +- [Operating System](https://cppcheatsheet.com/notes/tools/os.html) - System administration tools +- [Network](https://cppcheatsheet.com/notes/tools/net.html) - Network configuration and debugging +- [Hardware](https://cppcheatsheet.com/notes/tools/hardware.html) - Hardware inspection and monitoring +- [GPU](https://cppcheatsheet.com/notes/tools/gpu.html) - GPU tools and utilities +- [Systemd](https://cppcheatsheet.com/notes/tools/systemd.html) - Service management and system control + +## Debugging & Profiling +- [GDB](https://cppcheatsheet.com/notes/debug/gdb.html) - GNU debugger and debugging techniques +- [Valgrind](https://cppcheatsheet.com/notes/debug/valgrind.html) - Memory analysis and leak detection +- [Sanitizers](https://cppcheatsheet.com/notes/debug/sanitizers.html) - Address, thread, and undefined behavior sanitizers +- [Tracing](https://cppcheatsheet.com/notes/debug/tracing.html) - System and application tracing +- [Perf](https://cppcheatsheet.com/notes/debug/perf.html) - Performance analysis and profiling +- [Nsight Systems](https://cppcheatsheet.com/notes/debug/nsight.html) - NVIDIA profiling tools +- [Binary Tools](https://cppcheatsheet.com/notes/debug/binary_tools.html) - Binary inspection with nm, readelf, objdump, otool, ldd, strings, size + +## Rust for C++ Developers +- [Basics](https://cppcheatsheet.com/notes/rust/rust_basic.html) - Rust fundamentals for C++ developers +- [Ownership & Borrowing](https://cppcheatsheet.com/notes/rust/rust_ownership.html) - Memory safety without GC +- [RAII & Drop](https://cppcheatsheet.com/notes/rust/rust_raii.html) - Resource management patterns +- [Strings](https://cppcheatsheet.com/notes/rust/rust_string.html) - String handling and text processing +- [Collections](https://cppcheatsheet.com/notes/rust/rust_container.html) - Data structures and containers +- [Iterators](https://cppcheatsheet.com/notes/rust/rust_iterator.html) - Functional iteration patterns +- [Traits & Generics](https://cppcheatsheet.com/notes/rust/rust_traits.html) - Generic programming and interfaces +- [Casting](https://cppcheatsheet.com/notes/rust/rust_casting.html) - Type conversions and casting +- [Const Functions](https://cppcheatsheet.com/notes/rust/rust_constfn.html) - Compile-time computation +- [Closures](https://cppcheatsheet.com/notes/rust/rust_closure.html) - Anonymous functions and captures +- [Smart Pointers](https://cppcheatsheet.com/notes/rust/rust_smartptr.html) - Reference counting and ownership +- [Error Handling](https://cppcheatsheet.com/notes/rust/rust_error.html) - Result types and error propagation +- [Threads](https://cppcheatsheet.com/notes/rust/rust_thread.html) - Safe concurrency and parallelism +- [Modules](https://cppcheatsheet.com/notes/rust/rust_modules.html) - Code organization and visibility +- [FFI & C++ Bindings](https://cppcheatsheet.com/notes/rust/rust_ffi.html) - Interoperability with C/C++ diff --git a/.codex/skills/readable-cpp/SKILL.md b/.codex/skills/readable-cpp/SKILL.md new file mode 100644 index 0000000000..1f56f88374 --- /dev/null +++ b/.codex/skills/readable-cpp/SKILL.md @@ -0,0 +1,596 @@ +--- +name: readable-cpp +description: Readable C/C++/Rust/CUDA code rules inspired by The Art of Readable Code. Use when writing, reviewing, or refactoring C, C++, Rust, or CUDA code. Enforces short functions, flat control flow, clear naming, readable structure, and idiomatic patterns. +--- + +# Readable C/C++/Rust/CUDA Rules (/readable-cpp) + +Apply these rules when writing, reviewing, or refactoring C, C++, Rust, or CUDA code. Inspired by *The Art of Readable Code* by Dustin Boswell and Trevor Foucher. + +**Core principle: Code should be easy to understand.** The time it takes someone else (or future you) to understand the code is the ultimate metric. + +## 1. Keep Functions Short and Focused + +- A function should do **one thing**. If you can describe what it does with "and", split it. +- Aim for functions that fit on one screen (~15-25 lines). If it's longer, extract sub-tasks. +- Each function should operate at a **single level of abstraction** — don't mix high-level logic with low-level details in the same function. + +## 2. Flatten Control Flow — No Deep Nesting + +- **Never nest more than 2 levels deep.** If you have a loop inside a loop, or an `if` inside a loop inside an `if`, extract the inner block into a helper function with a descriptive name. +- Use **early returns / guard clauses** to handle edge cases at the top, keeping the main logic flat. +- Prefer `continue` or `break` to skip iterations rather than wrapping the body in a conditional. +- Replace complex conditionals with well-named helper functions or variables that explain the intent. + +```cpp +// Bad: nested and hard to follow +for (auto& user : users) { + if (user.is_active()) { + for (auto& order : user.orders()) { + if (order.is_pending()) { + process(order); + } + } + } +} + +// Good: flat, each function name explains what it does +auto active_users = get_active_users(users); +for (auto& user : active_users) { + process_pending_orders(user.orders()); +} +``` + +## 3. Name Things Clearly + +- **Pack information into names.** Use specific, concrete words — `fetch_page` not `get`, `num_retries` not `n`. +- **Avoid generic names** like `tmp`, `data`, `result`, `val`, `info`, `handle` — unless the scope is tiny (2-3 lines). +- **Use names that can't be misconstrued.** If a range is inclusive, say `max_items` not `limit`. If a boolean, use `is_`, `has_`, `should_`, `can_` prefixes. +- **Match the name length to the scope.** Short names for small scopes, descriptive names for wide scopes. +- **Don't use abbreviations** unless they're universally understood (`num`, `max`, `min`, `err` are fine; `svc_mgr_cfg` is not). + +## 4. Make Control Flow Easy to Follow + +- Put the **changing/interesting value on the left** side of comparisons: `if (length > 10)` not `if (10 < length)`. +- Order `if/else` blocks: **positive case first**, simpler case first, or the more interesting case first. +- Minimize the number of variables the reader has to track. Reduce the **mental footprint** of each block. +- Avoid deeply nested ternary operators — if it's not immediately obvious, use an `if/else`. + +## 5. Break Down Giant Expressions + +- Use **explaining variables** to break complex expressions into named pieces. +- Use **summary variables** to capture a long expression that's used more than once. +- Apply **De Morgan's laws** to simplify negated boolean expressions. + +```cpp +// Bad +if (!(age >= 18 && has_id && !is_banned)) { + deny(); +} + +// Good +bool is_eligible = age >= 18 && has_id && !is_banned; +if (!is_eligible) { + deny(); +} +``` + +## 6. Extract Unrelated Subproblems + +- If a block of code is solving a **subproblem unrelated to the main goal** of the function, extract it. +- The helper function should be **pure and self-contained** — it shouldn't need to know about the calling context. +- This is the single most effective way to improve readability: separate *what* you're doing from *how*. + +## 7. One Task at a Time + +- Each section of code should do **one task**. If a function is doing parsing AND validation AND transformation, split them into separate steps. +- List the tasks a function does. If there's more than one, reorganize so each task is in its own block or function. + +## 8. Reduce Variable Scope + +- **Declare variables close to where they're used.** Don't declare at the top of a function if it's only used 30 lines later. +- **Minimize the "live time" of a variable** — the fewer lines between its assignment and last use, the easier it is to follow. +- **Prefer write-once variables.** Variables that are assigned once and never modified are easier to reason about. +- **Eliminate unnecessary variables.** If a variable is used only once and doesn't clarify anything, inline it. + +## 9. No Magic Numbers or Strings + +- Replace **magic numbers and strings** with named constants: `if (retries > MAX_RETRIES)` not `if (retries > 3)`. +- If a value has meaning, give it a name. The name documents the intent. +- Group related constants together. + +## 10. Fewer Function Arguments + +- Aim for **3 or fewer arguments** per function. More than that is a smell. +- Group related arguments into a **struct, class, or tuple**. +- If a function needs many config-like options, pass a single config/options object. +- Boolean flag arguments are a sign the function does two things — split it instead. + +## 11. Consistency + +- If the codebase does something one way, **do it the same way**. Don't mix styles. +- Consistent naming patterns, consistent structure, consistent error handling. +- When joining an existing codebase, **match the existing conventions** even if you'd prefer a different style. +- Surprise is the enemy of readability — predictable code is readable code. + +## 12. Write Less Code + +- The best code is **no code at all**. Question whether a feature is truly needed before implementing. +- **Don't over-engineer.** Solve the problem at hand, not hypothetical future problems. +- Remove dead code. Commented-out code is dead code. +- Use standard libraries before writing custom solutions. + +## 13. Comments: Explain Why, Not What + +- Don't comment **what** the code does — the code already says that. Comment **why** it does it. +- Comment **flaws and workarounds**: `// TODO:`, `// HACK:`, `// XXX:` with explanation. +- Comment **surprising behavior** or non-obvious decisions — things where a reader would ask "why?". +- **Don't comment bad code — rewrite it.** If you need a comment to explain what a block does, extract it into a well-named function instead. + +## 14. Design Code to Survive Auto-Formatting + +- Write code that looks good **after** the auto-formatter runs. If a chained expression or repeated pattern would be broken across 4+ lines by the formatter, extract a helper function instead. +- **Prefer one-line helper calls** over long inline chains that the formatter will expand vertically. +- The formatter is your reader's first impression. Run it *before* committing — if the result looks ugly, that's a signal to refactor, not to disable the formatter. + +```rust +// Bad: rustfmt expands this to 4 lines per field — noisy and repetitive +fn from_dict(cfg: &Bound<'_, PyDict>) -> PyResult { + Ok(Self { + rom: cfg.get_item("rom")?.ok_or_else(|| missing("rom"))?.extract()?, + // ... each field becomes 4 lines after rustfmt + }) +} + +// Good: extract a helper so each field stays one clean line +fn get_required(cfg: &Bound<'_, PyDict>, key: &str) -> PyResult { + cfg.get_item(key)? + .ok_or_else(|| PyKeyError::new_err(key.to_string()))? + .extract() +} + +fn from_dict(cfg: &Bound<'_, PyDict>) -> PyResult { + Ok(Self { + rom: get_required(cfg, "rom")?, + actions: get_required(cfg, "actions")?, + }) +} +``` + +```cpp +// Bad: clang-format wraps this into a hard-to-scan block +auto result = container.find(key)->second.get_value().transform(func).value_or(default_val); + +// Good: name the intermediate step +auto& entry = container.find(key)->second; +auto result = entry.get_value().transform(func).value_or(default_val); +``` + +--- + +# C-Specific Rules + +## 15. RAII-Like Patterns with goto Cleanup + +- In C, use the **goto cleanup pattern** for resource management — allocate at the top, clean up at a single labeled block at the bottom. +- Never scatter `free()` calls across multiple return paths. A single cleanup section is easier to audit. +- Use `__attribute__((cleanup))` (GCC/Clang) when available for automatic cleanup. + +```c +// Good: single cleanup path +int process_file(const char *path) { + int ret = -1; + FILE *fp = fopen(path, "r"); + if (!fp) return -1; + + char *buf = malloc(BUF_SIZE); + if (!buf) goto cleanup_file; + + // ... do work ... + ret = 0; + +cleanup_buf: + free(buf); +cleanup_file: + fclose(fp); + return ret; +} +``` + +## 16. Use `const` Liberally + +- Mark pointers `const` when the function doesn't modify the pointed-to data: `const char *msg`. +- Mark local variables `const` when they don't change after initialization. +- This documents intent and helps the compiler catch mistakes. + +## 17. Prefer Sized Types for Data Structures + +- Use `` types (`uint32_t`, `int64_t`) for data that crosses boundaries (files, network, hardware). +- Use `size_t` for sizes and counts, `ptrdiff_t` for pointer differences. +- Use `int` and `unsigned` for simple loop counters and local arithmetic. + +## 18. Defensive Macro Hygiene + +- Wrap macro bodies in `do { ... } while(0)` for statement-like macros. +- Parenthesize all macro parameters: `#define SQUARE(x) ((x) * (x))`. +- Prefer `static inline` functions over macros when possible (type safety, debuggability). +- Use `_Generic` (C11) for type-safe "overloading" instead of macro tricks. + +--- + +# C++-Specific Rules + +## 19. Use RAII for All Resources + +- Every resource (memory, file handles, locks, sockets) should be owned by an RAII object. +- Use `std::unique_ptr` for exclusive ownership, `std::shared_ptr` only when shared ownership is genuinely needed. +- Write custom RAII wrappers for non-standard resources (e.g., C library handles). +- Never use raw `new`/`delete` in application code — let smart pointers and containers handle it. + +## 20. Prefer Value Semantics and Move + +- Pass small objects by value, large objects by `const&`. +- Return objects by value — rely on RVO/NRVO and move semantics. +- Implement move constructors/assignment for types that own resources. +- Use `std::move` only when you truly want to transfer ownership — don't `std::move` from things you'll use again. + +## 21. Use Modern C++ Over C Idioms + +- Use `std::array` over C arrays, `std::string` over `char*`, `std::vector` over `malloc`/`realloc`. +- Use `std::optional` over sentinel values, `std::variant` over type-unsafe unions. +- Use range-based `for` loops: `for (const auto& item : container)`. +- Use structured bindings (C++17): `auto [key, value] = *map.begin();`. +- Use `std::format` (C++20) or `fmt::format` over `sprintf` / string concatenation. + +## 22. Templates: Keep It Simple + +- Use concepts (C++20) to constrain templates — errors become readable. +- Prefer `if constexpr` over SFINAE when possible. +- Don't write template metaprogramming unless the benefit is clear and the team can maintain it. +- A non-template solution that's slightly less generic is often better than a template solution nobody understands. + +## 23. Use `constexpr` and `const` Aggressively + +- Mark functions `constexpr` when they can be evaluated at compile time. +- Use `constexpr` variables instead of `#define` for constants. +- Use `const` on member functions that don't modify state. +- `consteval` (C++20) for functions that *must* be compile-time evaluated. + +## 24. Error Handling: Pick One Pattern + +- Use exceptions for truly exceptional conditions, `std::expected` (C++23) or `std::optional` for expected failures. +- Don't mix error codes and exceptions in the same layer. +- If using exceptions, make them specific — derive from `std::runtime_error`, not `std::exception`. +- Use `noexcept` on functions that cannot throw (destructors, move operations). + +--- + +# Rust-Specific Rules + +## 25. Embrace the Ownership Model + +- Don't fight the borrow checker — redesign your data flow instead. +- Prefer passing references (`&T`, `&mut T`) over cloning. Clone only when ownership transfer is genuinely needed. +- Use lifetimes explicitly only when the compiler can't infer them — don't annotate unnecessarily. +- Prefer `&str` over `String` in function parameters when you don't need ownership. + +## 26. Use Iterators and Combinators + +- Prefer iterator chains (`.iter().filter().map().collect()`) over manual loops with indices. +- Use `for item in &collection` instead of `for i in 0..collection.len()`. +- Use `enumerate()`, `zip()`, `chain()`, `chunks()` — the iterator API is rich. +- Avoid `.unwrap()` in production code — use `?`, `unwrap_or`, `unwrap_or_else`, or pattern matching. + +```rust +// Bad: manual indexing +let mut names = Vec::new(); +for i in 0..users.len() { + if users[i].is_active { + names.push(users[i].name.clone()); + } +} + +// Good: idiomatic iterator chain +let names: Vec<_> = users.iter() + .filter(|u| u.is_active) + .map(|u| u.name.clone()) + .collect(); +``` + +## 27. Use Enums and Pattern Matching + +- Use `enum` with data variants instead of class hierarchies or tagged unions. +- Use `match` exhaustively — the compiler ensures you handle all cases. +- Use `if let` / `while let` for single-variant matching instead of full `match`. +- Prefer `Result` over panicking — make errors part of the type signature. + +## 28. Leverage the Type System + +- Use **newtype wrappers** (`struct UserId(u64)`) to prevent mixing up same-typed values. +- Use `Option` instead of sentinel values or null pointers. +- Use `#[must_use]` on functions whose return values shouldn't be ignored. +- Prefer `From`/`Into` traits for type conversions over manual conversion functions. + +## 29. Module Organization + +- Keep `pub` surfaces small — expose only what's needed. +- Use `pub(crate)` for crate-internal visibility instead of full `pub`. +- Group related types and functions in modules — one concept per module. +- Re-export key types at the crate root for ergonomic imports. + +--- + +# CUDA-Specific Rules + +## 30. Name Kernels and Device Functions Clearly + +- Kernel names should describe **what** they compute, not that they're kernels: `reduce_sum` not `kernel1` or `myKernel`. +- Use a consistent naming convention to distinguish execution spaces: e.g., `reduce_sum_kernel` for `__global__`, `warp_reduce` for `__device__` helpers. +- Name grid/block dimension variables descriptively: `threads_per_block`, `num_blocks` not `tpb`, `nb`, or bare `256`. + +```cuda +// Bad: opaque names, magic numbers +__global__ void k1(float *a, float *b, int n) { + int i = blockIdx.x * 256 + threadIdx.x; + if (i < n) b[i] = a[i] * 2.0f; +} +k1<<<(n+255)/256, 256>>>(d_in, d_out, n); + +// Good: clear intent, named constants +constexpr int THREADS_PER_BLOCK = 256; + +__global__ void scale_kernel(const float *input, float *output, + float scale_factor, int num_elements) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + output[idx] = input[idx] * scale_factor; + } +} + +const int num_blocks = (num_elements + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; +scale_kernel<<>>(d_input, d_output, 2.0f, num_elements); +``` + +## 31. Separate Host Logic from Device Logic + +- Keep **host orchestration** (memory allocation, transfers, kernel launches, synchronization) in separate functions from **device computation** (kernels and device helpers). +- Don't mix `cudaMalloc`/`cudaMemcpy` with application logic — wrap them in RAII classes or helper functions. +- Use a clear file organization: consider separating `.cu` kernel files from `.cpp` host logic files, or at minimum group host and device code into clearly labeled sections. + +```cpp +// Good: RAII wrapper hides allocation/deallocation +template +class DeviceBuffer { + T *ptr_ = nullptr; + size_t size_ = 0; +public: + explicit DeviceBuffer(size_t count) : size_(count) { + check_cuda(cudaMalloc(&ptr_, count * sizeof(T))); + } + ~DeviceBuffer() { cudaFree(ptr_); } + + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + DeviceBuffer(DeviceBuffer&& o) noexcept : ptr_(o.ptr_), size_(o.size_) { o.ptr_ = nullptr; } + + T *get() { return ptr_; } + const T *get() const { return ptr_; } + size_t size() const { return size_; } + + void copy_from_host(const T *host_data) { + check_cuda(cudaMemcpy(ptr_, host_data, size_ * sizeof(T), cudaMemcpyHostToDevice)); + } + void copy_to_host(T *host_data) const { + check_cuda(cudaMemcpy(host_data, ptr_, size_ * sizeof(T), cudaMemcpyDeviceToHost)); + } +}; +``` + +## 32. Always Check CUDA Errors + +- **Check every CUDA API call.** Silent failures are the #1 source of hard-to-debug CUDA issues. +- Use a `check_cuda` macro or inline function — not raw `if` blocks after every call. +- Check errors after kernel launches with `cudaGetLastError()` + `cudaDeviceSynchronize()` during development. +- In release builds, at minimum check allocations and memcpy — these are the most likely to fail at runtime. + +```cuda +// Good: concise, catches file/line info +inline void check_cuda(cudaError_t err, const char *file, int line) { + if (err != cudaSuccess) { + fprintf(stderr, "CUDA error at %s:%d — %s\n", + file, line, cudaGetErrorString(err)); + exit(EXIT_FAILURE); + } +} +#define check_cuda(err) check_cuda((err), __FILE__, __LINE__) + +// Usage +check_cuda(cudaMalloc(&d_ptr, size)); +my_kernel<<>>(d_ptr, n); +check_cuda(cudaGetLastError()); +check_cuda(cudaDeviceSynchronize()); +``` + +## 33. Make Thread Indexing Obvious + +- Compute the global thread index **once** at the top of the kernel and store it in a clearly named variable. +- Use **early return** for out-of-bounds threads — don't wrap the entire kernel body in an `if`. +- For 2D/3D grids, name dimensions explicitly: `row`, `col`, `depth` — not `x`, `y`, `z`. + +```cuda +// Bad: index computed inline, entire body wrapped +__global__ void process(float *data, int width, int height) { + if (blockIdx.x * blockDim.x + threadIdx.x < width && + blockIdx.y * blockDim.y + threadIdx.y < height) { + int idx = (blockIdx.y * blockDim.y + threadIdx.y) * width + + (blockIdx.x * blockDim.x + threadIdx.x); + data[idx] = data[idx] * 2.0f; + } +} + +// Good: named indices, early return +__global__ void process(float *data, int width, int height) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y * blockDim.y + threadIdx.y; + if (col >= width || row >= height) return; + + const int idx = row * width + col; + data[idx] = data[idx] * 2.0f; +} +``` + +## 34. Document Shared Memory Usage + +- Declare shared memory with a **descriptive name** that indicates what it holds: `shared_tile` not `smem` or `s`. +- Add a brief comment explaining the **size** and **purpose** of shared memory when it's dynamically allocated (`extern __shared__`). +- Keep the shared memory lifecycle short — load, sync, compute, sync — and make each phase visually distinct. + +```cuda +// Good: clear phases, descriptive names +__global__ void tiled_matmul_kernel(const float *A, const float *B, + float *C, int N) { + __shared__ float tile_A[TILE_SIZE][TILE_SIZE]; + __shared__ float tile_B[TILE_SIZE][TILE_SIZE]; + + const int row = blockIdx.y * TILE_SIZE + threadIdx.y; + const int col = blockIdx.x * TILE_SIZE + threadIdx.x; + float accumulator = 0.0f; + + for (int tile_idx = 0; tile_idx < N / TILE_SIZE; ++tile_idx) { + // Phase 1: Load tiles from global memory + tile_A[threadIdx.y][threadIdx.x] = A[row * N + tile_idx * TILE_SIZE + threadIdx.x]; + tile_B[threadIdx.y][threadIdx.x] = B[(tile_idx * TILE_SIZE + threadIdx.y) * N + col]; + __syncthreads(); + + // Phase 2: Compute partial dot product from tiles + for (int k = 0; k < TILE_SIZE; ++k) { + accumulator += tile_A[threadIdx.y][k] * tile_B[k][threadIdx.x]; + } + __syncthreads(); + } + + C[row * N + col] = accumulator; +} +``` + +## 35. Keep Kernels Short — Extract Device Helpers + +- Apply the same "one function, one task" rule to kernels. If a kernel does loading, computing, and reducing, extract `__device__` helper functions. +- Use `__forceinline__ __device__` for small helpers that you want inlined without relying on compiler heuristics. +- This makes kernels easier to read, test (via unit-testing device functions), and reuse. + +```cuda +// Good: kernel reads like pseudocode, details in helpers +__forceinline__ __device__ +float warp_reduce_sum(float val) { + for (int offset = warpSize / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +__forceinline__ __device__ +float block_reduce_sum(float val) { + __shared__ float warp_sums[32]; + const int lane = threadIdx.x % warpSize; + const int warp_id = threadIdx.x / warpSize; + + val = warp_reduce_sum(val); + if (lane == 0) warp_sums[warp_id] = val; + __syncthreads(); + + val = (threadIdx.x < blockDim.x / warpSize) ? warp_sums[lane] : 0.0f; + if (warp_id == 0) val = warp_reduce_sum(val); + return val; +} + +__global__ void reduce_sum_kernel(const float *input, float *output, int n) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const float val = (idx < n) ? input[idx] : 0.0f; + + const float block_sum = block_reduce_sum(val); + if (threadIdx.x == 0) atomicAdd(output, block_sum); +} +``` + +## 36. Be Explicit About Memory Spaces + +- Use `const` on kernel parameters for read-only device pointers — documents intent and enables compiler optimizations. +- Use `__restrict__` when pointers don't alias — but add a comment explaining the non-aliasing guarantee. +- When using unified memory (`cudaMallocManaged`), comment the expected access pattern (host-only init, device-only compute, etc.) — the implicit page migration behavior is not obvious. +- Prefer explicit memory copies over unified memory in performance-critical paths — be explicit about data movement. + +```cuda +// Good: const + restrict with clear intent +__global__ void vector_add_kernel( + const float *__restrict__ a, // read-only, no alias with output + const float *__restrict__ b, // read-only, no alias with output + float *__restrict__ output, // write-only + int num_elements) +{ + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_elements) return; + output[idx] = a[idx] + b[idx]; +} +``` + +## 37. Synchronization: Make It Visible and Minimal + +- Place `__syncthreads()` on its own line, never buried inside a conditional branch that not all threads take — this is **undefined behavior** and hard to spot. +- Add a brief comment before each `__syncthreads()` stating what invariant it establishes: "all threads have loaded their tile", "partial sums are written to shared memory". +- Minimize synchronization points — restructure algorithms to reduce the number of barriers. +- For warp-level operations, prefer warp intrinsics (`__shfl_sync`, `__ballot_sync`) with explicit masks over `__syncthreads()`. + +## 38. Launch Configuration: Make It Readable + +- Wrap kernel launches in a **host function** that computes and names the launch parameters. +- Never hardcode grid/block dimensions at the call site — compute them from the problem size. +- For complex launch configurations, use a struct or helper to make the 2D/3D grid/block shape clear. +- Query device properties at startup rather than assuming specific hardware limits. + +```cuda +// Bad: magic numbers, unclear intent +foo<<<(n+127)/128, 128, 0, stream>>>(d_ptr, n); + +// Good: named, computed, self-documenting +void launch_scale_kernel(float *d_data, float factor, int n, cudaStream_t stream) { + constexpr int BLOCK_SIZE = 256; + const int grid_size = (n + BLOCK_SIZE - 1) / BLOCK_SIZE; + scale_kernel<<>>(d_data, factor, n); + check_cuda(cudaGetLastError()); +} +``` + +## 39. Streams and Async: Comment the Dependency Graph + +- When using multiple CUDA streams, add a comment block showing the **dependency graph** — which operations must complete before others begin. +- Name streams after their purpose: `compute_stream`, `transfer_stream` — not `s1`, `s2`. +- Group related async operations visually and separate independent pipelines with blank lines. +- Always synchronize before reading results on the host — make the sync point explicit and commented. + +```cuda +// Good: dependency graph documented, streams named by purpose +// Dependency graph: +// upload (transfer_stream) --> compute (compute_stream) --> download (transfer_stream) +// Event 'upload_done' gates compute start. +// Event 'compute_done' gates download start. + +cudaStream_t transfer_stream, compute_stream; +cudaEvent_t upload_done, compute_done; + +// Stage 1: async upload +cudaMemcpyAsync(d_input, h_input, size, cudaMemcpyHostToDevice, transfer_stream); +cudaEventRecord(upload_done, transfer_stream); + +// Stage 2: compute waits for upload +cudaStreamWaitEvent(compute_stream, upload_done); +process_kernel<<>>(d_input, d_output, n); +cudaEventRecord(compute_done, compute_stream); + +// Stage 3: download waits for compute +cudaStreamWaitEvent(transfer_stream, compute_done); +cudaMemcpyAsync(h_output, d_output, size, cudaMemcpyDeviceToHost, transfer_stream); + +// Sync before host reads the result +cudaStreamSynchronize(transfer_stream); +``` diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 505b0c0973..21285ab3ad 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -207,13 +207,13 @@ static torch::Tensor get_paged_mqa_logits_metadata(const torch::Tensor& context_ const auto arch_major = device_runtime->get_arch_major(); if (is_varlen) { const auto& indices_tensor = indices.value(); - DG_HOST_ASSERT(arch_major == 10 and next_n == 1 and (block_kv == 64 or block_kv == 32)); + DG_HOST_ASSERT((arch_major == 10 or arch_major == 12) and next_n == 1 and (block_kv == 64 or block_kv == 32)); DG_HOST_ASSERT(indices_tensor.dim() == 1 and indices_tensor.size(0) == batch_size); DG_HOST_ASSERT(indices_tensor.is_contiguous()); DG_HOST_ASSERT(indices_tensor.scalar_type() == torch::kInt); smxx_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, next_n, block_kv, num_sms, is_context_lens_2d, true, indices_tensor.data_ptr()); - } else if (arch_major == 9 or arch_major == 10) { - DG_HOST_ASSERT(block_kv == 64 or (arch_major == 10 and block_kv == 32)); + } else if (arch_major == 9 or arch_major == 10 or arch_major == 12) { + DG_HOST_ASSERT(block_kv == 64 or ((arch_major == 10 or arch_major == 12) and block_kv == 32)); smxx_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, next_n, block_kv, num_sms, is_context_lens_2d, false, nullptr); } else { DG_HOST_UNREACHABLE("Unsupported architecture"); @@ -335,7 +335,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tupleget_arch_major(); const auto indices_tensor = indices.value_or(torch::Tensor()); if (is_varlen) { - DG_HOST_ASSERT(arch_major == 10 and next_n == 1); + DG_HOST_ASSERT((arch_major == 10 or arch_major == 12) and next_n == 1); DG_HOST_ASSERT(indices_tensor.dim() == 1 and indices_tensor.size(0) == batch_size); DG_HOST_ASSERT(indices_tensor.is_contiguous()); DG_HOST_ASSERT(indices_tensor.scalar_type() == torch::kInt); @@ -368,6 +368,10 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple { +public: + struct Args { + int batch_size; + int next_n; + int num_heads; + int head_dim; + int block_kv; + bool is_context_lens_2d; + int block_table_stride; + int logits_stride; + int tokens_per_block; + + int64_t q_batch_stride; + int64_t q_next_stride; + int64_t q_head_stride; + int64_t kv_block_stride; + int64_t kv_token_stride; + int64_t kv_scale_block_stride; + int64_t weights_row_stride; + + void* q; + void* kv_cache; + float* kv_cache_scales; + float* weights; + int* context_lens; + void* logits; + int* block_table; + at::ScalarType logits_dtype; + + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm120_fp8_paged_mqa_logits_reference< + {}, {}, {}, + {}, {}, + {} + >); +}}; +)", args.num_heads, args.head_dim, args.block_kv, + args.is_context_lens_2d ? "true" : "false", args.tokens_per_block, + to_string(args.logits_dtype)); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.batch_size, + args.next_n, + args.logits_stride, + args.block_table_stride, + args.q_batch_stride, + args.q_next_stride, + args.q_head_stride, + args.kv_block_stride, + args.kv_token_stride, + args.kv_scale_block_stride, + args.weights_row_stride, + args.q, + args.kv_cache, + args.kv_cache_scales, + args.weights, + args.context_lens, + args.logits, + args.block_table + )); + } +}; + +static void sm120_fp8_paged_mqa_logits_reference(const torch::Tensor& q, + const torch::Tensor& kv_cache, + const torch::Tensor& kv_cache_scales, + const torch::Tensor& weights, + const torch::Tensor& context_lens, + const torch::Tensor& logits, + const torch::Tensor& block_table, + const at::ScalarType& logits_dtype, + const int& batch_size, const int& next_n, + const int& num_heads, const int& head_dim, + const int& block_kv, + const bool& is_context_lens_2d, + const int& logits_stride, + const int& block_table_stride) { + constexpr int tokens_per_block = 128; + DG_HOST_ASSERT(num_heads == 32 or num_heads == 64); + DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128); + DG_HOST_ASSERT(block_kv == 32 or block_kv == 64); + + const SM120FP8PagedMQALogitsReferenceRuntime::Args args = { + .batch_size = batch_size, + .next_n = next_n, + .num_heads = num_heads, + .head_dim = head_dim, + .block_kv = block_kv, + .is_context_lens_2d = is_context_lens_2d, + .block_table_stride = block_table_stride, + .logits_stride = logits_stride, + .tokens_per_block = tokens_per_block, + .q_batch_stride = q.stride(0), + .q_next_stride = q.stride(1), + .q_head_stride = q.stride(2), + .kv_block_stride = kv_cache.stride(0), + .kv_token_stride = kv_cache.stride(1), + .kv_scale_block_stride = kv_cache_scales.stride(0), + .weights_row_stride = weights.stride(0), + .q = q.data_ptr(), + .kv_cache = kv_cache.data_ptr(), + .kv_cache_scales = kv_cache_scales.data_ptr(), + .weights = weights.data_ptr(), + .context_lens = context_lens.data_ptr(), + .logits = logits.data_ptr(), + .block_table = block_table.data_ptr(), + .logits_dtype = logits_dtype, + .launch_args = LaunchArgs(std::make_pair(batch_size * next_n, ceil_div(logits_stride, tokens_per_block)), + tokens_per_block) + }; + const auto code = SM120FP8PagedMQALogitsReferenceRuntime::generate(args); + const auto runtime = compiler->build("sm120_fp8_paged_mqa_logits_reference", code); + SM120FP8PagedMQALogitsReferenceRuntime::launch(runtime, args); +} + static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q, const torch::Tensor& kv_cache, const torch::Tensor& kv_cache_scales, diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh new file mode 100644 index 0000000000..28adffa1a8 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh @@ -0,0 +1,75 @@ +#pragma once + +#include + +#include + +#include + +namespace deep_gemm { + +template +CUTLASS_GLOBAL __launch_bounds__(kTokensPerBlock, 1) +void sm120_fp8_paged_mqa_logits_reference( + const uint32_t batch_size, + const uint32_t next_n, + const uint32_t logits_stride, + const uint32_t block_table_stride, + const uint64_t q_batch_stride, + const uint64_t q_next_stride, + const uint64_t q_head_stride, + const uint64_t kv_block_stride, + const uint64_t kv_token_stride, + const uint64_t kv_scale_block_stride, + const uint64_t weights_row_stride, + const __nv_fp8_e4m3* q, + const __nv_fp8_e4m3* kv_cache, + const float* kv_cache_scales, + const float* weights, + const uint32_t* context_lens, + logits_dtype_t* logits, + const uint32_t* block_table) { + const auto row = blockIdx.x; + const auto token_idx = blockIdx.y * kTokensPerBlock + threadIdx.x; + if (row >= batch_size * next_n or token_idx >= logits_stride) + return; + + const auto batch_idx = row / next_n; + const auto next_idx = row - batch_idx * next_n; + const auto context_lens_idx = kIsContextLens2D ? row : batch_idx; + const auto context_len = context_lens[context_lens_idx]; + const auto output_offset = row * static_cast(logits_stride) + token_idx; + + if (token_idx >= context_len) { + logits[output_offset] = static_cast(-__uint_as_float(0x7f800000u)); + return; + } + + const auto logical_block_idx = token_idx / BLOCK_KV; + const auto token_in_block = token_idx - logical_block_idx * BLOCK_KV; + const auto kv_block_idx = block_table[batch_idx * static_cast(block_table_stride) + logical_block_idx]; + const auto q_base = q + batch_idx * q_batch_stride + next_idx * q_next_stride; + const auto kv_base = kv_cache + kv_block_idx * kv_block_stride + token_in_block * kv_token_stride; + const auto kv_scale = kv_cache_scales[kv_block_idx * kv_scale_block_stride + token_in_block]; + const auto weight_base = weights + row * weights_row_stride; + + float total = 0.0f; + #pragma unroll + for (uint32_t head_idx = 0; head_idx < kNumHeads; ++head_idx) { + const auto q_head = q_base + head_idx * q_head_stride; + float score = 0.0f; + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; ++dim_idx) { + const auto q_value = static_cast(q_head[dim_idx]); + const auto kv_value = static_cast(kv_base[dim_idx]) * kv_scale; + score = fmaf(q_value, kv_value, score); + } + total += fmaxf(score, 0.0f) * weight_base[head_idx]; + } + + logits[output_offset] = static_cast(total); +} + +} // namespace deep_gemm diff --git a/tests/test_sm120_reference.py b/tests/test_sm120_reference.py new file mode 100644 index 0000000000..a8b31cc80a --- /dev/null +++ b/tests/test_sm120_reference.py @@ -0,0 +1,126 @@ +import torch + +import deep_gemm +from deep_gemm.testing import calc_diff, get_arch_major, test_filter as _test_filter + + +def _cast_kv_cache_to_fp8(kv_cache: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + num_blocks, block_kv, num_heads, head_dim = kv_cache.shape + assert num_heads == 1 + + scale = kv_cache.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) / 448.0 + kv_fp8 = (kv_cache * scale.reciprocal()).to(torch.float8_e4m3fn) + kv_simulated = kv_fp8.float() * scale + + fused = torch.empty( + (num_blocks, block_kv * (head_dim + 4)), + device=kv_cache.device, + dtype=torch.uint8, + ) + fused[:, : block_kv * head_dim] = kv_fp8.view(num_blocks, block_kv * head_dim).view(torch.uint8) + fused[:, block_kv * head_dim :] = scale.view(num_blocks, block_kv).view(torch.uint8) + return fused.view(num_blocks, block_kv, num_heads, head_dim + 4), kv_simulated + + +def _paged_mqa_reference( + q: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + max_context_len: int, +) -> torch.Tensor: + batch_size, next_n, num_heads, head_dim = q.shape + _, block_kv, _, _ = kv_cache.shape + logits = torch.full( + (batch_size * next_n, max_context_len), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + for batch_idx in range(batch_size): + for next_idx in range(next_n): + row = batch_idx * next_n + next_idx + context_len = int(context_lens[batch_idx, next_idx].item()) + for token_idx in range(context_len): + block_idx = int(block_table[batch_idx, token_idx // block_kv].item()) + kv = kv_cache[block_idx, token_idx % block_kv, 0].float() + score = (q[batch_idx, next_idx].float() * kv).sum(dim=1) + logits[row, token_idx] = (score.relu() * weights[row]).sum() + + return logits + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_hc_prenorm_gemm_reference_path() -> None: + torch.manual_seed(0) + + m, n, k = 5, 8, 64 + a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + b = torch.randn((n, k), dtype=torch.float32, device="cuda") + for num_splits in (None, 3): + d = torch.empty((m, n), dtype=torch.float32, device="cuda") if num_splits is None else \ + torch.empty((num_splits, m, n), dtype=torch.float32, device="cuda") + sqr_sum = torch.empty((m,), dtype=torch.float32, device="cuda") if num_splits is None else \ + torch.empty((num_splits, m), dtype=torch.float32, device="cuda") + + deep_gemm.tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits=num_splits) + + final_d = d if num_splits is None else d.sum(dim=0) + final_sqr_sum = sqr_sum if num_splits is None else sqr_sum.sum(dim=0) + torch.testing.assert_close(final_d, a.float() @ b.T, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(final_sqr_sum, a.float().square().sum(dim=-1), rtol=1e-4, atol=1e-4) + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_fp8_paged_mqa_logits_reference_path() -> None: + torch.manual_seed(1) + + batch_size, next_n, num_heads, head_dim = 2, 2, 32, 32 + block_kv = 64 + max_context_len = 96 + num_blocks = 4 + + q = torch.randn((batch_size, next_n, num_heads, head_dim), device="cuda", dtype=torch.bfloat16) + q_fp8 = q.to(torch.float8_e4m3fn) + q_simulated = q_fp8.float() + + kv_cache = torch.randn((num_blocks, block_kv, 1, head_dim), device="cuda", dtype=torch.bfloat16) + fused_kv_cache, kv_simulated = _cast_kv_cache_to_fp8(kv_cache) + + weights = torch.randn((batch_size * next_n, num_heads), device="cuda", dtype=torch.float32) + context_lens = torch.tensor([[17, 23], [31, 47]], device="cuda", dtype=torch.int32) + block_table = torch.tensor([[0], [1]], device="cuda", dtype=torch.int32) + + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens, + block_kv, + deep_gemm.get_num_sms(), + ) + expected = _paged_mqa_reference( + q_simulated, + kv_simulated, + weights, + context_lens, + block_table, + max_context_len, + ) + valid_mask = torch.arange(max_context_len, device="cuda").unsqueeze(0) < context_lens.view(-1, 1) + for logits_dtype in (torch.float32, torch.bfloat16): + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + q=(q_fp8, None), + kv_cache=fused_kv_cache, + weights=weights, + context_lens=context_lens, + block_table=block_table, + schedule_meta=schedule_meta, + max_context_len=max_context_len, + clean_logits=False, + logits_dtype=logits_dtype, + ) + + assert logits.dtype == logits_dtype + diff = calc_diff(logits.float()[valid_mask], expected[valid_mask]) + threshold = 1e-6 if logits_dtype == torch.float32 else 2e-6 + assert diff < threshold, f"{logits_dtype=}, {diff=}" From 959f1df759ed591cadb463ab9af68222428de5df Mon Sep 17 00:00:00 2001 From: jasl Date: Sat, 25 Apr 2026 08:03:58 +0800 Subject: [PATCH 03/14] Add SM120 FP8 einsum reference fallback DeepSeek V4 on SM120 currently reaches an einsum path where the activation side is laid out as bhr and the output projection is logically hdr, but the weight and scale tensors can arrive from the vLLM Cutlass path as flattened 2D tensors reshaped into that logical form. Add an SM120-specific reference implementation for the bhr,hdr->bhd FP8 einsum recipe using FP32 block scales and scalar FP32 accumulation. This is intentionally a correctness-first CUDA fallback rather than an optimized Tensor Core kernel, and is only dispatched for arch major 12 with recipe (1, 128, 128). Verified on 10.0.0.110 with the vLLM editable build against /home/jasl/tmp/DeepGEMM. The generated JIT compile used --gpu-architecture=sm_120f, and direct DeepGEMM plus vLLM wrapper checks matched the PyTorch reference with max_abs 0.0. The full DeepSeek-V4-Flash serve path now gets past the previous DeepGEMM einsum failure and reaches the separate FlashMLA backend blocker. --- csrc/apis/einsum.hpp | 8 + csrc/jit_kernels/impls/sm120_fp8_einsum.hpp | 149 ++++++++++++++++++ .../deep_gemm/impls/sm120_fp8_einsum.cuh | 64 ++++++++ 3 files changed, 221 insertions(+) create mode 100644 csrc/jit_kernels/impls/sm120_fp8_einsum.hpp create mode 100644 deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index f82331cacc..4ea1b185d3 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -14,6 +14,7 @@ #include "../jit_kernels/impls/sm100_bmk_bnk_mn.hpp" #include "../jit_kernels/impls/sm90_bf16_gemm.hpp" #include "../jit_kernels/impls/sm100_bf16_gemm.hpp" +#include "../jit_kernels/impls/sm120_fp8_einsum.hpp" #include "../jit_kernels/impls/smxx_cublaslt.hpp" #endif @@ -186,6 +187,13 @@ static void fp8_einsum(const std::string& expr, // Some hardcoded Einstein sum kernels const auto arch_major = device_runtime->get_arch_major(); if (expr == "bhr,hdr->bhd") { + if (arch_major == 12) { + DG_HOST_ASSERT(not c.has_value()); + DG_HOST_ASSERT(recipe == std::make_tuple(1, 128, 128)); + sm120_fp8_bhr_hdr_bhd_reference(a.first, a.second, b.first, b.second, d); + return; + } + // Permute dims to satisfy the order of (batch_size, m, n, k) // (batch_size, m, n, k): (h, b, d, r) const auto perm_a = a.first.permute({1, 0, 2}); diff --git a/csrc/jit_kernels/impls/sm120_fp8_einsum.hpp b/csrc/jit_kernels/impls/sm120_fp8_einsum.hpp new file mode 100644 index 0000000000..94777410e2 --- /dev/null +++ b/csrc/jit_kernels/impls/sm120_fp8_einsum.hpp @@ -0,0 +1,149 @@ +#pragma once + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "runtime_utils.hpp" + +namespace deep_gemm { + +class SM120FP8BhrHdrBhdReferenceRuntime final: public LaunchRuntime { +public: + struct Args { + int shape_b; + int shape_h; + int shape_d; + int shape_r; + + int64_t a_stride_b; + int64_t a_stride_h; + int64_t a_stride_r; + int64_t sfa_stride_b; + int64_t sfa_stride_h; + int64_t sfa_stride_r; + int64_t b_stride_h; + int64_t b_stride_d; + int64_t b_stride_r; + int64_t sfb_stride_h; + int64_t sfb_stride_d; + int64_t sfb_stride_r; + int64_t d_stride_b; + int64_t d_stride_h; + int64_t d_stride_d; + + void* a; + float* sfa; + void* b; + float* sfb; + void* d; + at::ScalarType output_dtype; + + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + constexpr int output_tile_d = 128; + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm120_fp8_bhr_hdr_bhd_reference< + {}, {} + >); +}}; +)", to_string(args.output_dtype), output_tile_d); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.shape_b, + args.shape_h, + args.shape_d, + args.shape_r, + args.a_stride_b, + args.a_stride_h, + args.a_stride_r, + args.sfa_stride_b, + args.sfa_stride_h, + args.sfa_stride_r, + args.b_stride_h, + args.b_stride_d, + args.b_stride_r, + args.sfb_stride_h, + args.sfb_stride_d, + args.sfb_stride_r, + args.d_stride_b, + args.d_stride_h, + args.d_stride_d, + args.a, + args.sfa, + args.b, + args.sfb, + args.d + )); + } +}; + +static void sm120_fp8_bhr_hdr_bhd_reference(const torch::Tensor& a, + const torch::Tensor& sfa, + const torch::Tensor& b, + const torch::Tensor& sfb, + const torch::Tensor& d) { + constexpr int gran_mn = 128; + constexpr int gran_k = 128; + constexpr int output_tile_d = 128; + + const auto [shape_b, shape_h, shape_r] = get_shape<3>(a); + const auto [shape_h_, shape_d, shape_r_] = get_shape<3>(b); + const auto [shape_b_, shape_h__, shape_d_] = get_shape<3>(d); + + DG_HOST_ASSERT(shape_b == shape_b_ and shape_h == shape_h_ and shape_h == shape_h__); + DG_HOST_ASSERT(shape_r == shape_r_ and shape_d == shape_d_); + DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(sfa.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(sfb.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(a.stride(2) == 1 and b.stride(2) == 1 and d.stride(2) == 1); + + check_sf_layout(sfa, shape_h, shape_r, 1, gran_k, shape_b); + check_sf_layout(sfb, shape_d, shape_r, gran_mn, gran_k, shape_h); + + const int num_d_tiles = ceil_div(shape_d, output_tile_d); + const int grid_dim = shape_b * shape_h * num_d_tiles; + const SM120FP8BhrHdrBhdReferenceRuntime::Args args = { + .shape_b = shape_b, + .shape_h = shape_h, + .shape_d = shape_d, + .shape_r = shape_r, + .a_stride_b = a.stride(0), + .a_stride_h = a.stride(1), + .a_stride_r = a.stride(2), + .sfa_stride_b = sfa.stride(0), + .sfa_stride_h = sfa.stride(1), + .sfa_stride_r = sfa.stride(2), + .b_stride_h = b.stride(0), + .b_stride_d = b.stride(1), + .b_stride_r = b.stride(2), + .sfb_stride_h = sfb.stride(0), + .sfb_stride_d = sfb.stride(1), + .sfb_stride_r = sfb.stride(2), + .d_stride_b = d.stride(0), + .d_stride_h = d.stride(1), + .d_stride_d = d.stride(2), + .a = a.data_ptr(), + .sfa = sfa.data_ptr(), + .b = b.data_ptr(), + .sfb = sfb.data_ptr(), + .d = d.data_ptr(), + .output_dtype = d.scalar_type(), + .launch_args = LaunchArgs(grid_dim, output_tile_d) + }; + const auto code = SM120FP8BhrHdrBhdReferenceRuntime::generate(args); + const auto runtime = compiler->build("sm120_fp8_bhr_hdr_bhd_reference", code); + SM120FP8BhrHdrBhdReferenceRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh new file mode 100644 index 0000000000..601410bd72 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh @@ -0,0 +1,64 @@ +#pragma once + +#include + +#include + +#include + +namespace deep_gemm { + +template +CUTLASS_GLOBAL void sm120_fp8_bhr_hdr_bhd_reference( + const uint32_t shape_b, + const uint32_t shape_h, + const uint32_t shape_d, + const uint32_t shape_r, + const uint64_t a_stride_b, + const uint64_t a_stride_h, + const uint64_t a_stride_r, + const uint64_t sfa_stride_b, + const uint64_t sfa_stride_h, + const uint64_t sfa_stride_r, + const uint64_t b_stride_h, + const uint64_t b_stride_d, + const uint64_t b_stride_r, + const uint64_t sfb_stride_h, + const uint64_t sfb_stride_d, + const uint64_t sfb_stride_r, + const uint64_t d_stride_b, + const uint64_t d_stride_h, + const uint64_t d_stride_d, + const __nv_fp8_e4m3* a, + const float* sfa, + const __nv_fp8_e4m3* b, + const float* sfb, + output_t* d) { + const uint32_t num_d_tiles = (shape_d + kOutputTileD - 1) / kOutputTileD; + const uint32_t tile_idx = blockIdx.x; + const uint32_t d_tile_idx = tile_idx % num_d_tiles; + const uint32_t h_idx = (tile_idx / num_d_tiles) % shape_h; + const uint32_t b_idx = tile_idx / (num_d_tiles * shape_h); + const uint32_t d_idx = d_tile_idx * kOutputTileD + threadIdx.x; + + if (b_idx >= shape_b or d_idx >= shape_d) + return; + + float accum = 0.0f; + for (uint32_t r_idx = 0; r_idx < shape_r; ++r_idx) { + const uint32_t r_scale_idx = r_idx / 128; + const auto a_offset = b_idx * a_stride_b + h_idx * a_stride_h + r_idx * a_stride_r; + const auto b_offset = h_idx * b_stride_h + d_idx * b_stride_d + r_idx * b_stride_r; + const auto sfa_offset = b_idx * sfa_stride_b + h_idx * sfa_stride_h + r_scale_idx * sfa_stride_r; + const auto sfb_offset = h_idx * sfb_stride_h + (d_idx / 128) * sfb_stride_d + r_scale_idx * sfb_stride_r; + + const float a_value = static_cast(a[a_offset]) * sfa[sfa_offset]; + const float b_value = static_cast(b[b_offset]) * sfb[sfb_offset]; + accum = fmaf(a_value, b_value, accum); + } + + const auto d_offset = b_idx * d_stride_b + h_idx * d_stride_h + d_idx * d_stride_d; + d[d_offset] = static_cast(accum); +} + +} // namespace deep_gemm From 23ec7b5662101edbba60cf82a6661512a6759907 Mon Sep 17 00:00:00 2001 From: jasl Date: Sat, 25 Apr 2026 11:01:05 +0800 Subject: [PATCH 04/14] Exclude local Codex config --- .codex/skills/cpp/SKILL.md | 111 ---- .codex/skills/cpp/references/guidelines.md | 152 ------ .codex/skills/cpp/references/structure.md | 89 --- .codex/skills/readable-cpp/SKILL.md | 596 --------------------- .gitignore | 5 +- 5 files changed, 4 insertions(+), 949 deletions(-) delete mode 100644 .codex/skills/cpp/SKILL.md delete mode 100644 .codex/skills/cpp/references/guidelines.md delete mode 100644 .codex/skills/cpp/references/structure.md delete mode 100644 .codex/skills/readable-cpp/SKILL.md diff --git a/.codex/skills/cpp/SKILL.md b/.codex/skills/cpp/SKILL.md deleted file mode 100644 index 2c61ec1e86..0000000000 --- a/.codex/skills/cpp/SKILL.md +++ /dev/null @@ -1,111 +0,0 @@ ---- -name: cpp -description: Comprehensive C/C++ programming reference covering everything from C11-C23 and C++11-C++23, system programming, CUDA GPU computing, debugging tools, Rust interop, and advanced topics. Use for: C/C++ questions, C/C++ interview preparation, modern language features, RAII/memory management, templates/generics, CUDA programming, system programming, debugging/profiling, performance optimization, cross-platform development, build systems, assembly, shell scripting, and any C/C++/CUDA development tasks. ---- - -# C/C++ Comprehensive Cheat Sheets (/cpp) - -Complete C/C++ development reference combining local documentation with live examples from cppcheatsheet.com, covering everything from basic syntax to advanced GPU programming and system-level development. - -## How It Works - -Help users write functional, correct C/C++ code and answer C/C++ questions by fetching proven patterns and examples from cppcheatsheet.com. - -When a user asks a C/C++ question or wants to write C/C++ code: - -1. Look up the relevant topic(s) in [Structure](references/structure.md) to find the matching URL(s) -2. **Always fetch** the URL(s) using WebFetch to get real examples and patterns from the site -3. Use the fetched content to: - - **Write code**: Apply the patterns to produce functional, correct code that solves the user's task - - **Answer questions**: Provide thorough explanations backed by the examples and information from the site -4. Follow the [Guidelines](references/guidelines.md) for code quality - -## Key Principle - -**Functionality first, cleanliness second.** The code must work correctly and handle the task properly. Fetching from cppcheatsheet.com ensures solutions use battle-tested patterns rather than guessing. The site contains rich examples covering edge cases, common pitfalls, and practical usage that go beyond basic documentation. - -## Coverage Areas - -### C/C++ Interview Preparation -**Interview Cheatsheet:** Curated C and C++ interview questions grouped by topic (memory management, move semantics, templates, STL, concurrency, debugging), each deep-linked to the section of the cheat sheet that answers it. Use for interview prep or quick topic review. - -### Modern C Programming (C11-C23) -**Core Language:** Syntax, types, memory management, preprocessor macros -**GNU Extensions:** Compiler-specific features and optimizations -**Build Systems:** Makefiles, compilation, linking -**Assembly:** X86 assembly integration and inline assembly - -### Modern C++ Programming (C++11-C++23) -**Core Features:** RAII, templates, STL containers, iterators, algorithms -**Modern Standards:** Move semantics, constexpr, lambdas, concepts, coroutines, modules -**Memory Management:** Smart pointers, resource management, optimization (RVO) -**Build Systems:** CMake, package management, cross-platform builds - -### System Programming -**Process Management:** POSIX processes, signals, process communication -**File Systems:** File I/O, directory operations, filesystem monitoring -**Networking:** Sockets, protocols, network programming patterns -**Threading:** Multithreading, synchronization, parallel programming -**IPC:** Inter-process communication, shared memory, message queues - -### CUDA Programming -**GPU Computing:** CUDA kernels, memory hierarchy, performance optimization -**Advanced CUDA:** libcu++, Thrust library, cooperative groups -**Multi-GPU:** GPU-GPU communication, hardware topology, NVSHMEM -**Async Programming:** CUDA pipelines, memory visibility, asynchronous execution - -### Debugging & Profiling -**Debug Tools:** GDB debugging, Valgrind memory analysis, sanitizers, binary inspection (nm, readelf, objdump, otool) -**Performance:** Perf profiling, tracing, performance optimization -**GPU Debugging:** Nsight Systems, CUDA debugging and profiling - -### System Tools & Automation -**Shell Scripting:** Bash programming, system administration -**System Tools:** OS utilities, networking tools, hardware inspection -**Service Management:** Systemd, process management, system monitoring - -### Cross-Language Development -**Rust Interop:** Rust for C++ developers, FFI, memory safety comparison -**Language Bridging:** C/C++ integration, foreign function interfaces - -### Advanced Topics -**Blog Content:** Deep-dive articles, RDMA networking, GPU-initiated communication -**Low-Level Programming:** Hardware interfaces, performance tuning -**Architecture:** System design, scalable applications - -## References - -For detailed information, I can access: -- **[Structure](references/structure.md)** - Complete topic-to-URL reference map -- **[Guidelines](references/guidelines.md)** - Code quality and best practices - -## Examples - -### C/C++ Interview Prep -- "What should I review for a C++ interview?" → Fetch https://cppcheatsheet.com/notes/interview/index.html and walk the reader through the topic groups -- "Common C++ interview questions on smart pointers" → Fetch https://cppcheatsheet.com/notes/interview/index.html and then drill into https://cppcheatsheet.com/notes/cpp/cpp_smartpointers.html for detailed answers - -### C/C++ Core -- "How do smart pointers work?" → Fetch https://cppcheatsheet.com/notes/cpp/cpp_smartpointers.html and explain with the site's examples -- "How does RAII work in C++?" → Fetch https://cppcheatsheet.com/notes/cpp/cpp_raii.html and explain with practical examples -- "How to use STL containers?" → Fetch https://cppcheatsheet.com/notes/cpp/cpp_container.html and explain with practical examples -- "Template metaprogramming techniques" → Fetch https://cppcheatsheet.com/notes/cpp/cpp_template.html and explain with practical examples - -### System Programming -- "POSIX socket server in C" → Fetch https://cppcheatsheet.com/notes/os/os_socket.html, use the patterns to write a working server -- "Multithreading with std::thread" → Fetch https://cppcheatsheet.com/notes/os/os_thread.html and explain with practical examples -- "Signal handling and process management" → Fetch https://cppcheatsheet.com/notes/os/os_signal.html and explain with practical examples - -### CUDA & GPU Programming -- "Write a CUDA kernel" → Fetch https://cppcheatsheet.com/notes/cuda/cuda_basics.html, use the patterns to write working GPU code -- "CUDA memory hierarchy and optimization" → Fetch https://cppcheatsheet.com/notes/cuda/cuda_memory_visibility.html and explain with practical examples -- "Multi-GPU communication with NCCL" → Fetch https://cppcheatsheet.com/notes/cuda/cuda_nccl.html and explain with practical examples - -### Debugging & Tools -- "Debug a segfault with GDB" → Fetch https://cppcheatsheet.com/notes/debug/gdb.html and explain with practical examples -- "Valgrind memory leak detection" → Fetch https://cppcheatsheet.com/notes/debug/valgrind.html and explain with practical examples -- "Performance profiling with perf" → Fetch https://cppcheatsheet.com/notes/debug/perf.html and explain with practical examples - -### Build & Development -- "CMake cross-platform build systems" → Fetch https://cppcheatsheet.com/notes/cpp/cpp_cmake.html and explain with practical examples -- "Makefile patterns and best practices" → Fetch https://cppcheatsheet.com/notes/c/make.html and explain with practical examples \ No newline at end of file diff --git a/.codex/skills/cpp/references/guidelines.md b/.codex/skills/cpp/references/guidelines.md deleted file mode 100644 index 9af3ec1a39..0000000000 --- a/.codex/skills/cpp/references/guidelines.md +++ /dev/null @@ -1,152 +0,0 @@ -# C/C++ Comprehensive Development Guidelines - -Always fetch relevant examples from cppcheatsheet.com first to ensure correctness, then apply these guidelines when writing code. - -## Modern C Programming (C11-C23) - -### Code Quality & Standards -- Follow C11/C17/C23 standards for modern C development -- Use static analysis tools: clang-static-analyzer, cppcheck, scan-build -- Implement proper error handling with return codes and errno -- Use const correctness and restrict pointers for optimization -- Leverage compound literals and designated initializers (C99+) - -### Memory Management -- Always pair malloc/free, avoid memory leaks -- Use valgrind and AddressSanitizer for memory debugging -- Consider memory alignment for performance-critical code -- Use restrict keyword for pointer aliasing optimization - -### Build Systems & Portability -- Write portable Makefiles with proper dependency tracking -- Use feature test macros for cross-platform compatibility -- Leverage GNU extensions judiciously with fallbacks - -## Modern C++ Programming (C++11-C++23) - -### Resource Management (RAII) -- Always use RAII for resource management -- Prefer smart pointers: unique_ptr, shared_ptr, weak_ptr -- Use custom deleters for non-standard resources -- Implement move semantics for expensive-to-copy objects - -### Modern Language Features -- Use auto for type deduction, but maintain readability -- Leverage constexpr for compile-time computation -- Use structured bindings (C++17+) for multiple return values -- Implement concepts (C++20+) for better template constraints -- Use coroutines (C++20+) for asynchronous programming -- Adopt modules (C++20+) for better build times and encapsulation - -### Template Programming -- Use SFINAE and concepts for template constraints -- Prefer if constexpr over template specialization where applicable -- Use fold expressions (C++17+) for variadic templates -- Implement type traits for generic programming - -### Performance Optimization -- Profile before optimizing (use perf, vtune, or similar) -- Understand memory hierarchy and cache-friendly algorithms -- Use move semantics and perfect forwarding -- Leverage Return Value Optimization (RVO) and copy elision -- Consider noexcept specifications for optimization opportunities - -## System Programming - -### POSIX Best Practices -- Always check return values from system calls -- Handle EINTR properly in system call loops -- Use signal-safe functions in signal handlers -- Implement proper cleanup with RAII wrappers - -### Multithreading & Concurrency -- Use std::thread and C++11+ concurrency primitives -- Avoid data races with proper synchronization -- Prefer lock-free algorithms when appropriate -- Use thread-local storage for per-thread data -- Consider memory ordering for atomic operations - -### Network Programming -- Handle partial reads/writes in socket programming -- Use non-blocking I/O with proper error handling -- Implement timeout mechanisms for robust networking -- Consider endianness for network protocols - -## CUDA Programming - -### Performance Guidelines -- Optimize memory access patterns (coalesced access) -- Use shared memory effectively for data reuse -- Minimize divergence in warp execution -- Consider occupancy vs. resource usage tradeoffs - -### Memory Management -- Use CUDA unified memory judiciously -- Understand memory hierarchy: global, shared, constant, texture -- Implement proper error checking with cudaGetLastError() -- Use CUDA streams for overlapping computation and communication - -### Advanced CUDA -- Leverage cooperative groups for flexible thread cooperation -- Use libcu++ for standard library features on GPU -- Implement multi-GPU communication with NVSHMEM or NCCL -- Consider GPU-initiated communication for advanced patterns - -## Debugging & Profiling - -### Debug Tools Usage -- Use GDB with proper debug symbols (-g flag) -- Leverage Valgrind for memory error detection -- Use sanitizers: AddressSanitizer, ThreadSanitizer, UBSan -- Implement proper logging with different verbosity levels - -### Performance Analysis -- Use perf for CPU profiling and cache analysis -- Implement custom trace points for application-specific profiling -- Use Nsight Systems/Compute for CUDA application analysis -- Profile both CPU and GPU components in heterogeneous applications - -## Build Systems & DevOps - -### CMake Best Practices -- Use modern CMake (3.15+) with target-based design -- Properly handle dependencies with find_package or FetchContent -- Implement proper install rules and export configurations -- Use generator expressions for configuration-specific settings - -### Cross-Platform Development -- Write portable code with proper feature detection -- Use standard library features over platform-specific ones -- Handle different compiler behaviors (GCC, Clang, MSVC) -- Test on multiple platforms and architectures - -## Security & Safety - -### Memory Safety -- Use static analysis tools to catch buffer overflows -- Implement bounds checking in debug builds -- Consider using safer alternatives (std::array vs C arrays) -- Use string handling functions that prevent overflows - -### Code Review & Quality -- Implement automated testing with appropriate coverage -- Use continuous integration for multi-platform testing -- Enforce coding standards with clang-format and linters -- Document public APIs with clear contracts and examples - -## Cross-Language Integration - -### C/C++ Interoperability -- Use extern "C" linkage for C++ code called from C -- Handle exceptions at C++ boundaries when interfacing with C -- Use proper name mangling considerations - -### Rust Integration -- Understand ownership model differences between C++ and Rust -- Use proper FFI patterns for safe interoperation -- Handle error propagation between language boundaries -- Consider memory safety implications at interfaces - -## Related Documentation - -This skill is based on the comprehensive C/C++ reference available at https://cppcheatsheet.com/ which includes working code snippets, performance benchmarks, real-world patterns, and integration guides. The reference is continuously updated with the latest C/C++ features and best practices. \ No newline at end of file diff --git a/.codex/skills/cpp/references/structure.md b/.codex/skills/cpp/references/structure.md deleted file mode 100644 index be32cc2faa..0000000000 --- a/.codex/skills/cpp/references/structure.md +++ /dev/null @@ -1,89 +0,0 @@ -# C/C++ Topics Reference Map - -## C/C++ Interview Cheatsheet -- [C/C++ Interview Cheatsheet](https://cppcheatsheet.com/notes/interview/index.html) - Curated interview questions, each deep-linked to the section that answers it - -## Modern C Programming -- [C Basics](https://cppcheatsheet.com/notes/c/c_basic.html) - Core C language features from C11 to C23 -- [Memory](https://cppcheatsheet.com/notes/c/c_memory.html) - Memory management and allocation -- [Preprocessor & GNU Extensions](https://cppcheatsheet.com/notes/c/c_macro.html) - Macros and GNU extensions -- [Makefile](https://cppcheatsheet.com/notes/c/make.html) - Build system configuration -- [X86 Assembly](https://cppcheatsheet.com/notes/c/asm.html) - Assembly language integration - -## Modern C++ Programming -- [C++ Basics](https://cppcheatsheet.com/notes/cpp/cpp_basic.html) - Modern C++ features from C++11 to C++23 -- [Resource Management](https://cppcheatsheet.com/notes/cpp/cpp_raii.html) - RAII and resource management -- [String](https://cppcheatsheet.com/notes/cpp/cpp_string.html) - String handling and operations -- [Container](https://cppcheatsheet.com/notes/cpp/cpp_container.html) - STL containers and data structures -- [Iterator](https://cppcheatsheet.com/notes/cpp/cpp_iterator.html) - Iterator patterns and usage -- [Template](https://cppcheatsheet.com/notes/cpp/cpp_template.html) - Generic programming and metaprogramming -- [Casting](https://cppcheatsheet.com/notes/cpp/cpp_casting.html) - Type casting and conversions -- [Constexpr](https://cppcheatsheet.com/notes/cpp/cpp_constexpr.html) - Compile-time computation -- [Lambda](https://cppcheatsheet.com/notes/cpp/cpp_lambda.html) - Lambda expressions and functional programming -- [Concepts](https://cppcheatsheet.com/notes/cpp/cpp_concepts.html) - C++20 concepts and constraints -- [Requires](https://cppcheatsheet.com/notes/cpp/cpp_requires.html) - Requirements and constraints -- [Time](https://cppcheatsheet.com/notes/cpp/cpp_time.html) - Time and chrono utilities -- [Smart Pointers](https://cppcheatsheet.com/notes/cpp/cpp_smartpointers.html) - Modern memory management -- [Return Value Optimization](https://cppcheatsheet.com/notes/cpp/cpp_rvo.html) - RVO and move semantics -- [Algorithm](https://cppcheatsheet.com/notes/cpp/cpp_algorithm.html) - STL algorithms and functional patterns -- [Coroutine](https://cppcheatsheet.com/notes/cpp/cpp_coroutine.html) - C++20 coroutines -- [Modules](https://cppcheatsheet.com/notes/cpp/cpp_modules.html) - C++20 modules system -- [CMake](https://cppcheatsheet.com/notes/cpp/cpp_cmake.html) - Build system and project configuration - -## System Programming -- [Process](https://cppcheatsheet.com/notes/os/os_process.html) - Process management and control -- [File](https://cppcheatsheet.com/notes/os/os_file.html) - File I/O and filesystem operations -- [Signal](https://cppcheatsheet.com/notes/os/os_signal.html) - Signal handling and IPC -- [Socket](https://cppcheatsheet.com/notes/os/os_socket.html) - Network programming and sockets -- [Thread](https://cppcheatsheet.com/notes/os/os_thread.html) - Threading and synchronization -- [IPC](https://cppcheatsheet.com/notes/os/os_ipc.html) - Inter-process communication - -## CUDA Programming -- [CUDA Basics](https://cppcheatsheet.com/notes/cuda/cuda_basics.html) - GPU programming fundamentals -- [CUDA C++ (libcu++)](https://cppcheatsheet.com/notes/cuda/cuda_cpp.html) - Modern CUDA C++ features -- [Thrust](https://cppcheatsheet.com/notes/cuda/cuda_thrust.html) - High-level parallel algorithms -- [Cooperative Groups](https://cppcheatsheet.com/notes/cuda/cuda_coop_groups.html) - Thread cooperation patterns -- [Memory Visibility](https://cppcheatsheet.com/notes/cuda/cuda_memory_visibility.html) - Memory models and synchronization -- [CUDA Graph](https://cppcheatsheet.com/notes/cuda/cuda_graph.html) - Graph-based execution and optimization -- [CUDA Quantization](https://cppcheatsheet.com/notes/cuda/cuda_quantization.html) - Quantization techniques on GPU -- [Pipelines](https://cppcheatsheet.com/notes/cuda/cuda_pipelines.html) - Asynchronous execution pipelines -- [NCCL](https://cppcheatsheet.com/notes/cuda/cuda_nccl.html) - Multi-GPU collective communication -- [GPU-GPU Communication](https://cppcheatsheet.com/notes/cuda/cuda_ipc.html) - Multi-GPU programming -- [Hardware Topology](https://cppcheatsheet.com/notes/cuda/cuda_hwloc.html) - Hardware locality and topology - -## Blog -- [Building NVSHMEM from Scratch](https://cppcheatsheet.com/notes/blog/nvshmem.html) - GPU-initiated networking deep dive - -## Bash & System Tools -- [Bash](https://cppcheatsheet.com/notes/tools/bash.html) - Shell scripting and command-line tools -- [Operating System](https://cppcheatsheet.com/notes/tools/os.html) - System administration tools -- [Network](https://cppcheatsheet.com/notes/tools/net.html) - Network configuration and debugging -- [Hardware](https://cppcheatsheet.com/notes/tools/hardware.html) - Hardware inspection and monitoring -- [GPU](https://cppcheatsheet.com/notes/tools/gpu.html) - GPU tools and utilities -- [Systemd](https://cppcheatsheet.com/notes/tools/systemd.html) - Service management and system control - -## Debugging & Profiling -- [GDB](https://cppcheatsheet.com/notes/debug/gdb.html) - GNU debugger and debugging techniques -- [Valgrind](https://cppcheatsheet.com/notes/debug/valgrind.html) - Memory analysis and leak detection -- [Sanitizers](https://cppcheatsheet.com/notes/debug/sanitizers.html) - Address, thread, and undefined behavior sanitizers -- [Tracing](https://cppcheatsheet.com/notes/debug/tracing.html) - System and application tracing -- [Perf](https://cppcheatsheet.com/notes/debug/perf.html) - Performance analysis and profiling -- [Nsight Systems](https://cppcheatsheet.com/notes/debug/nsight.html) - NVIDIA profiling tools -- [Binary Tools](https://cppcheatsheet.com/notes/debug/binary_tools.html) - Binary inspection with nm, readelf, objdump, otool, ldd, strings, size - -## Rust for C++ Developers -- [Basics](https://cppcheatsheet.com/notes/rust/rust_basic.html) - Rust fundamentals for C++ developers -- [Ownership & Borrowing](https://cppcheatsheet.com/notes/rust/rust_ownership.html) - Memory safety without GC -- [RAII & Drop](https://cppcheatsheet.com/notes/rust/rust_raii.html) - Resource management patterns -- [Strings](https://cppcheatsheet.com/notes/rust/rust_string.html) - String handling and text processing -- [Collections](https://cppcheatsheet.com/notes/rust/rust_container.html) - Data structures and containers -- [Iterators](https://cppcheatsheet.com/notes/rust/rust_iterator.html) - Functional iteration patterns -- [Traits & Generics](https://cppcheatsheet.com/notes/rust/rust_traits.html) - Generic programming and interfaces -- [Casting](https://cppcheatsheet.com/notes/rust/rust_casting.html) - Type conversions and casting -- [Const Functions](https://cppcheatsheet.com/notes/rust/rust_constfn.html) - Compile-time computation -- [Closures](https://cppcheatsheet.com/notes/rust/rust_closure.html) - Anonymous functions and captures -- [Smart Pointers](https://cppcheatsheet.com/notes/rust/rust_smartptr.html) - Reference counting and ownership -- [Error Handling](https://cppcheatsheet.com/notes/rust/rust_error.html) - Result types and error propagation -- [Threads](https://cppcheatsheet.com/notes/rust/rust_thread.html) - Safe concurrency and parallelism -- [Modules](https://cppcheatsheet.com/notes/rust/rust_modules.html) - Code organization and visibility -- [FFI & C++ Bindings](https://cppcheatsheet.com/notes/rust/rust_ffi.html) - Interoperability with C/C++ diff --git a/.codex/skills/readable-cpp/SKILL.md b/.codex/skills/readable-cpp/SKILL.md deleted file mode 100644 index 1f56f88374..0000000000 --- a/.codex/skills/readable-cpp/SKILL.md +++ /dev/null @@ -1,596 +0,0 @@ ---- -name: readable-cpp -description: Readable C/C++/Rust/CUDA code rules inspired by The Art of Readable Code. Use when writing, reviewing, or refactoring C, C++, Rust, or CUDA code. Enforces short functions, flat control flow, clear naming, readable structure, and idiomatic patterns. ---- - -# Readable C/C++/Rust/CUDA Rules (/readable-cpp) - -Apply these rules when writing, reviewing, or refactoring C, C++, Rust, or CUDA code. Inspired by *The Art of Readable Code* by Dustin Boswell and Trevor Foucher. - -**Core principle: Code should be easy to understand.** The time it takes someone else (or future you) to understand the code is the ultimate metric. - -## 1. Keep Functions Short and Focused - -- A function should do **one thing**. If you can describe what it does with "and", split it. -- Aim for functions that fit on one screen (~15-25 lines). If it's longer, extract sub-tasks. -- Each function should operate at a **single level of abstraction** — don't mix high-level logic with low-level details in the same function. - -## 2. Flatten Control Flow — No Deep Nesting - -- **Never nest more than 2 levels deep.** If you have a loop inside a loop, or an `if` inside a loop inside an `if`, extract the inner block into a helper function with a descriptive name. -- Use **early returns / guard clauses** to handle edge cases at the top, keeping the main logic flat. -- Prefer `continue` or `break` to skip iterations rather than wrapping the body in a conditional. -- Replace complex conditionals with well-named helper functions or variables that explain the intent. - -```cpp -// Bad: nested and hard to follow -for (auto& user : users) { - if (user.is_active()) { - for (auto& order : user.orders()) { - if (order.is_pending()) { - process(order); - } - } - } -} - -// Good: flat, each function name explains what it does -auto active_users = get_active_users(users); -for (auto& user : active_users) { - process_pending_orders(user.orders()); -} -``` - -## 3. Name Things Clearly - -- **Pack information into names.** Use specific, concrete words — `fetch_page` not `get`, `num_retries` not `n`. -- **Avoid generic names** like `tmp`, `data`, `result`, `val`, `info`, `handle` — unless the scope is tiny (2-3 lines). -- **Use names that can't be misconstrued.** If a range is inclusive, say `max_items` not `limit`. If a boolean, use `is_`, `has_`, `should_`, `can_` prefixes. -- **Match the name length to the scope.** Short names for small scopes, descriptive names for wide scopes. -- **Don't use abbreviations** unless they're universally understood (`num`, `max`, `min`, `err` are fine; `svc_mgr_cfg` is not). - -## 4. Make Control Flow Easy to Follow - -- Put the **changing/interesting value on the left** side of comparisons: `if (length > 10)` not `if (10 < length)`. -- Order `if/else` blocks: **positive case first**, simpler case first, or the more interesting case first. -- Minimize the number of variables the reader has to track. Reduce the **mental footprint** of each block. -- Avoid deeply nested ternary operators — if it's not immediately obvious, use an `if/else`. - -## 5. Break Down Giant Expressions - -- Use **explaining variables** to break complex expressions into named pieces. -- Use **summary variables** to capture a long expression that's used more than once. -- Apply **De Morgan's laws** to simplify negated boolean expressions. - -```cpp -// Bad -if (!(age >= 18 && has_id && !is_banned)) { - deny(); -} - -// Good -bool is_eligible = age >= 18 && has_id && !is_banned; -if (!is_eligible) { - deny(); -} -``` - -## 6. Extract Unrelated Subproblems - -- If a block of code is solving a **subproblem unrelated to the main goal** of the function, extract it. -- The helper function should be **pure and self-contained** — it shouldn't need to know about the calling context. -- This is the single most effective way to improve readability: separate *what* you're doing from *how*. - -## 7. One Task at a Time - -- Each section of code should do **one task**. If a function is doing parsing AND validation AND transformation, split them into separate steps. -- List the tasks a function does. If there's more than one, reorganize so each task is in its own block or function. - -## 8. Reduce Variable Scope - -- **Declare variables close to where they're used.** Don't declare at the top of a function if it's only used 30 lines later. -- **Minimize the "live time" of a variable** — the fewer lines between its assignment and last use, the easier it is to follow. -- **Prefer write-once variables.** Variables that are assigned once and never modified are easier to reason about. -- **Eliminate unnecessary variables.** If a variable is used only once and doesn't clarify anything, inline it. - -## 9. No Magic Numbers or Strings - -- Replace **magic numbers and strings** with named constants: `if (retries > MAX_RETRIES)` not `if (retries > 3)`. -- If a value has meaning, give it a name. The name documents the intent. -- Group related constants together. - -## 10. Fewer Function Arguments - -- Aim for **3 or fewer arguments** per function. More than that is a smell. -- Group related arguments into a **struct, class, or tuple**. -- If a function needs many config-like options, pass a single config/options object. -- Boolean flag arguments are a sign the function does two things — split it instead. - -## 11. Consistency - -- If the codebase does something one way, **do it the same way**. Don't mix styles. -- Consistent naming patterns, consistent structure, consistent error handling. -- When joining an existing codebase, **match the existing conventions** even if you'd prefer a different style. -- Surprise is the enemy of readability — predictable code is readable code. - -## 12. Write Less Code - -- The best code is **no code at all**. Question whether a feature is truly needed before implementing. -- **Don't over-engineer.** Solve the problem at hand, not hypothetical future problems. -- Remove dead code. Commented-out code is dead code. -- Use standard libraries before writing custom solutions. - -## 13. Comments: Explain Why, Not What - -- Don't comment **what** the code does — the code already says that. Comment **why** it does it. -- Comment **flaws and workarounds**: `// TODO:`, `// HACK:`, `// XXX:` with explanation. -- Comment **surprising behavior** or non-obvious decisions — things where a reader would ask "why?". -- **Don't comment bad code — rewrite it.** If you need a comment to explain what a block does, extract it into a well-named function instead. - -## 14. Design Code to Survive Auto-Formatting - -- Write code that looks good **after** the auto-formatter runs. If a chained expression or repeated pattern would be broken across 4+ lines by the formatter, extract a helper function instead. -- **Prefer one-line helper calls** over long inline chains that the formatter will expand vertically. -- The formatter is your reader's first impression. Run it *before* committing — if the result looks ugly, that's a signal to refactor, not to disable the formatter. - -```rust -// Bad: rustfmt expands this to 4 lines per field — noisy and repetitive -fn from_dict(cfg: &Bound<'_, PyDict>) -> PyResult { - Ok(Self { - rom: cfg.get_item("rom")?.ok_or_else(|| missing("rom"))?.extract()?, - // ... each field becomes 4 lines after rustfmt - }) -} - -// Good: extract a helper so each field stays one clean line -fn get_required(cfg: &Bound<'_, PyDict>, key: &str) -> PyResult { - cfg.get_item(key)? - .ok_or_else(|| PyKeyError::new_err(key.to_string()))? - .extract() -} - -fn from_dict(cfg: &Bound<'_, PyDict>) -> PyResult { - Ok(Self { - rom: get_required(cfg, "rom")?, - actions: get_required(cfg, "actions")?, - }) -} -``` - -```cpp -// Bad: clang-format wraps this into a hard-to-scan block -auto result = container.find(key)->second.get_value().transform(func).value_or(default_val); - -// Good: name the intermediate step -auto& entry = container.find(key)->second; -auto result = entry.get_value().transform(func).value_or(default_val); -``` - ---- - -# C-Specific Rules - -## 15. RAII-Like Patterns with goto Cleanup - -- In C, use the **goto cleanup pattern** for resource management — allocate at the top, clean up at a single labeled block at the bottom. -- Never scatter `free()` calls across multiple return paths. A single cleanup section is easier to audit. -- Use `__attribute__((cleanup))` (GCC/Clang) when available for automatic cleanup. - -```c -// Good: single cleanup path -int process_file(const char *path) { - int ret = -1; - FILE *fp = fopen(path, "r"); - if (!fp) return -1; - - char *buf = malloc(BUF_SIZE); - if (!buf) goto cleanup_file; - - // ... do work ... - ret = 0; - -cleanup_buf: - free(buf); -cleanup_file: - fclose(fp); - return ret; -} -``` - -## 16. Use `const` Liberally - -- Mark pointers `const` when the function doesn't modify the pointed-to data: `const char *msg`. -- Mark local variables `const` when they don't change after initialization. -- This documents intent and helps the compiler catch mistakes. - -## 17. Prefer Sized Types for Data Structures - -- Use `` types (`uint32_t`, `int64_t`) for data that crosses boundaries (files, network, hardware). -- Use `size_t` for sizes and counts, `ptrdiff_t` for pointer differences. -- Use `int` and `unsigned` for simple loop counters and local arithmetic. - -## 18. Defensive Macro Hygiene - -- Wrap macro bodies in `do { ... } while(0)` for statement-like macros. -- Parenthesize all macro parameters: `#define SQUARE(x) ((x) * (x))`. -- Prefer `static inline` functions over macros when possible (type safety, debuggability). -- Use `_Generic` (C11) for type-safe "overloading" instead of macro tricks. - ---- - -# C++-Specific Rules - -## 19. Use RAII for All Resources - -- Every resource (memory, file handles, locks, sockets) should be owned by an RAII object. -- Use `std::unique_ptr` for exclusive ownership, `std::shared_ptr` only when shared ownership is genuinely needed. -- Write custom RAII wrappers for non-standard resources (e.g., C library handles). -- Never use raw `new`/`delete` in application code — let smart pointers and containers handle it. - -## 20. Prefer Value Semantics and Move - -- Pass small objects by value, large objects by `const&`. -- Return objects by value — rely on RVO/NRVO and move semantics. -- Implement move constructors/assignment for types that own resources. -- Use `std::move` only when you truly want to transfer ownership — don't `std::move` from things you'll use again. - -## 21. Use Modern C++ Over C Idioms - -- Use `std::array` over C arrays, `std::string` over `char*`, `std::vector` over `malloc`/`realloc`. -- Use `std::optional` over sentinel values, `std::variant` over type-unsafe unions. -- Use range-based `for` loops: `for (const auto& item : container)`. -- Use structured bindings (C++17): `auto [key, value] = *map.begin();`. -- Use `std::format` (C++20) or `fmt::format` over `sprintf` / string concatenation. - -## 22. Templates: Keep It Simple - -- Use concepts (C++20) to constrain templates — errors become readable. -- Prefer `if constexpr` over SFINAE when possible. -- Don't write template metaprogramming unless the benefit is clear and the team can maintain it. -- A non-template solution that's slightly less generic is often better than a template solution nobody understands. - -## 23. Use `constexpr` and `const` Aggressively - -- Mark functions `constexpr` when they can be evaluated at compile time. -- Use `constexpr` variables instead of `#define` for constants. -- Use `const` on member functions that don't modify state. -- `consteval` (C++20) for functions that *must* be compile-time evaluated. - -## 24. Error Handling: Pick One Pattern - -- Use exceptions for truly exceptional conditions, `std::expected` (C++23) or `std::optional` for expected failures. -- Don't mix error codes and exceptions in the same layer. -- If using exceptions, make them specific — derive from `std::runtime_error`, not `std::exception`. -- Use `noexcept` on functions that cannot throw (destructors, move operations). - ---- - -# Rust-Specific Rules - -## 25. Embrace the Ownership Model - -- Don't fight the borrow checker — redesign your data flow instead. -- Prefer passing references (`&T`, `&mut T`) over cloning. Clone only when ownership transfer is genuinely needed. -- Use lifetimes explicitly only when the compiler can't infer them — don't annotate unnecessarily. -- Prefer `&str` over `String` in function parameters when you don't need ownership. - -## 26. Use Iterators and Combinators - -- Prefer iterator chains (`.iter().filter().map().collect()`) over manual loops with indices. -- Use `for item in &collection` instead of `for i in 0..collection.len()`. -- Use `enumerate()`, `zip()`, `chain()`, `chunks()` — the iterator API is rich. -- Avoid `.unwrap()` in production code — use `?`, `unwrap_or`, `unwrap_or_else`, or pattern matching. - -```rust -// Bad: manual indexing -let mut names = Vec::new(); -for i in 0..users.len() { - if users[i].is_active { - names.push(users[i].name.clone()); - } -} - -// Good: idiomatic iterator chain -let names: Vec<_> = users.iter() - .filter(|u| u.is_active) - .map(|u| u.name.clone()) - .collect(); -``` - -## 27. Use Enums and Pattern Matching - -- Use `enum` with data variants instead of class hierarchies or tagged unions. -- Use `match` exhaustively — the compiler ensures you handle all cases. -- Use `if let` / `while let` for single-variant matching instead of full `match`. -- Prefer `Result` over panicking — make errors part of the type signature. - -## 28. Leverage the Type System - -- Use **newtype wrappers** (`struct UserId(u64)`) to prevent mixing up same-typed values. -- Use `Option` instead of sentinel values or null pointers. -- Use `#[must_use]` on functions whose return values shouldn't be ignored. -- Prefer `From`/`Into` traits for type conversions over manual conversion functions. - -## 29. Module Organization - -- Keep `pub` surfaces small — expose only what's needed. -- Use `pub(crate)` for crate-internal visibility instead of full `pub`. -- Group related types and functions in modules — one concept per module. -- Re-export key types at the crate root for ergonomic imports. - ---- - -# CUDA-Specific Rules - -## 30. Name Kernels and Device Functions Clearly - -- Kernel names should describe **what** they compute, not that they're kernels: `reduce_sum` not `kernel1` or `myKernel`. -- Use a consistent naming convention to distinguish execution spaces: e.g., `reduce_sum_kernel` for `__global__`, `warp_reduce` for `__device__` helpers. -- Name grid/block dimension variables descriptively: `threads_per_block`, `num_blocks` not `tpb`, `nb`, or bare `256`. - -```cuda -// Bad: opaque names, magic numbers -__global__ void k1(float *a, float *b, int n) { - int i = blockIdx.x * 256 + threadIdx.x; - if (i < n) b[i] = a[i] * 2.0f; -} -k1<<<(n+255)/256, 256>>>(d_in, d_out, n); - -// Good: clear intent, named constants -constexpr int THREADS_PER_BLOCK = 256; - -__global__ void scale_kernel(const float *input, float *output, - float scale_factor, int num_elements) { - const int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < num_elements) { - output[idx] = input[idx] * scale_factor; - } -} - -const int num_blocks = (num_elements + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; -scale_kernel<<>>(d_input, d_output, 2.0f, num_elements); -``` - -## 31. Separate Host Logic from Device Logic - -- Keep **host orchestration** (memory allocation, transfers, kernel launches, synchronization) in separate functions from **device computation** (kernels and device helpers). -- Don't mix `cudaMalloc`/`cudaMemcpy` with application logic — wrap them in RAII classes or helper functions. -- Use a clear file organization: consider separating `.cu` kernel files from `.cpp` host logic files, or at minimum group host and device code into clearly labeled sections. - -```cpp -// Good: RAII wrapper hides allocation/deallocation -template -class DeviceBuffer { - T *ptr_ = nullptr; - size_t size_ = 0; -public: - explicit DeviceBuffer(size_t count) : size_(count) { - check_cuda(cudaMalloc(&ptr_, count * sizeof(T))); - } - ~DeviceBuffer() { cudaFree(ptr_); } - - DeviceBuffer(const DeviceBuffer&) = delete; - DeviceBuffer& operator=(const DeviceBuffer&) = delete; - DeviceBuffer(DeviceBuffer&& o) noexcept : ptr_(o.ptr_), size_(o.size_) { o.ptr_ = nullptr; } - - T *get() { return ptr_; } - const T *get() const { return ptr_; } - size_t size() const { return size_; } - - void copy_from_host(const T *host_data) { - check_cuda(cudaMemcpy(ptr_, host_data, size_ * sizeof(T), cudaMemcpyHostToDevice)); - } - void copy_to_host(T *host_data) const { - check_cuda(cudaMemcpy(host_data, ptr_, size_ * sizeof(T), cudaMemcpyDeviceToHost)); - } -}; -``` - -## 32. Always Check CUDA Errors - -- **Check every CUDA API call.** Silent failures are the #1 source of hard-to-debug CUDA issues. -- Use a `check_cuda` macro or inline function — not raw `if` blocks after every call. -- Check errors after kernel launches with `cudaGetLastError()` + `cudaDeviceSynchronize()` during development. -- In release builds, at minimum check allocations and memcpy — these are the most likely to fail at runtime. - -```cuda -// Good: concise, catches file/line info -inline void check_cuda(cudaError_t err, const char *file, int line) { - if (err != cudaSuccess) { - fprintf(stderr, "CUDA error at %s:%d — %s\n", - file, line, cudaGetErrorString(err)); - exit(EXIT_FAILURE); - } -} -#define check_cuda(err) check_cuda((err), __FILE__, __LINE__) - -// Usage -check_cuda(cudaMalloc(&d_ptr, size)); -my_kernel<<>>(d_ptr, n); -check_cuda(cudaGetLastError()); -check_cuda(cudaDeviceSynchronize()); -``` - -## 33. Make Thread Indexing Obvious - -- Compute the global thread index **once** at the top of the kernel and store it in a clearly named variable. -- Use **early return** for out-of-bounds threads — don't wrap the entire kernel body in an `if`. -- For 2D/3D grids, name dimensions explicitly: `row`, `col`, `depth` — not `x`, `y`, `z`. - -```cuda -// Bad: index computed inline, entire body wrapped -__global__ void process(float *data, int width, int height) { - if (blockIdx.x * blockDim.x + threadIdx.x < width && - blockIdx.y * blockDim.y + threadIdx.y < height) { - int idx = (blockIdx.y * blockDim.y + threadIdx.y) * width + - (blockIdx.x * blockDim.x + threadIdx.x); - data[idx] = data[idx] * 2.0f; - } -} - -// Good: named indices, early return -__global__ void process(float *data, int width, int height) { - const int col = blockIdx.x * blockDim.x + threadIdx.x; - const int row = blockIdx.y * blockDim.y + threadIdx.y; - if (col >= width || row >= height) return; - - const int idx = row * width + col; - data[idx] = data[idx] * 2.0f; -} -``` - -## 34. Document Shared Memory Usage - -- Declare shared memory with a **descriptive name** that indicates what it holds: `shared_tile` not `smem` or `s`. -- Add a brief comment explaining the **size** and **purpose** of shared memory when it's dynamically allocated (`extern __shared__`). -- Keep the shared memory lifecycle short — load, sync, compute, sync — and make each phase visually distinct. - -```cuda -// Good: clear phases, descriptive names -__global__ void tiled_matmul_kernel(const float *A, const float *B, - float *C, int N) { - __shared__ float tile_A[TILE_SIZE][TILE_SIZE]; - __shared__ float tile_B[TILE_SIZE][TILE_SIZE]; - - const int row = blockIdx.y * TILE_SIZE + threadIdx.y; - const int col = blockIdx.x * TILE_SIZE + threadIdx.x; - float accumulator = 0.0f; - - for (int tile_idx = 0; tile_idx < N / TILE_SIZE; ++tile_idx) { - // Phase 1: Load tiles from global memory - tile_A[threadIdx.y][threadIdx.x] = A[row * N + tile_idx * TILE_SIZE + threadIdx.x]; - tile_B[threadIdx.y][threadIdx.x] = B[(tile_idx * TILE_SIZE + threadIdx.y) * N + col]; - __syncthreads(); - - // Phase 2: Compute partial dot product from tiles - for (int k = 0; k < TILE_SIZE; ++k) { - accumulator += tile_A[threadIdx.y][k] * tile_B[k][threadIdx.x]; - } - __syncthreads(); - } - - C[row * N + col] = accumulator; -} -``` - -## 35. Keep Kernels Short — Extract Device Helpers - -- Apply the same "one function, one task" rule to kernels. If a kernel does loading, computing, and reducing, extract `__device__` helper functions. -- Use `__forceinline__ __device__` for small helpers that you want inlined without relying on compiler heuristics. -- This makes kernels easier to read, test (via unit-testing device functions), and reuse. - -```cuda -// Good: kernel reads like pseudocode, details in helpers -__forceinline__ __device__ -float warp_reduce_sum(float val) { - for (int offset = warpSize / 2; offset > 0; offset /= 2) { - val += __shfl_down_sync(0xffffffff, val, offset); - } - return val; -} - -__forceinline__ __device__ -float block_reduce_sum(float val) { - __shared__ float warp_sums[32]; - const int lane = threadIdx.x % warpSize; - const int warp_id = threadIdx.x / warpSize; - - val = warp_reduce_sum(val); - if (lane == 0) warp_sums[warp_id] = val; - __syncthreads(); - - val = (threadIdx.x < blockDim.x / warpSize) ? warp_sums[lane] : 0.0f; - if (warp_id == 0) val = warp_reduce_sum(val); - return val; -} - -__global__ void reduce_sum_kernel(const float *input, float *output, int n) { - const int idx = blockIdx.x * blockDim.x + threadIdx.x; - const float val = (idx < n) ? input[idx] : 0.0f; - - const float block_sum = block_reduce_sum(val); - if (threadIdx.x == 0) atomicAdd(output, block_sum); -} -``` - -## 36. Be Explicit About Memory Spaces - -- Use `const` on kernel parameters for read-only device pointers — documents intent and enables compiler optimizations. -- Use `__restrict__` when pointers don't alias — but add a comment explaining the non-aliasing guarantee. -- When using unified memory (`cudaMallocManaged`), comment the expected access pattern (host-only init, device-only compute, etc.) — the implicit page migration behavior is not obvious. -- Prefer explicit memory copies over unified memory in performance-critical paths — be explicit about data movement. - -```cuda -// Good: const + restrict with clear intent -__global__ void vector_add_kernel( - const float *__restrict__ a, // read-only, no alias with output - const float *__restrict__ b, // read-only, no alias with output - float *__restrict__ output, // write-only - int num_elements) -{ - const int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= num_elements) return; - output[idx] = a[idx] + b[idx]; -} -``` - -## 37. Synchronization: Make It Visible and Minimal - -- Place `__syncthreads()` on its own line, never buried inside a conditional branch that not all threads take — this is **undefined behavior** and hard to spot. -- Add a brief comment before each `__syncthreads()` stating what invariant it establishes: "all threads have loaded their tile", "partial sums are written to shared memory". -- Minimize synchronization points — restructure algorithms to reduce the number of barriers. -- For warp-level operations, prefer warp intrinsics (`__shfl_sync`, `__ballot_sync`) with explicit masks over `__syncthreads()`. - -## 38. Launch Configuration: Make It Readable - -- Wrap kernel launches in a **host function** that computes and names the launch parameters. -- Never hardcode grid/block dimensions at the call site — compute them from the problem size. -- For complex launch configurations, use a struct or helper to make the 2D/3D grid/block shape clear. -- Query device properties at startup rather than assuming specific hardware limits. - -```cuda -// Bad: magic numbers, unclear intent -foo<<<(n+127)/128, 128, 0, stream>>>(d_ptr, n); - -// Good: named, computed, self-documenting -void launch_scale_kernel(float *d_data, float factor, int n, cudaStream_t stream) { - constexpr int BLOCK_SIZE = 256; - const int grid_size = (n + BLOCK_SIZE - 1) / BLOCK_SIZE; - scale_kernel<<>>(d_data, factor, n); - check_cuda(cudaGetLastError()); -} -``` - -## 39. Streams and Async: Comment the Dependency Graph - -- When using multiple CUDA streams, add a comment block showing the **dependency graph** — which operations must complete before others begin. -- Name streams after their purpose: `compute_stream`, `transfer_stream` — not `s1`, `s2`. -- Group related async operations visually and separate independent pipelines with blank lines. -- Always synchronize before reading results on the host — make the sync point explicit and commented. - -```cuda -// Good: dependency graph documented, streams named by purpose -// Dependency graph: -// upload (transfer_stream) --> compute (compute_stream) --> download (transfer_stream) -// Event 'upload_done' gates compute start. -// Event 'compute_done' gates download start. - -cudaStream_t transfer_stream, compute_stream; -cudaEvent_t upload_done, compute_done; - -// Stage 1: async upload -cudaMemcpyAsync(d_input, h_input, size, cudaMemcpyHostToDevice, transfer_stream); -cudaEventRecord(upload_done, transfer_stream); - -// Stage 2: compute waits for upload -cudaStreamWaitEvent(compute_stream, upload_done); -process_kernel<<>>(d_input, d_output, n); -cudaEventRecord(compute_done, compute_stream); - -// Stage 3: download waits for compute -cudaStreamWaitEvent(transfer_stream, compute_done); -cudaMemcpyAsync(h_output, d_output, size, cudaMemcpyDeviceToHost, transfer_stream); - -// Sync before host reads the result -cudaStreamSynchronize(transfer_stream); -``` diff --git a/.gitignore b/.gitignore index d0cdf6ca42..dd11d39075 100644 --- a/.gitignore +++ b/.gitignore @@ -17,8 +17,11 @@ deep_gemm/include/cutlass /.clang* /.cache +# Local Codex skills and agent configuration +/.codex/ + # Generated stub files stubs/ # Symlinks to compiled extensions -deep_gemm/*.so \ No newline at end of file +deep_gemm/*.so From 64978f1ca9a51b3072eaf05a3b689813b657ab7f Mon Sep 17 00:00:00 2001 From: ergodic-flow <277504589+ergodic-flow@users.noreply.github.com> Date: Fri, 24 Apr 2026 14:31:25 -0400 Subject: [PATCH 05/14] Implement SM120 kernel for tf32_hc_prenorm_gemm This change introduces a new kernel for SM120 architecture (RTX 5090 and 6000 Pro) specifically for tf32_hc_prenorm_gemm . --- csrc/apis/hyperconnection.hpp | 28 +- csrc/indexing/main.cu | 1 + csrc/jit_kernels/heuristics/sm120.hpp | 243 +++++++++++++++ .../impls/sm120_tf32_hc_prenorm_gemm.hpp | 153 ++++++++++ .../impls/sm120_tf32_hc_prenorm_gemm.cuh | 288 ++++++++++++++++++ deep_gemm/include/deep_gemm/mma/sm120.cuh | 62 ++++ 6 files changed, 749 insertions(+), 26 deletions(-) create mode 100644 csrc/jit_kernels/heuristics/sm120.hpp create mode 100644 csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp create mode 100644 deep_gemm/include/deep_gemm/impls/sm120_tf32_hc_prenorm_gemm.cuh create mode 100644 deep_gemm/include/deep_gemm/mma/sm120.cuh diff --git a/csrc/apis/hyperconnection.hpp b/csrc/apis/hyperconnection.hpp index 7ee74b6cab..97c4b36fd2 100644 --- a/csrc/apis/hyperconnection.hpp +++ b/csrc/apis/hyperconnection.hpp @@ -5,34 +5,12 @@ #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp" #include "../jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp" +#include "../jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp" #endif namespace deep_gemm::hyperconnection { #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE -static void sm12x_tf32_hc_prenorm_gemm_reference(const torch::Tensor& a, - const torch::Tensor& b, - const torch::Tensor& d, - const torch::Tensor& sqr_sum, - const int& num_splits, - const bool& has_split_outputs) { - auto a_float = a.to(torch::kFloat); - auto d_value = at::matmul(a_float, b.transpose(0, 1)); - auto sqr_sum_value = (a_float * a_float).sum(-1); - - if (has_split_outputs) { - d.zero_(); - sqr_sum.zero_(); - d.select(0, 0).copy_(d_value); - sqr_sum.select(0, 0).copy_(sqr_sum_value); - return; - } - - DG_HOST_ASSERT(num_splits == 1); - d.copy_(d_value); - sqr_sum.copy_(sqr_sum_value); -} - static void tf32_hc_prenorm_gemm(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, @@ -76,9 +54,7 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a, } else if (arch_major == 10) { sm100_tf32_hc_prenorm_gemm(a, b, d, sqr_sum, m, n, k, num_splits.has_value() ? num_splits.value() : 1); } else if (arch_major == 12) { - sm12x_tf32_hc_prenorm_gemm_reference(a, b, d, sqr_sum, - num_splits.has_value() ? num_splits.value() : 1, - num_splits.has_value()); + sm120_tf32_hc_prenorm_gemm(a, b, d, sqr_sum, m, n, k, num_splits.has_value() ? num_splits.value() : 1); } else { DG_HOST_UNREACHABLE("Unsupported architecture"); } diff --git a/csrc/indexing/main.cu b/csrc/indexing/main.cu index a42b66f9e6..24c188f674 100644 --- a/csrc/indexing/main.cu +++ b/csrc/indexing/main.cu @@ -20,6 +20,7 @@ // Hyperconnection kernels #include #include +#include // Layout kernels #include diff --git a/csrc/jit_kernels/heuristics/sm120.hpp b/csrc/jit_kernels/heuristics/sm120.hpp new file mode 100644 index 0000000000..a204b61b65 --- /dev/null +++ b/csrc/jit_kernels/heuristics/sm120.hpp @@ -0,0 +1,243 @@ +#pragma once + +#include +#include + +#include "common.hpp" +#include "utils.hpp" +#include "../../utils/exception.hpp" + +namespace deep_gemm { + +struct SM120ArchSpec { + + static constexpr int smem_capacity = 232448; + + static std::vector get_layout_candidates(const GemmDesc& desc) { + // Block M candidates + std::vector block_m_candidates; + if (desc.gemm_type == GemmType::Normal or + desc.gemm_type == GemmType::Batched or + desc.gemm_type == GemmType::KGroupedContiguous) { + block_m_candidates = {64, 128}; + if (desc.m <= 16) block_m_candidates.push_back(16); + if (desc.m <= 32) block_m_candidates.push_back(32); + + // BF16 output GEMM supports 256 + if (desc.cd_dtype != torch::kFloat) + block_m_candidates.push_back(256); + } else if (desc.gemm_type == GemmType::MGroupedContiguous or + desc.gemm_type == GemmType::MGroupedContiguousWithPsumLayout) { + block_m_candidates = std::vector{heuristics_runtime->get_mk_alignment_for_contiguous_layout()}; + } else if (desc.gemm_type == GemmType::MGroupedMasked) { + block_m_candidates = {64, 128}; + } + + // Block N candidates + std::vector block_n_candidates; + int step = std::lcm(16, heuristics_runtime->get_block_n_multiple_of()); + int start = step; + // Avoid bank conflicts for 1D1D kernel FP32 output + if (desc.kernel_type == KernelType::Kernel1D1D and desc.cd_dtype == torch::kFloat) { + DG_HOST_ASSERT(desc.major_a == cute::UMMA::Major::K); + DG_HOST_ASSERT(desc.major_b == cute::UMMA::Major::K); + start = 24; + block_n_candidates.push_back(16); + } + // Register spills + int end = 256; + if (desc.kernel_type == KernelType::Kernel1D2D) + end = 192; + if (desc.kernel_type == KernelType::Kernel1D1D) + end = 160; + // Enumerate + for (int i = start; i <= end; i += step) + block_n_candidates.push_back(i); + + // Block K is always in a fixed manner + const int block_k = 128 / get_element_size(desc.get_mma_kind()); + + // Disable multicast for performance + const bool disable_multicast = + // The number of k-groups is large (a heuristic) + (desc.gemm_type == GemmType::KGroupedContiguous and desc.num_groups > 4) or + // Not supported + (desc.gemm_type == GemmType::Batched); + + // Enumerate all candidates + std::vector candidates; + for (int cluster_m = 1; cluster_m <= (disable_multicast ? 1 : 2); ++ cluster_m) { + for (int cluster_n = 1; cluster_n <= (disable_multicast ? 1 : 2); ++ cluster_n) { + // We only support cluster 2 + if (cluster_m * cluster_n > 2) + continue; + + // SM count must be divisible + if (desc.num_sms % (cluster_m * cluster_n) != 0) + continue; + + for (int block_m: block_m_candidates) { + for (int block_n: block_n_candidates) { + // 1D2D kernel unroll requirement + if (desc.kernel_type == KernelType::Kernel1D2D and block_n > block_k and (block_n % (block_n - block_k) != 0 and block_k % (block_n - block_k) != 0)) + continue; + + // Multicast legality for masked layout + // TODO: add some comments about it + if ((desc.gemm_type == GemmType::MGroupedMasked or desc.gemm_type == GemmType::MGroupedContiguousWithPsumLayout) and + ceil_div(desc.n, block_n) % (cluster_m * cluster_n) != 0) + continue; + + // The block sizes cannot be too large (for enough registers), so at least one dim less than 128 + if (block_m > 128 and block_n > 128) + continue; + + // Calculate swizzling + const auto layout = Layout{0, block_m, block_n, block_k, cluster_m, cluster_n}; + const auto storage_config = get_storage_config(desc, layout); + + // Make sure swizzling is large enough (32B's performance is low) + if (storage_config.swizzle_a_mode % 64 != 0 or storage_config.swizzle_b_mode % 64 != 0) + continue; + + // To hide TMA latency, the stage count should be at least 3; for small matrices, at least 4 + int num_stages = get_pipeline_config(desc, layout, storage_config).num_stages; + if (num_stages < 3 or (block_m * block_n < 128 * 192 and num_stages < 4)) + continue; + + candidates.push_back(layout); + } + } + } + } + + DG_HOST_ASSERT(not candidates.empty()); + return candidates; + } + + static StorageConfig get_storage_config(const GemmDesc& desc, const Layout& layout) { + constexpr int mma_m = 64; + + // Load/store block sizes (w/o consideration of swizzling atoms, w/ consideration of loop atoms) + // TODO: support swap AB + DG_HOST_ASSERT(layout.swap_ab == 0); + const auto load_block_m = layout.block_m; + const auto load_block_n = layout.block_n; + // 1D1D kernel will do single warp-group stores + const auto store_block_m = desc.kernel_type == KernelType::Kernel1D1D ? mma_m : layout.block_m; + const auto store_block_n = layout.block_n; + + // Decide swizzling by the inner dim + const auto swizzle_mode_a = get_swizzle_mode( + desc.major_a == cute::UMMA::Major::K ? layout.block_k : load_block_m, c10::elementSize(desc.a_dtype)); + const auto swizzle_mode_b = get_swizzle_mode( + desc.major_b == cute::UMMA::Major::K ? layout.block_k : load_block_n, c10::elementSize(desc.b_dtype)); + // We only enable swizzling for non-FP32 outputs + const auto swizzle_mode_cd = desc.cd_dtype != torch::kFloat ? + get_swizzle_mode(store_block_n, c10::elementSize(desc.cd_dtype)) : 0; + + return { + load_block_m, load_block_n, + store_block_m, store_block_n, + swizzle_mode_a, swizzle_mode_b, swizzle_mode_cd + }; + } + + static PipelineConfig get_pipeline_config(const GemmDesc& desc, const Layout& layout, const StorageConfig& storage_config) { + constexpr int kNumMaxStages = 16; + + // TODO: consider swap AB + // C/D for TMA stores + // NOTES: 1024 is for TMA swizzling alignment requirement + const int smem_cd = + align(layout.block_m * layout.block_n * static_cast(c10::elementSize(desc.cd_dtype)), 1024); + const int smem_barriers = kNumMaxStages * 8 * 2; + + // Calculate A/B per stages + const int smem_a_per_stage = storage_config.load_block_m * layout.block_k * c10::elementSize(desc.a_dtype); + const int smem_b_per_stage = storage_config.load_block_n * layout.block_k * c10::elementSize(desc.b_dtype); + + // Calculate SF A/B per stages + const int smem_sfa_per_stage = desc.kernel_type == KernelType::KernelNoSF ? + 0 : align(layout.block_m * static_cast(sizeof(float)), 128); + const int smem_sfb_per_stage = desc.kernel_type != KernelType::Kernel1D1D ? + 0 : align(layout.block_n * static_cast(sizeof(float)), 128); + + // Extra SFB sizes for 1D2D kernels + const int use_uniform_sfb = layout.block_k % layout.block_n == 0 ? 1 : 2; + const int smem_extra_sfb = desc.kernel_type != KernelType::Kernel1D2D ? + 0 : align(ceil_div(desc.k, layout.block_k) * static_cast(sizeof(float)) * use_uniform_sfb, 8); + + // Extra tensormap for 1D1D kernels + const int smem_tensormap = + desc.gemm_type == GemmType::KGroupedContiguous ? 4 * static_cast(sizeof(CUtensorMap)) : 0; + + // Calculate stages + const int smem_extra = smem_cd + smem_barriers + smem_extra_sfb + smem_tensormap; + const int smem_per_stage = smem_a_per_stage + smem_b_per_stage + smem_sfa_per_stage + smem_sfb_per_stage; + const int num_stages = std::min( + (smem_capacity - smem_extra) / smem_per_stage, + kNumMaxStages); + return { + smem_extra + num_stages * smem_per_stage, + num_stages + }; + } + + static LaunchConfig get_launch_config(const GemmDesc& desc, const Layout& layout) { + const int num_tma_threads = 128; + const int num_math_threads = layout.block_m <= 64 ? 128 : 256; + return { + desc.num_sms, + layout.get_cluster_size(), + num_tma_threads + num_math_threads, + num_tma_threads, num_math_threads, + 0, 0 // Meaningless for SM120 + }; + } + + static LayoutInfo get_layout_info(const GemmDesc& desc, const Layout& layout) { + const auto num_blocks = + ceil_div(desc.get_expected_m(), layout.block_m) * + ceil_div(desc.get_expected_n(), layout.block_n) * + desc.get_expected_num_groups(); + const auto num_waves = ceil_div(num_blocks, desc.num_sms); + const auto num_last_blocks = num_blocks % desc.num_sms; + const auto last_wave_util = num_last_blocks == 0 ? desc.num_sms : num_last_blocks; + + // Utils + const int l2_bandwidth_per_cycle = std::min(64. * desc.num_sms, 8e6 / (1.3e3)); // B/cycle + const int l1_bandwidth_per_cycle = 128 * desc.num_sms; // B/cycle + const int mma_m = 64; + const int elem_size_ab = c10::elementSize(desc.a_dtype); + const int elem_size_cd = c10::elementSize(desc.cd_dtype); + DG_HOST_ASSERT(desc.a_dtype == desc.b_dtype); + + // Data movement per block + int64_t expected_k = desc.get_expected_k(); + int64_t num_bytes_l2_ab = expected_k * (layout.block_m / layout.cluster_n + layout.block_n / layout.cluster_m) * elem_size_ab; + int64_t num_bytes_l1_ab = expected_k * (layout.block_m + layout.block_n) * elem_size_ab; + int64_t num_bytes_l1_tc = expected_k * (std::max(mma_m, layout.block_m) + layout.block_n) * elem_size_ab + + layout.block_m * layout.block_n * elem_size_cd; + int64_t num_bytes_l1_l2_cd = layout.block_m * layout.block_n * elem_size_cd * (desc.with_accumulation ? 2 : 1); + + // HBM bandwidth and total compute (Tensor/CUDA cores) are constant across configs + // We only model L1/L2 cycles as they are the primary variables between configs + int64_t num_l2_cycles = (num_bytes_l2_ab + num_bytes_l1_l2_cd) * num_blocks / l2_bandwidth_per_cycle; + int64_t num_l1_cycles = (num_bytes_l1_ab + num_bytes_l1_tc + num_bytes_l1_l2_cd) * num_blocks / l1_bandwidth_per_cycle; + float wave_efficiency = static_cast(num_blocks) / (num_waves * desc.num_sms); + int64_t num_cycles = std::max(num_l1_cycles, num_l2_cycles) / wave_efficiency; + + // Disable multicasting if only one wave exists + if (layout.cluster_n * layout.cluster_m > 1 and num_waves <= 1) + num_cycles = std::numeric_limits::max(); + + return {num_waves, last_wave_util, num_cycles, layout}; + } + + static bool compare(const LayoutInfo& a, const LayoutInfo& b) { + return a.num_cycles < b.num_cycles; + } +}; + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp new file mode 100644 index 0000000000..c37e1ab253 --- /dev/null +++ b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp @@ -0,0 +1,153 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "../../utils/math.hpp" +#include "../heuristics/sm120.hpp" +#include "runtime_utils.hpp" + +namespace deep_gemm { + +class SM120BF16HCPrenormGemmRuntime final: public LaunchRuntime { +public: + struct Args { + int m, n, k; + int block_m, block_n, block_k; + int num_splits; + int swizzle_cd_mode; + int num_stages; + int num_math_threads, num_tma_threads; + + LaunchArgs launch_args; + + CUtensorMap tensor_map_a; + CUtensorMap tensor_map_b; + CUtensorMap tensor_map_d; + float* sqr_sum; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm120_tf32_hc_prenorm_gemm_impl< + {}, {}, + {}, {}, {}, + {}, + {}, + {}, + {}, {} + >); +}}; +)", + args.n, args.k, + args.block_m, args.block_n, args.block_k, + args.num_splits, + args.swizzle_cd_mode, + args.num_stages, + args.num_math_threads, args.num_tma_threads); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.m, args.tensor_map_a, args.tensor_map_b, args.tensor_map_d, args.sqr_sum)); + } +}; + +static void sm120_tf32_hc_prenorm_gemm(const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& d, + const torch::Tensor& sqr_sum, + const int& m, const int& n, const int& k, + const int& num_splits) { + + // Note: may need to change these later + constexpr int block_m = 64; + constexpr int block_k = 64; + constexpr int num_math_threads = 128; + constexpr int num_tma_threads = 128; + constexpr int num_threads = num_math_threads + num_tma_threads; + + const int block_n = align(n, 16); + DG_HOST_ASSERT(n <= block_n); + // Only support small N for now + DG_HOST_ASSERT(n <= 32 and n % 8 == 0); + DG_HOST_ASSERT(k % block_k == 0); + + const auto swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float)); + const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k, + block_m, block_k, + static_cast(a.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1, + get_swizzle_mode(block_k, a.element_size()), 0, + true); + const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k, + block_n, block_k, + static_cast(b.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1, + get_swizzle_mode(block_k, b.element_size()), 0, + true); + const auto tensor_map_d = num_splits == 1 ? make_tma_cd_desc(d, m, n, + block_m, block_n, + static_cast(d.stride(-2)), 1, + swizzle_cd_mode) + : make_tma_3d_desc(d, n, m, num_splits, + block_n, block_m, 1, + static_cast(d.stride(-2)), + static_cast(d.stride(-3)), + swizzle_cd_mode); + + // Calculate stages + int num_stages = 12, smem_size = 0; + while (num_stages > 0) { + const int smem_a_per_stage = block_m * block_k * static_cast(sizeof(nv_bfloat16)); + const int smem_b_per_stage = block_n * block_k * static_cast(sizeof(float)); + const int smem_cd = block_m * swizzle_cd_mode; + const int smem_barriers = num_stages * 2 * 8; + smem_size = (smem_a_per_stage + smem_b_per_stage) * num_stages + + smem_cd + smem_barriers; + + const int smem_capacity = device_runtime->get_prop()->sharedMemPerBlockOptin; + if (smem_size <= smem_capacity) + break; + -- num_stages; + } + DG_HOST_ASSERT(num_stages > 0); + + if (get_env("DG_JIT_DEBUG", 0)) { + printf("M: %d, N: %d, K: %d -> " + "block M: %d, block N: %d, block K: %d, split K: %d" + "stages: %d, shared memory: %d, swizzle CD: %d\n", + m, n, k, block_m, block_n, block_k, num_splits, + num_stages, smem_size, swizzle_cd_mode); + } + + smem_size = device_runtime->get_prop()->sharedMemPerBlockOptin; + + const SM120BF16HCPrenormGemmRuntime::Args& args = { + .m = m, .n = n, .k = k, + .block_m = block_m, .block_n = block_n, .block_k = block_k, + .num_splits = num_splits, + .swizzle_cd_mode = swizzle_cd_mode, + .num_stages = num_stages, + .num_math_threads = num_math_threads, + .num_tma_threads = num_tma_threads, + .launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_threads, smem_size), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_d = tensor_map_d, + .sqr_sum = sqr_sum.data_ptr() + }; + + const auto code = SM120BF16HCPrenormGemmRuntime::generate(args); + const auto runtime = compiler->build("sm120_tf32_hc_prenorm_gemm", code); + SM120BF16HCPrenormGemmRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/impls/sm120_tf32_hc_prenorm_gemm.cuh b/deep_gemm/include/deep_gemm/impls/sm120_tf32_hc_prenorm_gemm.cuh new file mode 100644 index 0000000000..d7d659daf1 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm120_tf32_hc_prenorm_gemm.cuh @@ -0,0 +1,288 @@ +#pragma once +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include + +#include +#include +#include +#include + +#include +#include +#include + +namespace deep_gemm { + +template +CUTLASS_DEVICE +uint32_t get_sm120_swizzled_bank_group_idx(const uint32_t& offset, const uint32_t& lane_idx) { + constexpr uint32_t kGroupsInSwizzleRange = kSwizzleMode / kSwizzleBase; + + const auto bank_group_idx = offset + lane_idx * kGroupsInSwizzleRange; + + constexpr uint32_t kNumBankGroups = 128 / kSwizzleBase; + constexpr bool kHasShortcut = kGroupsInSwizzleRange == kNumBankGroups; + auto row = kHasShortcut ? (offset / kNumBankGroups + lane_idx) : (bank_group_idx / kNumBankGroups); + auto col = kHasShortcut ? (offset) : (bank_group_idx % kNumBankGroups); + col ^= row % kGroupsInSwizzleRange; + + return (row * kNumBankGroups + col) % kGroupsInSwizzleRange; +} + +template +CUTLASS_GLOBAL void __launch_bounds__(kNumMathThreads + kNumTMAThreads, 1) +sm120_tf32_hc_prenorm_gemm_impl(const uint32_t shape_m, + const __grid_constant__ cute::TmaDescriptor tensor_map_a, + const __grid_constant__ cute::TmaDescriptor tensor_map_b, + const __grid_constant__ cute::TmaDescriptor tensor_map_d, + float* sqr_sum) { + +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1200)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + constexpr uint32_t kSwizzleAMode = cute::min(BLOCK_K * sizeof(nv_bfloat16), 128); + constexpr uint32_t kSwizzleBMode = cute::min(BLOCK_K * sizeof(float), 128); + DG_STATIC_ASSERT(BLOCK_K == 64, "Invalid block K"); + DG_STATIC_ASSERT(kSwizzleAMode == 128, "Invalid swizzle A mode"); + DG_STATIC_ASSERT(kSwizzleBMode == 128, "Invalid swizzle B mode"); + + DG_STATIC_ASSERT(kSwizzleCDMode / sizeof(float) == BLOCK_N, "Invalid block N"); + DG_STATIC_ASSERT(kNumMathThreads == 128, "Invalid MMA threads"); + + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto lane_idx = ptx::get_lane_idx(); + + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + + constexpr uint32_t SMEM_CD_SIZE = BLOCK_M * kSwizzleCDMode; + constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(float); + DG_STATIC_ASSERT(SMEM_CD_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes"); + + if (warp_idx == 0 and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_a); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_d); + } + + auto smem_cd = reinterpret_cast(smem_buffer); + auto smem_a = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE)); + }); + auto smem_b = utils::PatternVisitor([&](const uint32_t& i) { + return reinterpret_cast(smem_buffer + (SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE)); + }); + + auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE)); + auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (i); }); + auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages + i); }); + + if (warp_idx == 1 and cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + full_barriers[i]->init(1); + empty_barriers[i]->init(128); + } + cutlass::arch::fence_barrier_init(); + } + __syncthreads(); + + constexpr uint32_t kNumKBlocks = math::constexpr_ceil_div(SHAPE_K, BLOCK_K); + constexpr uint32_t kNumKBlocksPerSplit = kNumKBlocks / kNumSplits; + constexpr uint32_t kRemainKBlocks = kNumKBlocks % kNumSplits; + const uint32_t block_idx = __shfl_sync(0xffffffff, blockIdx.x, 0); + const uint32_t m_block_idx = block_idx / kNumSplits; + const uint32_t k_split_idx = block_idx % kNumSplits; + const uint32_t k_offset = (k_split_idx * kNumKBlocksPerSplit + cute::min(k_split_idx, kRemainKBlocks)) * BLOCK_K; + const uint32_t m_offset = shape_m * k_split_idx; + const uint32_t num_total_stages = kNumKBlocksPerSplit + (k_split_idx < kRemainKBlocks); + + constexpr uint32_t kNumTMARegisters = 40; + + cudaGridDependencySynchronize(); + + if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) { + cutlass::arch::warpgroup_reg_dealloc(); + for (uint32_t s = 0; s < num_total_stages; ++ s) { + const auto stage_idx = s % kNumStages; + empty_barriers[stage_idx]->wait(((s / kNumStages) & 1) ^ 1); + + uint32_t m_idx = m_block_idx * BLOCK_M; + uint32_t k_idx = k_offset + s * BLOCK_K; + + tma::copy(&tensor_map_a, full_barriers[stage_idx], smem_a[stage_idx], k_idx, m_idx); + tma::copy(&tensor_map_b, full_barriers[stage_idx], smem_b[stage_idx], k_idx, 0); + + constexpr uint32_t kNumArrivalBytes = SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE; + full_barriers[stage_idx]->arrive_and_expect_tx(kNumArrivalBytes); + } + + for (uint32_t s = num_total_stages; s < num_total_stages + kNumStages; ++ s) { + const auto stage_idx = s % kNumStages; + empty_barriers[stage_idx]->wait(((s / kNumStages) & 1) ^ 1); + } + } else if (warp_idx < kNumMathThreads / 32) { + + DG_STATIC_ASSERT(BLOCK_M == 64, "Invalid block M"); + DG_STATIC_ASSERT(BLOCK_K * sizeof(nv_bfloat16) == kSwizzleAMode, "Invalid block K"); + constexpr uint32_t BLOCK_M_PER_WARP = BLOCK_M / 4; + constexpr uint32_t WGMMA_N = BLOCK_N; + + using MMASelector = mma::sm120::TF32MMASelector; + float accum[MMASelector::N_atoms * MMASelector::kNumAccumPerAtom] = {0}; + + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(nv_bfloat16); + constexpr uint32_t kNumLoads = BLOCK_K / kNumElemsPerBankGroup; + float sqr_sum_acc_0 = 0; + float sqr_sum_acc_1 = 0; + + #pragma unroll kNumStages < 8 ? kNumStages : kNumStages / 2 + for (uint32_t s = 0; s < num_total_stages; ++ s) { + const auto& stage_idx = s % kNumStages; + full_barriers[stage_idx]->wait((s / kNumStages) & 1); + + // Using MMA atom dimensions mapped to registers + constexpr uint32_t kNumRegPerMMA = MMASelector::M * MMASelector::K / 32; // 4 floats + constexpr uint32_t kNumMMAPerBlockK = BLOCK_K / MMASelector::K; + + float a[kNumRegPerMMA * kNumMMAPerBlockK]; + DG_STATIC_ASSERT(kSwizzleAMode == 128, "Invalid swizzle A mode"); + + uint32_t row = warp_idx * 16 + lane_idx / 4; + + #pragma unroll + for (uint32_t i = 0; i < kNumLoads; ++ i) { + uint32_t bank_group_idx = (row ^ i) % 8; + nv_bfloat16* a_bf16_smem_ptr_upper = smem_a[stage_idx] + row * BLOCK_K + bank_group_idx * kNumElemsPerBankGroup; + nv_bfloat16* a_bf16_smem_ptr_lower = smem_a[stage_idx] + (row + 8) * BLOCK_K + bank_group_idx * kNumElemsPerBankGroup; + + uint32_t elem_offset = lane_idx % 4; + + nv_bfloat16 a_bf16[kNumRegPerMMA]; + a_bf16[0] = a_bf16_smem_ptr_upper[elem_offset]; + a_bf16[1] = a_bf16_smem_ptr_lower[elem_offset]; + a_bf16[2] = a_bf16_smem_ptr_upper[elem_offset + 4]; + a_bf16[3] = a_bf16_smem_ptr_lower[elem_offset + 4]; + + auto a_bf16x2_ptr = reinterpret_cast(a_bf16); + auto a_float2_ptr = reinterpret_cast(a); + + float2 a_float2_0 = __bfloat1622float2(a_bf16x2_ptr[0]); + float2 a_float2_1 = __bfloat1622float2(a_bf16x2_ptr[1]); + + a_float2_ptr[i * 2 + 0] = a_float2_0; + a_float2_ptr[i * 2 + 1] = a_float2_1; + + sqr_sum_acc_0 += a_float2_0.x * a_float2_0.x + a_float2_1.x * a_float2_1.x; + sqr_sum_acc_1 += a_float2_0.y * a_float2_0.y + a_float2_1.y * a_float2_1.y; + } + + __syncwarp(); + if (s > 0) + empty_barriers[(s - 1) % kNumStages]->arrive(); + + constexpr int kNumElemsInSwizzleRange = 128 / sizeof(float); + constexpr uint32_t kNumAtomsInSwizzleRange = kNumElemsInSwizzleRange / MMASelector::K; + DG_STATIC_ASSERT(BLOCK_K % kNumElemsInSwizzleRange == 0, "Invalid block K"); + + #pragma unroll + for (int i = 0; i < BLOCK_K / kNumElemsInSwizzleRange; i++) { + #pragma unroll + for (int k = 0; k < kNumAtomsInSwizzleRange; k++) { + + float* a_step = a + (i * kNumAtomsInSwizzleRange + k) * kNumRegPerMMA; + + #pragma unroll + for (int n = 0; n < MMASelector::N_atoms; n++) { + + uint32_t atom_n = lane_idx / 4; + uint32_t atom_k = lane_idx % 4; + + uint32_t global_n = n * 8 + atom_n; + uint32_t global_k = (i * kNumAtomsInSwizzleRange + k) * 8 + atom_k; + + uint32_t b_atom_idx = global_k / kNumElemsInSwizzleRange; + uint32_t b_atom_k = global_k % kNumElemsInSwizzleRange; + uint32_t atom_linear_idx_0 = global_n * kNumElemsInSwizzleRange + b_atom_k; + uint32_t atom_linear_idx_1 = atom_linear_idx_0 + 4; + + uint32_t swizzle_xor = ((atom_linear_idx_0 >> 5) & 7) << 2; + uint32_t atom_swizzled_idx_0 = atom_linear_idx_0 ^ swizzle_xor; + uint32_t atom_swizzled_idx_1 = atom_linear_idx_1 ^ swizzle_xor; + uint32_t atom_base_idx = b_atom_idx * BLOCK_N * kNumElemsInSwizzleRange; + + float b0 = smem_b[stage_idx][atom_base_idx + atom_swizzled_idx_0]; + float b1 = smem_b[stage_idx][atom_base_idx + atom_swizzled_idx_1]; + + MMASelector::type::fma( + accum[n * 4 + 0], accum[n * 4 + 1], accum[n * 4 + 2], accum[n * 4 + 3], + a_step[0], a_step[1], a_step[2], a_step[3], + b0, b1, + accum[n * 4 + 0], accum[n * 4 + 1], accum[n * 4 + 2], accum[n * 4 + 3] + ); + } + } + } + } + + const auto& reduced_sum_0 = math::warp_reduce_sum<4>(sqr_sum_acc_0); + const auto& reduced_sum_1 = math::warp_reduce_sum<4>(sqr_sum_acc_1); + + const auto& m_idx = m_block_idx * BLOCK_M + (warp_idx * BLOCK_M_PER_WARP + lane_idx / 4); + if (lane_idx % 4 == 0) { + if (m_idx < shape_m) + sqr_sum[m_offset + m_idx] = reduced_sum_0; + if (m_idx + 8 < shape_m) + sqr_sum[m_offset + m_idx + 8] = reduced_sum_1; + } + + __syncwarp(); + empty_barriers[(num_total_stages-1) % kNumStages]->arrive(); + + uint32_t is_odd_pair = lane_idx / 2 % 2; + uint32_t row_idx = lane_idx / 4; + uint32_t reordered_pair_idx = is_odd_pair * 8 + row_idx; + + auto shifted_smem_ptr = reinterpret_cast(smem_cd) + + (warp_idx * BLOCK_M_PER_WARP + row_idx) * kSwizzleCDMode + + lane_idx % 2 * 8; + + #pragma unroll + for (uint32_t i = 0; i < (kSwizzleCDMode / sizeof(float)) / 4; i += 2) { + uint32_t bank_group_idx = get_sm120_swizzled_bank_group_idx(i + is_odd_pair, reordered_pair_idx); + auto smem_ptr = shifted_smem_ptr + bank_group_idx * kNumBankGroupBytes; + + auto values = reinterpret_cast(accum + i * 2); + ptx::st_shared(smem_ptr, values[0], values[1]); + ptx::st_shared(smem_ptr + 8 * kSwizzleCDMode, values[2], values[3]); + } + cute::tma_store_fence(); + cutlass::arch::NamedBarrier::sync(128, 1); + + if (warp_idx == 0 and cute::elect_one_sync()) { + if constexpr (kNumSplits == 1) { + cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_cd, 0, m_block_idx * BLOCK_M); + } else { + cute::SM90_TMA_STORE_3D::copy(&tensor_map_d, smem_cd, 0, m_block_idx * BLOCK_M, k_split_idx); + } + cute::tma_store_arrive(); + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only support sm_120"); +#endif +} + +} // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/deep_gemm/include/deep_gemm/mma/sm120.cuh b/deep_gemm/include/deep_gemm/mma/sm120.cuh new file mode 100644 index 0000000000..91eec0efa0 --- /dev/null +++ b/deep_gemm/include/deep_gemm/mma/sm120.cuh @@ -0,0 +1,62 @@ +#pragma once + +#include +#include // Required for __float_as_uint + +namespace deep_gemm::mma::sm120 { + +CUTLASS_DEVICE void mma_m16n8k8_f32_tf32accum( + float& d0, float& d1, float& d2, float& d3, + float a0, float a1, float a2, float a3, + float b0, float b1, + float c0, float c1, float c2, float c3) { + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 " + "{%0, %1, %2, %3}, " + "{%4, %5, %6, %7}, " + "{%8, %9}, " + "{%10, %11, %12, %13};\n" + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) + // tf32 multiplicands expect .b32 registers ("r") + : "r"(__float_as_uint(a0)), "r"(__float_as_uint(a1)), "r"(__float_as_uint(a2)), "r"(__float_as_uint(a3)), + "r"(__float_as_uint(b0)), "r"(__float_as_uint(b1)), + // f32 accumulators expect .f32 registers ("f") + "f"(c0), "f"(c1), "f"(c2), "f"(c3)); +} + +template +struct TF32MMASync { + static constexpr int MMA_M = 16; + static constexpr int MMA_N = 8; + static constexpr int MMA_K = 8; + static constexpr int kNumAccum = M * N / 32; + + static_assert(M == 16 and N == 8 and K == 8, "SM120 TF32 mma.sync atom is 16x8x8"); + + CUTLASS_DEVICE static void fma( + float& d0, float& d1, float& d2, float& d3, + float a0, float a1, float a2, float a3, + float b0, float b1, + float c0, float c1, float c2, float c3) { + mma_m16n8k8_f32_tf32accum(d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, c0, c1, c2, c3); + } +}; + +template +struct TF32MMASelector { + static constexpr auto select_type() { + static_assert(N == 8 or N == 16 or N == 32, "SM120 TF32 hc_prenorm supports N <= 32"); + return TF32MMASync<16, 8, 8>{}; + } + + using type = decltype(select_type()); + + static constexpr int M = 16; + static constexpr int K = 8; + static constexpr int kNumAccumPerAtom = 4; + + static constexpr int MMA_N_per_atom() { return 8; } + static constexpr int N_atoms = N / MMA_N_per_atom(); +}; + +} // namespace deep_gemm::mma::sm120 From dd05d89b038165f1c85648d239c2b5ff14fd4ba7 Mon Sep 17 00:00:00 2001 From: jasl Date: Sun, 26 Apr 2026 03:37:43 +0800 Subject: [PATCH 06/14] Update SM120 HC regression for kernel path --- tests/test_sm120_reference.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test_sm120_reference.py b/tests/test_sm120_reference.py index a8b31cc80a..3fdad99508 100644 --- a/tests/test_sm120_reference.py +++ b/tests/test_sm120_reference.py @@ -53,7 +53,7 @@ def _paged_mqa_reference( @_test_filter(lambda: get_arch_major() >= 12) -def test_sm120_hc_prenorm_gemm_reference_path() -> None: +def test_sm120_hc_prenorm_gemm_kernel_path() -> None: torch.manual_seed(0) m, n, k = 5, 8, 64 @@ -69,8 +69,10 @@ def test_sm120_hc_prenorm_gemm_reference_path() -> None: final_d = d if num_splits is None else d.sum(dim=0) final_sqr_sum = sqr_sum if num_splits is None else sqr_sum.sum(dim=0) - torch.testing.assert_close(final_d, a.float() @ b.T, rtol=1e-4, atol=1e-4) - torch.testing.assert_close(final_sqr_sum, a.float().square().sum(dim=-1), rtol=1e-4, atol=1e-4) + ref_d = a.float() @ b.T + ref_sqr_sum = a.float().square().sum(dim=-1) + diff = max(calc_diff(final_d, ref_d), calc_diff(final_sqr_sum, ref_sqr_sum)) + assert diff < 1e-6, f"{num_splits=}, {diff=}" @_test_filter(lambda: get_arch_major() >= 12) From 5355f5196429e009748f59d599fbf39c848ce263 Mon Sep 17 00:00:00 2001 From: jasl Date: Sun, 26 Apr 2026 03:53:38 +0800 Subject: [PATCH 07/14] Strengthen SM120 fallback regression coverage Add SM120 coverage for the varlen FP8 paged MQA logits path and the FP8 bhr,hdr->bhd einsum fallback used by the DeepSeek V4 bring-up. These tests exercise the existing fallback kernels against PyTorch references with physical block tables, repeated varlen rows, and FP8 scaling-factor layouts. Also tighten the SM120 HyperConnection JIT wrapper so zero-stage split-K launches are rejected on the host. The regression now uses a valid split shape with enough K blocks, which keeps the test aligned with the kernel's split-K contract. --- .../impls/sm120_tf32_hc_prenorm_gemm.hpp | 1 + tests/test_sm120_reference.py | 105 +++++++++++++++++- 2 files changed, 105 insertions(+), 1 deletion(-) diff --git a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp index c37e1ab253..21723a23c7 100644 --- a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp @@ -81,6 +81,7 @@ static void sm120_tf32_hc_prenorm_gemm(const torch::Tensor& a, // Only support small N for now DG_HOST_ASSERT(n <= 32 and n % 8 == 0); DG_HOST_ASSERT(k % block_k == 0); + DG_HOST_ASSERT(num_splits >= 1 and num_splits <= k / block_k); const auto swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float)); const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k, diff --git a/tests/test_sm120_reference.py b/tests/test_sm120_reference.py index 3fdad99508..1b92fa4bcd 100644 --- a/tests/test_sm120_reference.py +++ b/tests/test_sm120_reference.py @@ -2,6 +2,7 @@ import deep_gemm from deep_gemm.testing import calc_diff, get_arch_major, test_filter as _test_filter +from deep_gemm.utils.math import ceil_div, per_block_cast_to_fp8, per_token_cast_to_fp8 def _cast_kv_cache_to_fp8(kv_cache: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -56,7 +57,7 @@ def _paged_mqa_reference( def test_sm120_hc_prenorm_gemm_kernel_path() -> None: torch.manual_seed(0) - m, n, k = 5, 8, 64 + m, n, k = 5, 8, 256 a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") b = torch.randn((n, k), dtype=torch.float32, device="cuda") for num_splits in (None, 3): @@ -126,3 +127,105 @@ def test_sm120_fp8_paged_mqa_logits_reference_path() -> None: diff = calc_diff(logits.float()[valid_mask], expected[valid_mask]) threshold = 1e-6 if logits_dtype == torch.float32 else 2e-6 assert diff < threshold, f"{logits_dtype=}, {diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_fp8_paged_mqa_logits_varlen_reference_path() -> None: + torch.manual_seed(2) + + raw_batch_size, num_heads, head_dim = 3, 32, 32 + block_kv = 32 + max_context_len = 80 + num_blocks = 8 + tokens_per_seq = torch.tensor([1, 2, 1], device="cuda", dtype=torch.int32) + indices = torch.arange(raw_batch_size, device="cuda", dtype=torch.int32) + indices = indices.repeat_interleave(tokens_per_seq) + batch_size = int(tokens_per_seq.sum().item()) + next_n = 1 + + base_context_lens = torch.tensor([11, 19, 7], device="cuda", dtype=torch.int32) + offsets = torch.cat([ + torch.arange(int(token_count.item()), device="cuda", dtype=torch.int32) + for token_count in tokens_per_seq + ]) + context_lens = (base_context_lens.repeat_interleave(tokens_per_seq) + offsets).view(-1, 1) + + q = torch.randn( + (batch_size, next_n, num_heads, head_dim), + device="cuda", + dtype=torch.bfloat16, + ) + q_fp8 = q.to(torch.float8_e4m3fn) + q_simulated = q_fp8.float() + kv_cache = torch.randn((num_blocks, block_kv, 1, head_dim), device="cuda", dtype=torch.bfloat16) + fused_kv_cache, kv_simulated = _cast_kv_cache_to_fp8(kv_cache) + weights = torch.randn((batch_size * next_n, num_heads), device="cuda", dtype=torch.float32) + + base_block_table = torch.tensor( + [ + [0, 1, 2], + [3, 4, 5], + [6, 7, 0], + ], + device="cuda", + dtype=torch.int32, + ) + block_table = base_block_table.repeat_interleave(tokens_per_seq, dim=0) + + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens, + block_kv, + deep_gemm.get_num_sms(), + indices=indices, + ) + expected = _paged_mqa_reference( + q_simulated, + kv_simulated, + weights, + context_lens, + block_table, + max_context_len, + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + q=(q_fp8, None), + kv_cache=fused_kv_cache, + weights=weights, + context_lens=context_lens, + block_table=block_table, + schedule_meta=schedule_meta, + max_context_len=max_context_len, + clean_logits=False, + logits_dtype=torch.float32, + indices=indices, + ) + + valid_mask = torch.arange(max_context_len, device="cuda").unsqueeze(0) < context_lens + diff = calc_diff(logits.float()[valid_mask], expected[valid_mask]) + assert diff < 1e-6, f"{diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_fp8_einsum_reference_path() -> None: + torch.manual_seed(3) + + batch_size, num_heads, rank, output_dim = 3, 4, 256, 128 + x = torch.randn((batch_size, num_heads, rank), device="cuda", dtype=torch.bfloat16) + y = torch.randn((num_heads, output_dim, rank), device="cuda", dtype=torch.bfloat16) + expected = torch.einsum("bhr,hdr->bhd", x, y) + + x_fp8, x_sf = per_token_cast_to_fp8(x.view(-1, rank), use_ue8m0=False) + x_in = x_fp8.view(batch_size, num_heads, rank), x_sf.view(batch_size, num_heads, ceil_div(rank, 128)) + y_fp8 = torch.empty_like(y, dtype=torch.float8_e4m3fn) + y_sf = torch.empty( + (num_heads, ceil_div(output_dim, 128), ceil_div(rank, 128)), + device="cuda", + dtype=torch.float32, + ) + for head_idx in range(num_heads): + y_fp8[head_idx], y_sf[head_idx] = per_block_cast_to_fp8(y[head_idx], use_ue8m0=False) + + actual = torch.empty((batch_size, num_heads, output_dim), device="cuda", dtype=torch.bfloat16) + deep_gemm.fp8_einsum("bhr,hdr->bhd", x_in, (y_fp8, y_sf), actual) + + diff = calc_diff(actual, expected) + assert diff < 1e-3, f"{diff=}" From 6d624fb63346c2f4e83d075dcc75f112601a9cb0 Mon Sep 17 00:00:00 2001 From: jasl Date: Sun, 26 Apr 2026 04:06:50 +0800 Subject: [PATCH 08/14] Vectorize SM120 FP8 paged MQA dot loop Replace the scalar per-dimension FP8 dot-product loop in the SM120 paged MQA fallback with fp8x4 loads and conversion. The launch geometry stays unchanged, which avoids the CTA-count regression observed with the warp-per-logit experiment while reducing conversion/load overhead inside each logit calculation. Add a V4-shaped SM120 regression with 64 heads and 128-dimensional FP8 Q/KV to keep the vectorized path covered alongside the existing small-shape and varlen cases. --- .../impls/sm120_fp8_paged_mqa_logits.cuh | 28 ++++++++--- tests/test_sm120_reference.py | 47 +++++++++++++++++++ 2 files changed, 68 insertions(+), 7 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh index 28adffa1a8..ef1f48f66c 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh @@ -8,6 +8,26 @@ namespace deep_gemm { +template +CUTLASS_DEVICE float sm120_fp8_dot_scaled( + const __nv_fp8_e4m3* q, + const __nv_fp8_e4m3* kv, + const float kv_scale) { + float score = 0.0f; + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += 4) { + const auto q_values = static_cast( + *reinterpret_cast(q + dim_idx)); + const auto kv_values = static_cast( + *reinterpret_cast(kv + dim_idx)); + score = fmaf(q_values.x, kv_values.x * kv_scale, score); + score = fmaf(q_values.y, kv_values.y * kv_scale, score); + score = fmaf(q_values.z, kv_values.z * kv_scale, score); + score = fmaf(q_values.w, kv_values.w * kv_scale, score); + } + return score; +} + template @@ -59,13 +79,7 @@ void sm120_fp8_paged_mqa_logits_reference( #pragma unroll for (uint32_t head_idx = 0; head_idx < kNumHeads; ++head_idx) { const auto q_head = q_base + head_idx * q_head_stride; - float score = 0.0f; - #pragma unroll - for (uint32_t dim_idx = 0; dim_idx < kHeadDim; ++dim_idx) { - const auto q_value = static_cast(q_head[dim_idx]); - const auto kv_value = static_cast(kv_base[dim_idx]) * kv_scale; - score = fmaf(q_value, kv_value, score); - } + const auto score = sm120_fp8_dot_scaled(q_head, kv_base, kv_scale); total += fmaxf(score, 0.0f) * weight_base[head_idx]; } diff --git a/tests/test_sm120_reference.py b/tests/test_sm120_reference.py index 1b92fa4bcd..9d3537213d 100644 --- a/tests/test_sm120_reference.py +++ b/tests/test_sm120_reference.py @@ -129,6 +129,53 @@ def test_sm120_fp8_paged_mqa_logits_reference_path() -> None: assert diff < threshold, f"{logits_dtype=}, {diff=}" +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_fp8_paged_mqa_logits_v4_shape_reference_path() -> None: + torch.manual_seed(4) + + batch_size, next_n, num_heads, head_dim = 1, 1, 64, 128 + block_kv = 64 + max_context_len = 65 + num_blocks = 2 + + q = torch.randn((batch_size, next_n, num_heads, head_dim), device="cuda", dtype=torch.bfloat16) + q_fp8 = q.to(torch.float8_e4m3fn) + q_simulated = q_fp8.float() + kv_cache = torch.randn((num_blocks, block_kv, 1, head_dim), device="cuda", dtype=torch.bfloat16) + fused_kv_cache, kv_simulated = _cast_kv_cache_to_fp8(kv_cache) + weights = torch.randn((batch_size * next_n, num_heads), device="cuda", dtype=torch.float32) + context_lens = torch.tensor([[max_context_len]], device="cuda", dtype=torch.int32) + block_table = torch.tensor([[0, 1]], device="cuda", dtype=torch.int32) + + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens, + block_kv, + deep_gemm.get_num_sms(), + ) + expected = _paged_mqa_reference( + q_simulated, + kv_simulated, + weights, + context_lens, + block_table, + max_context_len, + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + q=(q_fp8, None), + kv_cache=fused_kv_cache, + weights=weights, + context_lens=context_lens, + block_table=block_table, + schedule_meta=schedule_meta, + max_context_len=max_context_len, + clean_logits=False, + logits_dtype=torch.float32, + ) + + diff = calc_diff(logits.float(), expected) + assert diff < 1e-6, f"{diff=}" + + @_test_filter(lambda: get_arch_major() >= 12) def test_sm120_fp8_paged_mqa_logits_varlen_reference_path() -> None: torch.manual_seed(2) From 47369c1d2df7ed36c685dec77cec217a90a94a18 Mon Sep 17 00:00:00 2001 From: jasl Date: Sun, 26 Apr 2026 04:20:30 +0800 Subject: [PATCH 09/14] Vectorize SM120 FP8 einsum fallback Use fp8x4 loads and conversion in the SM120 bhr,hdr->bhd fallback while keeping the existing launch geometry. The vectorized path is gated on contiguous rank strides and a rank size divisible by four, with the original scalar path retained for irregular shapes. Add a V4-shaped regression covering 64 heads, rank 128, and output dim 128. On GB10 this improves representative event-timed shapes from roughly 0.99 TFLOPS to about 4.0 TFLOPS. --- .../deep_gemm/impls/sm120_fp8_einsum.cuh | 86 ++++++++++++++++--- tests/test_sm120_reference.py | 23 +++++ 2 files changed, 97 insertions(+), 12 deletions(-) diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh index 601410bd72..2e09a876c4 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh @@ -8,6 +8,70 @@ namespace deep_gemm { +CUTLASS_DEVICE float sm120_fp8_einsum_dot_scalar( + const __nv_fp8_e4m3* a, + const __nv_fp8_e4m3* b, + const float* sfa, + const float* sfb, + const uint32_t shape_r, + const uint64_t sfa_stride_r, + const uint64_t sfb_stride_r) { + float accum = 0.0f; + for (uint32_t r_idx = 0; r_idx < shape_r; ++r_idx) { + const uint32_t r_scale_idx = r_idx / 128; + const float scale = sfa[r_scale_idx * sfa_stride_r] * sfb[r_scale_idx * sfb_stride_r]; + const float a_value = static_cast(a[r_idx]); + const float b_value = static_cast(b[r_idx]); + accum = fmaf(a_value, b_value * scale, accum); + } + return accum; +} + +CUTLASS_DEVICE float sm120_fp8_einsum_dot_fp8x4( + const __nv_fp8_e4m3* a, + const __nv_fp8_e4m3* b, + const float* sfa, + const float* sfb, + const uint32_t shape_r, + const uint64_t sfa_stride_r, + const uint64_t sfb_stride_r) { + float accum = 0.0f; + const uint32_t num_full_scale_blocks = shape_r / 128; + for (uint32_t scale_idx = 0; scale_idx < num_full_scale_blocks; ++scale_idx) { + const float scale = sfa[scale_idx * sfa_stride_r] * sfb[scale_idx * sfb_stride_r]; + const uint32_t r_start = scale_idx * 128; + #pragma unroll + for (uint32_t r_offset = 0; r_offset < 128; r_offset += 4) { + const uint32_t r_idx = r_start + r_offset; + const auto a_values = static_cast( + *reinterpret_cast(a + r_idx)); + const auto b_values = static_cast( + *reinterpret_cast(b + r_idx)); + accum = fmaf(a_values.x, b_values.x * scale, accum); + accum = fmaf(a_values.y, b_values.y * scale, accum); + accum = fmaf(a_values.z, b_values.z * scale, accum); + accum = fmaf(a_values.w, b_values.w * scale, accum); + } + } + const uint32_t tail_start = num_full_scale_blocks * 128; + if (tail_start >= shape_r) + return accum; + + const float tail_scale = sfa[num_full_scale_blocks * sfa_stride_r] * + sfb[num_full_scale_blocks * sfb_stride_r]; + for (uint32_t r_idx = tail_start; r_idx < shape_r; r_idx += 4) { + const auto a_values = static_cast( + *reinterpret_cast(a + r_idx)); + const auto b_values = static_cast( + *reinterpret_cast(b + r_idx)); + accum = fmaf(a_values.x, b_values.x * tail_scale, accum); + accum = fmaf(a_values.y, b_values.y * tail_scale, accum); + accum = fmaf(a_values.z, b_values.z * tail_scale, accum); + accum = fmaf(a_values.w, b_values.w * tail_scale, accum); + } + return accum; +} + template CUTLASS_GLOBAL void sm120_fp8_bhr_hdr_bhd_reference( const uint32_t shape_b, @@ -44,18 +108,16 @@ CUTLASS_GLOBAL void sm120_fp8_bhr_hdr_bhd_reference( if (b_idx >= shape_b or d_idx >= shape_d) return; - float accum = 0.0f; - for (uint32_t r_idx = 0; r_idx < shape_r; ++r_idx) { - const uint32_t r_scale_idx = r_idx / 128; - const auto a_offset = b_idx * a_stride_b + h_idx * a_stride_h + r_idx * a_stride_r; - const auto b_offset = h_idx * b_stride_h + d_idx * b_stride_d + r_idx * b_stride_r; - const auto sfa_offset = b_idx * sfa_stride_b + h_idx * sfa_stride_h + r_scale_idx * sfa_stride_r; - const auto sfb_offset = h_idx * sfb_stride_h + (d_idx / 128) * sfb_stride_d + r_scale_idx * sfb_stride_r; - - const float a_value = static_cast(a[a_offset]) * sfa[sfa_offset]; - const float b_value = static_cast(b[b_offset]) * sfb[sfb_offset]; - accum = fmaf(a_value, b_value, accum); - } + const auto a_base = b_idx * a_stride_b + h_idx * a_stride_h; + const auto b_base = h_idx * b_stride_h + d_idx * b_stride_d; + const auto sfa_base = b_idx * sfa_stride_b + h_idx * sfa_stride_h; + const auto sfb_base = h_idx * sfb_stride_h + (d_idx / 128) * sfb_stride_d; + const bool can_vectorize = (shape_r % 4 == 0) and (a_stride_r == 1) and (b_stride_r == 1); + const float accum = can_vectorize ? + sm120_fp8_einsum_dot_fp8x4(a + a_base, b + b_base, sfa + sfa_base, sfb + sfb_base, + shape_r, sfa_stride_r, sfb_stride_r) : + sm120_fp8_einsum_dot_scalar(a + a_base, b + b_base, sfa + sfa_base, sfb + sfb_base, + shape_r, sfa_stride_r, sfb_stride_r); const auto d_offset = b_idx * d_stride_b + h_idx * d_stride_h + d_idx * d_stride_d; d[d_offset] = static_cast(accum); diff --git a/tests/test_sm120_reference.py b/tests/test_sm120_reference.py index 9d3537213d..641fa328e6 100644 --- a/tests/test_sm120_reference.py +++ b/tests/test_sm120_reference.py @@ -276,3 +276,26 @@ def test_sm120_fp8_einsum_reference_path() -> None: diff = calc_diff(actual, expected) assert diff < 1e-3, f"{diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_fp8_einsum_v4_shape_reference_path() -> None: + torch.manual_seed(5) + + batch_size, num_heads, rank, output_dim = 5, 64, 128, 128 + x = torch.randn((batch_size, num_heads, rank), device="cuda", dtype=torch.bfloat16) + y = torch.randn((num_heads, output_dim, rank), device="cuda", dtype=torch.bfloat16) + expected = torch.einsum("bhr,hdr->bhd", x, y) + + x_fp8, x_sf = per_token_cast_to_fp8(x.view(-1, rank), use_ue8m0=False) + x_in = x_fp8.view(batch_size, num_heads, rank), x_sf.view(batch_size, num_heads, 1) + y_fp8 = torch.empty_like(y, dtype=torch.float8_e4m3fn) + y_sf = torch.empty((num_heads, 1, 1), device="cuda", dtype=torch.float32) + for head_idx in range(num_heads): + y_fp8[head_idx], y_sf[head_idx] = per_block_cast_to_fp8(y[head_idx], use_ue8m0=False) + + actual = torch.empty((batch_size, num_heads, output_dim), device="cuda", dtype=torch.bfloat16) + deep_gemm.fp8_einsum("bhr,hdr->bhd", x_in, (y_fp8, y_sf), actual) + + diff = calc_diff(actual, expected) + assert diff < 1e-3, f"{diff=}" From d863f7bfd8c1e53a958a42dc5c62248d56062bff Mon Sep 17 00:00:00 2001 From: jasl Date: Sun, 26 Apr 2026 04:39:29 +0800 Subject: [PATCH 10/14] Cover larger SM120 HC prenorm GEMM tiles Add regression coverage for the SM120 tf32_hc_prenorm_gemm kernel beyond the original small m/n shape. The new cases cover n=16 and n=32, multiple 64-row M tiles, and split-K accumulation so future HC tuning has a stronger correctness guard. --- tests/test_sm120_reference.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_sm120_reference.py b/tests/test_sm120_reference.py index 641fa328e6..2953f8f8f8 100644 --- a/tests/test_sm120_reference.py +++ b/tests/test_sm120_reference.py @@ -76,6 +76,28 @@ def test_sm120_hc_prenorm_gemm_kernel_path() -> None: assert diff < 1e-6, f"{num_splits=}, {diff=}" +@_test_filter(lambda: get_arch_major() >= 12) +def test_sm120_hc_prenorm_gemm_larger_tiles() -> None: + torch.manual_seed(6) + + for m, n, k, num_splits in ( + (70, 16, 512, 2), + (70, 32, 512, 4), + (129, 32, 1024, 4), + ): + a = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") + b = torch.randn((n, k), dtype=torch.float32, device="cuda") + d = torch.empty((num_splits, m, n), dtype=torch.float32, device="cuda") + sqr_sum = torch.empty((num_splits, m), dtype=torch.float32, device="cuda") + + deep_gemm.tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits=num_splits) + + ref_d = a.float() @ b.T + ref_sqr_sum = a.float().square().sum(dim=-1) + diff = max(calc_diff(d.sum(dim=0), ref_d), calc_diff(sqr_sum.sum(dim=0), ref_sqr_sum)) + assert diff < 1e-6, f"{m=}, {n=}, {k=}, {num_splits=}, {diff=}" + + @_test_filter(lambda: get_arch_major() >= 12) def test_sm120_fp8_paged_mqa_logits_reference_path() -> None: torch.manual_seed(1) From e4194ac143e2cdcfa8e14c31376678d162aab59f Mon Sep 17 00:00:00 2001 From: jasl Date: Sun, 26 Apr 2026 06:55:38 +0800 Subject: [PATCH 11/14] Add experimental SM120 tiled paged MQA kernel Use explicit NVCC gencode targets for JIT builds so SM12x family-specific MMA instructions do not get rejected while compiling generic PTX. Add an env-gated SM120 FP8 paged MQA tiled kernel for the V4 shape: 64 heads, head_dim 128, block_kv 64, and float logits. The tiled path uses SM120 E4M3 QMMA, supports token group sizes 1/2/4/8, and defaults to a Q-cached 4-group schedule under DG_SM120_PAGED_MQA_TILED=1. Keep the scalar reference fallback as the default path for now and add CUDA parity coverage for all tiled variants. On DGX Spark/GB10, the V4-shaped component benchmark improved from 1.918 ms to 0.677 ms for batch=512 ctx=512, and from 1.906 ms to 0.678 ms for batch=128 ctx=2048. --- csrc/jit/compiler.hpp | 5 +- .../impls/smxx_fp8_fp4_paged_mqa_logits.hpp | 78 ++++++ .../impls/sm120_fp8_paged_mqa_logits.cuh | 228 ++++++++++++++++++ tests/test_sm120_reference.py | 56 ++++- 4 files changed, 364 insertions(+), 3 deletions(-) diff --git a/csrc/jit/compiler.hpp b/csrc/jit/compiler.hpp index 7d85a5f556..e1d1f20114 100644 --- a/csrc/jit/compiler.hpp +++ b/csrc/jit/compiler.hpp @@ -202,10 +202,11 @@ class NVCCCompiler final: public Compiler { // The override the compiler flags // Only NVCC >= 12.9 supports arch-specific family suffix const auto arch = device_runtime->get_arch(false, nvcc_major > 12 or nvcc_minor >= 9); - flags = fmt::format("{} -I{} --gpu-architecture=sm_{} " + const auto arch_flag = fmt::format("-gencode=arch=compute_{},code=sm_{}", arch, arch); + flags = fmt::format("{} -I{} {} " "--compiler-options=-fPIC,-O3,-fconcepts,-Wno-deprecated-declarations,-Wno-abi " "-O3 --expt-relaxed-constexpr --expt-extended-lambda", - flags, library_include_path.c_str(), arch); + flags, library_include_path.c_str(), arch_flag); } void compile(const std::string &code, const std::filesystem::path& dir_path, diff --git a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp index d966449678..b35ef4f587 100644 --- a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp @@ -178,6 +178,8 @@ class SM120FP8PagedMQALogitsReferenceRuntime final: public LaunchRuntime { +public: + using Args = SM120FP8PagedMQALogitsReferenceRuntime::Args; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm120_fp8_paged_mqa_logits_mma_tiled< + {}, + {}, + {}, + {}, + {} + >); +}}; +)", args.block_kv, args.is_context_lens_2d ? "true" : "false", + args.token_groups, + args.cache_q ? "true" : "false", + to_string(args.logits_dtype)); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.batch_size, + args.next_n, + args.logits_stride, + args.block_table_stride, + args.q_batch_stride, + args.q_next_stride, + args.q_head_stride, + args.kv_block_stride, + args.kv_token_stride, + args.kv_scale_block_stride, + args.weights_row_stride, + args.q, + args.kv_cache, + args.kv_cache_scales, + args.weights, + args.context_lens, + args.logits, + args.block_table + )); + } +}; + static void sm120_fp8_paged_mqa_logits_reference(const torch::Tensor& q, const torch::Tensor& kv_cache, const torch::Tensor& kv_cache_scales, @@ -270,6 +321,8 @@ static void sm120_fp8_paged_mqa_logits_reference(const torch::Tensor& q, .block_table_stride = block_table_stride, .logits_stride = logits_stride, .tokens_per_block = tokens_per_block, + .token_groups = 1, + .cache_q = false, .q_batch_stride = q.stride(0), .q_next_stride = q.stride(1), .q_head_stride = q.stride(2), @@ -288,6 +341,31 @@ static void sm120_fp8_paged_mqa_logits_reference(const torch::Tensor& q, .launch_args = LaunchArgs(std::make_pair(batch_size * next_n, ceil_div(logits_stride, tokens_per_block)), tokens_per_block) }; + const bool use_tiled = get_env("DG_SM120_PAGED_MQA_TILED", 0) != 0; + if (use_tiled and num_heads == 64 and head_dim == 128 and block_kv == 64 and logits_dtype == torch::kFloat32) { + const int token_groups = get_env("DG_SM120_PAGED_MQA_TILED_GROUPS", 4); + DG_HOST_ASSERT(token_groups == 1 or token_groups == 2 or token_groups == 4 or token_groups == 8); + constexpr int tiled_token_tile = 8; + constexpr int tiled_num_threads = 128; + const int tokens_per_tiled_cta = tiled_token_tile * token_groups; + const int tiled_smem_size = 64 * tokens_per_tiled_cta * sizeof(float); + auto tiled_args = args; + tiled_args.token_groups = token_groups; + const bool should_cache_q_by_default = token_groups == 4; + tiled_args.cache_q = get_env( + "DG_SM120_PAGED_MQA_TILED_CACHE_Q", + should_cache_q_by_default ? 1 : 0) != 0; + DG_HOST_ASSERT(not tiled_args.cache_q or token_groups == 4); + tiled_args.launch_args = LaunchArgs( + std::make_pair(batch_size * next_n, ceil_div(logits_stride, tokens_per_tiled_cta)), + tiled_num_threads, + tiled_smem_size); + const auto code = SM120FP8PagedMQALogitsTiledRuntime::generate(tiled_args); + const auto runtime = compiler->build("sm120_fp8_paged_mqa_logits_mma_tiled", code); + SM120FP8PagedMQALogitsTiledRuntime::launch(runtime, tiled_args); + return; + } + const auto code = SM120FP8PagedMQALogitsReferenceRuntime::generate(args); const auto runtime = compiler->build("sm120_fp8_paged_mqa_logits_reference", code); SM120FP8PagedMQALogitsReferenceRuntime::launch(runtime, args); diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh index ef1f48f66c..4ed1b2bf65 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh @@ -2,6 +2,9 @@ #include +#include +#include + #include #include @@ -86,4 +89,229 @@ void sm120_fp8_paged_mqa_logits_reference( logits[output_offset] = static_cast(total); } +template +CUTLASS_GLOBAL __launch_bounds__(128, 1) +void sm120_fp8_paged_mqa_logits_mma_tiled( + const uint32_t batch_size, + const uint32_t next_n, + const uint32_t logits_stride, + const uint32_t block_table_stride, + const uint64_t q_batch_stride, + const uint64_t q_next_stride, + const uint64_t q_head_stride, + const uint64_t kv_block_stride, + const uint64_t kv_token_stride, + const uint64_t kv_scale_block_stride, + const uint64_t weights_row_stride, + const __nv_fp8_e4m3* q, + const __nv_fp8_e4m3* kv_cache, + const float* kv_cache_scales, + const float* weights, + const uint32_t* context_lens, + logits_dtype_t* logits, + const uint32_t* block_table) { + using Element = cute::float_e4m3_t; + using Accumulator = float; + + constexpr uint32_t kNumHeads = 64; + constexpr uint32_t kHeadDim = 128; + constexpr uint32_t kHeadsPerWarp = 16; + constexpr uint32_t kTokensPerTile = 8; + constexpr uint32_t kTokensPerCta = kTokensPerTile * kTokenGroups; + constexpr uint32_t kMmaK = 32; + constexpr uint32_t kHeadTiles = kNumHeads / kHeadsPerWarp; + + extern __shared__ float shared_scores[]; + + const auto row = blockIdx.x; + const auto token_tile_start = blockIdx.y * kTokensPerCta; + if (row >= batch_size * next_n) + return; + + const auto batch_idx = row / next_n; + const auto next_idx = row - batch_idx * next_n; + const auto context_lens_idx = kIsContextLens2D ? row : batch_idx; + const auto context_len = context_lens[context_lens_idx]; + + if (token_tile_start >= context_len) { + if (threadIdx.x < kTokensPerCta) { + const auto token_idx = token_tile_start + threadIdx.x; + if (token_idx < logits_stride) { + const auto output_offset = row * static_cast(logits_stride) + token_idx; + logits[output_offset] = static_cast(-__uint_as_float(0x7f800000u)); + } + } + return; + } + + const auto lane_idx = threadIdx.x % 32; + const auto warp_idx = threadIdx.x / 32; + const auto q_base = q + batch_idx * q_batch_stride + next_idx * q_next_stride; + + auto tiled_mma = cute::make_tiled_mma( + cute::SM120_16x8x32_TN{}); + auto thread_mma = tiled_mma.get_slice(lane_idx); + + if constexpr (kCacheQ and kTokenGroups == 4) { + auto score_tile0 = cute::make_tensor( + cute::make_smem_ptr(shared_scores + warp_idx * kHeadsPerWarp * kTokensPerCta), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(cute::Int{}, cute::Int<1>{})); + auto score_tile1 = cute::make_tensor( + cute::make_smem_ptr(shared_scores + warp_idx * kHeadsPerWarp * kTokensPerCta + kTokensPerTile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(cute::Int{}, cute::Int<1>{})); + auto score_tile2 = cute::make_tensor( + cute::make_smem_ptr(shared_scores + warp_idx * kHeadsPerWarp * kTokensPerCta + 2 * kTokensPerTile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(cute::Int{}, cute::Int<1>{})); + auto score_tile3 = cute::make_tensor( + cute::make_smem_ptr(shared_scores + warp_idx * kHeadsPerWarp * kTokensPerCta + 3 * kTokensPerTile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(cute::Int{}, cute::Int<1>{})); + auto tCgC0 = thread_mma.partition_C(score_tile0); + auto tCgC1 = thread_mma.partition_C(score_tile1); + auto tCgC2 = thread_mma.partition_C(score_tile2); + auto tCgC3 = thread_mma.partition_C(score_tile3); + auto tCrC0 = thread_mma.make_fragment_C(tCgC0); + auto tCrC1 = thread_mma.make_fragment_C(tCgC1); + auto tCrC2 = thread_mma.make_fragment_C(tCgC2); + auto tCrC3 = thread_mma.make_fragment_C(tCgC3); + cute::clear(tCrC0); + cute::clear(tCrC1); + cute::clear(tCrC2); + cute::clear(tCrC3); + + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += kMmaK) { + const auto q_tile = reinterpret_cast( + q_base + (warp_idx * kHeadsPerWarp) * q_head_stride + dim_idx); + auto q_tensor = cute::make_tensor( + cute::make_gmem_ptr(q_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(q_head_stride, cute::Int<1>{})); + auto tAgA = thread_mma.partition_A(q_tensor); + auto tArA = thread_mma.partition_fragment_A(q_tensor); + cute::copy(tAgA, tArA); + + auto accumulate_group = [&](auto token_group, auto& tCrC) { + constexpr uint32_t kGroup = decltype(token_group)::value; + const auto group_token_start = token_tile_start + kGroup * kTokensPerTile; + if (group_token_start >= context_len or group_token_start >= logits_stride) + return; + + const auto logical_block_idx = group_token_start / BLOCK_KV; + const auto token_in_block = group_token_start - logical_block_idx * BLOCK_KV; + const auto kv_block_idx = block_table[ + batch_idx * static_cast(block_table_stride) + logical_block_idx]; + const auto kv_base = kv_cache + kv_block_idx * kv_block_stride + token_in_block * kv_token_stride; + const auto kv_tile = reinterpret_cast(kv_base + dim_idx); + auto kv_tensor = cute::make_tensor( + cute::make_gmem_ptr(kv_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(kv_token_stride, cute::Int<1>{})); + auto tBgB = thread_mma.partition_B(kv_tensor); + auto tBrB = thread_mma.partition_fragment_B(kv_tensor); + cute::copy(tBgB, tBrB); + cute::gemm(tiled_mma, tArA, tBrB, tCrC); + }; + accumulate_group(cute::Int<0>{}, tCrC0); + accumulate_group(cute::Int<1>{}, tCrC1); + accumulate_group(cute::Int<2>{}, tCrC2); + accumulate_group(cute::Int<3>{}, tCrC3); + } + + cute::copy(tCrC0, tCgC0); + cute::copy(tCrC1, tCgC1); + cute::copy(tCrC2, tCgC2); + cute::copy(tCrC3, tCgC3); + } else { + #pragma unroll + for (uint32_t token_group = 0; token_group < kTokenGroups; ++token_group) { + const auto group_token_start = token_tile_start + token_group * kTokensPerTile; + if (group_token_start >= context_len or group_token_start >= logits_stride) + continue; + + const auto logical_block_idx = group_token_start / BLOCK_KV; + const auto token_in_block = group_token_start - logical_block_idx * BLOCK_KV; + const auto kv_block_idx = block_table[ + batch_idx * static_cast(block_table_stride) + logical_block_idx]; + const auto kv_base = kv_cache + kv_block_idx * kv_block_stride + token_in_block * kv_token_stride; + + auto score_tile = cute::make_tensor( + cute::make_smem_ptr( + shared_scores + + warp_idx * kHeadsPerWarp * kTokensPerCta + + token_group * kTokensPerTile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(cute::Int{}, cute::Int<1>{})); + auto tCgC = thread_mma.partition_C(score_tile); + auto tCrC = thread_mma.make_fragment_C(tCgC); + cute::clear(tCrC); + + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += kMmaK) { + const auto q_tile = reinterpret_cast( + q_base + (warp_idx * kHeadsPerWarp) * q_head_stride + dim_idx); + const auto kv_tile = reinterpret_cast(kv_base + dim_idx); + + auto q_tensor = cute::make_tensor( + cute::make_gmem_ptr(q_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(q_head_stride, cute::Int<1>{})); + auto kv_tensor = cute::make_tensor( + cute::make_gmem_ptr(kv_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(kv_token_stride, cute::Int<1>{})); + + auto tAgA = thread_mma.partition_A(q_tensor); + auto tBgB = thread_mma.partition_B(kv_tensor); + auto tArA = thread_mma.partition_fragment_A(q_tensor); + auto tBrB = thread_mma.partition_fragment_B(kv_tensor); + + cute::copy(tAgA, tArA); + cute::copy(tBgB, tBrB); + cute::gemm(tiled_mma, tArA, tBrB, tCrC); + } + + cute::copy(tCrC, tCgC); + } + } + __syncthreads(); + + const auto token_idx = token_tile_start + threadIdx.x; + if (threadIdx.x >= kTokensPerCta or token_idx >= logits_stride) + return; + + const auto output_offset = row * static_cast(logits_stride) + token_idx; + if (token_idx >= context_len) { + logits[output_offset] = static_cast(-__uint_as_float(0x7f800000u)); + return; + } + + const auto logical_block_idx = token_idx / BLOCK_KV; + const auto token_in_block = token_idx - logical_block_idx * BLOCK_KV; + const auto kv_block_idx = block_table[ + batch_idx * static_cast(block_table_stride) + logical_block_idx]; + const auto kv_scale = kv_cache_scales[kv_block_idx * kv_scale_block_stride + token_in_block]; + const auto weight_base = weights + row * weights_row_stride; + + float total = 0.0f; + #pragma unroll + for (uint32_t head_tile = 0; head_tile < kHeadTiles; ++head_tile) { + #pragma unroll + for (uint32_t head = 0; head < kHeadsPerWarp; ++head) { + const auto head_idx = head_tile * kHeadsPerWarp + head; + const auto score = shared_scores[ + head_tile * kHeadsPerWarp * kTokensPerCta + + head * kTokensPerCta + threadIdx.x] * kv_scale; + total += fmaxf(score, 0.0f) * weight_base[head_idx]; + } + } + + logits[output_offset] = static_cast(total); +} + } // namespace deep_gemm diff --git a/tests/test_sm120_reference.py b/tests/test_sm120_reference.py index 2953f8f8f8..b1131b33f4 100644 --- a/tests/test_sm120_reference.py +++ b/tests/test_sm120_reference.py @@ -1,4 +1,5 @@ import torch +import pytest import deep_gemm from deep_gemm.testing import calc_diff, get_arch_major, test_filter as _test_filter @@ -194,7 +195,60 @@ def test_sm120_fp8_paged_mqa_logits_v4_shape_reference_path() -> None: logits_dtype=torch.float32, ) - diff = calc_diff(logits.float(), expected) + valid_mask = torch.arange(max_context_len, device="cuda").unsqueeze(0) < context_lens + diff = calc_diff(logits.float()[valid_mask], expected[valid_mask]) + assert diff < 1e-6, f"{diff=}" + + +@_test_filter(lambda: get_arch_major() >= 12) +@pytest.mark.parametrize("token_groups, cache_q", ((1, False), (2, False), (4, False), (8, False), (4, True))) +def test_sm120_fp8_paged_mqa_logits_tiled_path(monkeypatch, token_groups: int, cache_q: bool) -> None: + torch.manual_seed(5) + monkeypatch.setenv("DG_SM120_PAGED_MQA_TILED", "1") + monkeypatch.setenv("DG_SM120_PAGED_MQA_TILED_GROUPS", str(token_groups)) + monkeypatch.setenv("DG_SM120_PAGED_MQA_TILED_CACHE_Q", "1" if cache_q else "0") + + batch_size, next_n, num_heads, head_dim = 2, 1, 64, 128 + block_kv = 64 + max_context_len = 73 + num_blocks = 4 + + q = torch.randn((batch_size, next_n, num_heads, head_dim), device="cuda", dtype=torch.bfloat16) + q_fp8 = q.to(torch.float8_e4m3fn) + q_simulated = q_fp8.float() + kv_cache = torch.randn((num_blocks, block_kv, 1, head_dim), device="cuda", dtype=torch.bfloat16) + fused_kv_cache, kv_simulated = _cast_kv_cache_to_fp8(kv_cache) + weights = torch.randn((batch_size * next_n, num_heads), device="cuda", dtype=torch.float32) + context_lens = torch.tensor([[max_context_len], [37]], device="cuda", dtype=torch.int32) + block_table = torch.tensor([[0, 1], [2, 3]], device="cuda", dtype=torch.int32) + + schedule_meta = deep_gemm.get_paged_mqa_logits_metadata( + context_lens, + block_kv, + deep_gemm.get_num_sms(), + ) + expected = _paged_mqa_reference( + q_simulated, + kv_simulated, + weights, + context_lens, + block_table, + max_context_len, + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + q=(q_fp8, None), + kv_cache=fused_kv_cache, + weights=weights, + context_lens=context_lens, + block_table=block_table, + schedule_meta=schedule_meta, + max_context_len=max_context_len, + clean_logits=False, + logits_dtype=torch.float32, + ) + + valid_mask = torch.arange(max_context_len, device="cuda").unsqueeze(0) < context_lens + diff = calc_diff(logits.float()[valid_mask], expected[valid_mask]) assert diff < 1e-6, f"{diff=}" From 3acb661381a0f0ffea12dd62914824aa5ffcb399 Mon Sep 17 00:00:00 2001 From: jasl Date: Sun, 26 Apr 2026 07:11:09 +0800 Subject: [PATCH 12/14] Tune SM120 tiled paged MQA Q-cache schedule Extend the explicit Q-cache schedule to token group sizes 2 and 8, keeping env overrides for all tested group sizes. The runtime default now uses eight token groups with Q caching enabled, which reduces repeated Q fragment loads while keeping the scalar fallback unchanged unless DG_SM120_PAGED_MQA_TILED=1 is set. DGX Spark benchmark after this change: batch=512 ctx=512 improves from 1.909 ms reference to 0.506 ms default tiled, batch=128 ctx=2048 improves from 1.877 ms to 0.504 ms, and batch=4 ctx=32768 improves from 0.970 ms to 0.251 ms. --- .../impls/smxx_fp8_fp4_paged_mqa_logits.hpp | 6 +- .../impls/sm120_fp8_paged_mqa_logits.cuh | 208 ++++++++++++------ tests/test_sm120_reference.py | 5 +- 3 files changed, 150 insertions(+), 69 deletions(-) diff --git a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp index b35ef4f587..91a3fbaeaa 100644 --- a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp @@ -343,7 +343,7 @@ static void sm120_fp8_paged_mqa_logits_reference(const torch::Tensor& q, }; const bool use_tiled = get_env("DG_SM120_PAGED_MQA_TILED", 0) != 0; if (use_tiled and num_heads == 64 and head_dim == 128 and block_kv == 64 and logits_dtype == torch::kFloat32) { - const int token_groups = get_env("DG_SM120_PAGED_MQA_TILED_GROUPS", 4); + const int token_groups = get_env("DG_SM120_PAGED_MQA_TILED_GROUPS", 8); DG_HOST_ASSERT(token_groups == 1 or token_groups == 2 or token_groups == 4 or token_groups == 8); constexpr int tiled_token_tile = 8; constexpr int tiled_num_threads = 128; @@ -351,11 +351,11 @@ static void sm120_fp8_paged_mqa_logits_reference(const torch::Tensor& q, const int tiled_smem_size = 64 * tokens_per_tiled_cta * sizeof(float); auto tiled_args = args; tiled_args.token_groups = token_groups; - const bool should_cache_q_by_default = token_groups == 4; + const bool should_cache_q_by_default = token_groups != 1; tiled_args.cache_q = get_env( "DG_SM120_PAGED_MQA_TILED_CACHE_Q", should_cache_q_by_default ? 1 : 0) != 0; - DG_HOST_ASSERT(not tiled_args.cache_q or token_groups == 4); + DG_HOST_ASSERT(not tiled_args.cache_q or token_groups == 2 or token_groups == 4 or token_groups == 8); tiled_args.launch_args = LaunchArgs( std::make_pair(batch_size * next_n, ceil_div(logits_stride, tokens_per_tiled_cta)), tiled_num_threads, diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh index 4ed1b2bf65..b3544517e8 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh @@ -154,79 +154,157 @@ void sm120_fp8_paged_mqa_logits_mma_tiled( cute::SM120_16x8x32_TN{}); auto thread_mma = tiled_mma.get_slice(lane_idx); - if constexpr (kCacheQ and kTokenGroups == 4) { - auto score_tile0 = cute::make_tensor( - cute::make_smem_ptr(shared_scores + warp_idx * kHeadsPerWarp * kTokensPerCta), - cute::make_shape(cute::Int{}, cute::Int{}), - cute::make_stride(cute::Int{}, cute::Int<1>{})); - auto score_tile1 = cute::make_tensor( - cute::make_smem_ptr(shared_scores + warp_idx * kHeadsPerWarp * kTokensPerCta + kTokensPerTile), - cute::make_shape(cute::Int{}, cute::Int{}), - cute::make_stride(cute::Int{}, cute::Int<1>{})); - auto score_tile2 = cute::make_tensor( - cute::make_smem_ptr(shared_scores + warp_idx * kHeadsPerWarp * kTokensPerCta + 2 * kTokensPerTile), - cute::make_shape(cute::Int{}, cute::Int{}), - cute::make_stride(cute::Int{}, cute::Int<1>{})); - auto score_tile3 = cute::make_tensor( - cute::make_smem_ptr(shared_scores + warp_idx * kHeadsPerWarp * kTokensPerCta + 3 * kTokensPerTile), - cute::make_shape(cute::Int{}, cute::Int{}), - cute::make_stride(cute::Int{}, cute::Int<1>{})); + if constexpr (kCacheQ and (kTokenGroups == 2 or kTokenGroups == 4 or kTokenGroups == 8)) { + auto make_score_tile = [&](auto token_group) { + constexpr uint32_t kGroup = decltype(token_group)::value; + return cute::make_tensor( + cute::make_smem_ptr( + shared_scores + + warp_idx * kHeadsPerWarp * kTokensPerCta + + kGroup * kTokensPerTile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(cute::Int{}, cute::Int<1>{})); + }; + + auto score_tile0 = make_score_tile(cute::Int<0>{}); + auto score_tile1 = make_score_tile(cute::Int<1>{}); auto tCgC0 = thread_mma.partition_C(score_tile0); auto tCgC1 = thread_mma.partition_C(score_tile1); - auto tCgC2 = thread_mma.partition_C(score_tile2); - auto tCgC3 = thread_mma.partition_C(score_tile3); auto tCrC0 = thread_mma.make_fragment_C(tCgC0); auto tCrC1 = thread_mma.make_fragment_C(tCgC1); - auto tCrC2 = thread_mma.make_fragment_C(tCgC2); - auto tCrC3 = thread_mma.make_fragment_C(tCgC3); cute::clear(tCrC0); cute::clear(tCrC1); - cute::clear(tCrC2); - cute::clear(tCrC3); - #pragma unroll - for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += kMmaK) { - const auto q_tile = reinterpret_cast( - q_base + (warp_idx * kHeadsPerWarp) * q_head_stride + dim_idx); - auto q_tensor = cute::make_tensor( - cute::make_gmem_ptr(q_tile), - cute::make_shape(cute::Int{}, cute::Int{}), - cute::make_stride(q_head_stride, cute::Int<1>{})); - auto tAgA = thread_mma.partition_A(q_tensor); - auto tArA = thread_mma.partition_fragment_A(q_tensor); - cute::copy(tAgA, tArA); - - auto accumulate_group = [&](auto token_group, auto& tCrC) { - constexpr uint32_t kGroup = decltype(token_group)::value; - const auto group_token_start = token_tile_start + kGroup * kTokensPerTile; - if (group_token_start >= context_len or group_token_start >= logits_stride) - return; - - const auto logical_block_idx = group_token_start / BLOCK_KV; - const auto token_in_block = group_token_start - logical_block_idx * BLOCK_KV; - const auto kv_block_idx = block_table[ - batch_idx * static_cast(block_table_stride) + logical_block_idx]; - const auto kv_base = kv_cache + kv_block_idx * kv_block_stride + token_in_block * kv_token_stride; - const auto kv_tile = reinterpret_cast(kv_base + dim_idx); - auto kv_tensor = cute::make_tensor( - cute::make_gmem_ptr(kv_tile), - cute::make_shape(cute::Int{}, cute::Int{}), - cute::make_stride(kv_token_stride, cute::Int<1>{})); - auto tBgB = thread_mma.partition_B(kv_tensor); - auto tBrB = thread_mma.partition_fragment_B(kv_tensor); - cute::copy(tBgB, tBrB); - cute::gemm(tiled_mma, tArA, tBrB, tCrC); - }; - accumulate_group(cute::Int<0>{}, tCrC0); - accumulate_group(cute::Int<1>{}, tCrC1); - accumulate_group(cute::Int<2>{}, tCrC2); - accumulate_group(cute::Int<3>{}, tCrC3); - } + auto accumulate_group = [&](const uint32_t dim_idx, const auto& tArA, auto token_group, auto& tCrC) { + constexpr uint32_t kGroup = decltype(token_group)::value; + const auto group_token_start = token_tile_start + kGroup * kTokensPerTile; + if (group_token_start >= context_len or group_token_start >= logits_stride) + return; + + const auto logical_block_idx = group_token_start / BLOCK_KV; + const auto token_in_block = group_token_start - logical_block_idx * BLOCK_KV; + const auto kv_block_idx = block_table[ + batch_idx * static_cast(block_table_stride) + logical_block_idx]; + const auto kv_base = kv_cache + kv_block_idx * kv_block_stride + token_in_block * kv_token_stride; + const auto kv_tile = reinterpret_cast(kv_base + dim_idx); + auto kv_tensor = cute::make_tensor( + cute::make_gmem_ptr(kv_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(kv_token_stride, cute::Int<1>{})); + auto tBgB = thread_mma.partition_B(kv_tensor); + auto tBrB = thread_mma.partition_fragment_B(kv_tensor); + cute::copy(tBgB, tBrB); + cute::gemm(tiled_mma, tArA, tBrB, tCrC); + }; + + if constexpr (kTokenGroups == 2) { + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += kMmaK) { + const auto q_tile = reinterpret_cast( + q_base + (warp_idx * kHeadsPerWarp) * q_head_stride + dim_idx); + auto q_tensor = cute::make_tensor( + cute::make_gmem_ptr(q_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(q_head_stride, cute::Int<1>{})); + auto tAgA = thread_mma.partition_A(q_tensor); + auto tArA = thread_mma.partition_fragment_A(q_tensor); + cute::copy(tAgA, tArA); + + accumulate_group(dim_idx, tArA, cute::Int<0>{}, tCrC0); + accumulate_group(dim_idx, tArA, cute::Int<1>{}, tCrC1); + } + + cute::copy(tCrC0, tCgC0); + cute::copy(tCrC1, tCgC1); + } else if constexpr (kTokenGroups == 4) { + auto score_tile2 = make_score_tile(cute::Int<2>{}); + auto score_tile3 = make_score_tile(cute::Int<3>{}); + auto tCgC2 = thread_mma.partition_C(score_tile2); + auto tCgC3 = thread_mma.partition_C(score_tile3); + auto tCrC2 = thread_mma.make_fragment_C(tCgC2); + auto tCrC3 = thread_mma.make_fragment_C(tCgC3); + cute::clear(tCrC2); + cute::clear(tCrC3); + + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += kMmaK) { + const auto q_tile = reinterpret_cast( + q_base + (warp_idx * kHeadsPerWarp) * q_head_stride + dim_idx); + auto q_tensor = cute::make_tensor( + cute::make_gmem_ptr(q_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(q_head_stride, cute::Int<1>{})); + auto tAgA = thread_mma.partition_A(q_tensor); + auto tArA = thread_mma.partition_fragment_A(q_tensor); + cute::copy(tAgA, tArA); + + accumulate_group(dim_idx, tArA, cute::Int<0>{}, tCrC0); + accumulate_group(dim_idx, tArA, cute::Int<1>{}, tCrC1); + accumulate_group(dim_idx, tArA, cute::Int<2>{}, tCrC2); + accumulate_group(dim_idx, tArA, cute::Int<3>{}, tCrC3); + } + + cute::copy(tCrC0, tCgC0); + cute::copy(tCrC1, tCgC1); + cute::copy(tCrC2, tCgC2); + cute::copy(tCrC3, tCgC3); + } else { + auto score_tile2 = make_score_tile(cute::Int<2>{}); + auto score_tile3 = make_score_tile(cute::Int<3>{}); + auto score_tile4 = make_score_tile(cute::Int<4>{}); + auto score_tile5 = make_score_tile(cute::Int<5>{}); + auto score_tile6 = make_score_tile(cute::Int<6>{}); + auto score_tile7 = make_score_tile(cute::Int<7>{}); + auto tCgC2 = thread_mma.partition_C(score_tile2); + auto tCgC3 = thread_mma.partition_C(score_tile3); + auto tCgC4 = thread_mma.partition_C(score_tile4); + auto tCgC5 = thread_mma.partition_C(score_tile5); + auto tCgC6 = thread_mma.partition_C(score_tile6); + auto tCgC7 = thread_mma.partition_C(score_tile7); + auto tCrC2 = thread_mma.make_fragment_C(tCgC2); + auto tCrC3 = thread_mma.make_fragment_C(tCgC3); + auto tCrC4 = thread_mma.make_fragment_C(tCgC4); + auto tCrC5 = thread_mma.make_fragment_C(tCgC5); + auto tCrC6 = thread_mma.make_fragment_C(tCgC6); + auto tCrC7 = thread_mma.make_fragment_C(tCgC7); + cute::clear(tCrC2); + cute::clear(tCrC3); + cute::clear(tCrC4); + cute::clear(tCrC5); + cute::clear(tCrC6); + cute::clear(tCrC7); - cute::copy(tCrC0, tCgC0); - cute::copy(tCrC1, tCgC1); - cute::copy(tCrC2, tCgC2); - cute::copy(tCrC3, tCgC3); + #pragma unroll + for (uint32_t dim_idx = 0; dim_idx < kHeadDim; dim_idx += kMmaK) { + const auto q_tile = reinterpret_cast( + q_base + (warp_idx * kHeadsPerWarp) * q_head_stride + dim_idx); + auto q_tensor = cute::make_tensor( + cute::make_gmem_ptr(q_tile), + cute::make_shape(cute::Int{}, cute::Int{}), + cute::make_stride(q_head_stride, cute::Int<1>{})); + auto tAgA = thread_mma.partition_A(q_tensor); + auto tArA = thread_mma.partition_fragment_A(q_tensor); + cute::copy(tAgA, tArA); + + accumulate_group(dim_idx, tArA, cute::Int<0>{}, tCrC0); + accumulate_group(dim_idx, tArA, cute::Int<1>{}, tCrC1); + accumulate_group(dim_idx, tArA, cute::Int<2>{}, tCrC2); + accumulate_group(dim_idx, tArA, cute::Int<3>{}, tCrC3); + accumulate_group(dim_idx, tArA, cute::Int<4>{}, tCrC4); + accumulate_group(dim_idx, tArA, cute::Int<5>{}, tCrC5); + accumulate_group(dim_idx, tArA, cute::Int<6>{}, tCrC6); + accumulate_group(dim_idx, tArA, cute::Int<7>{}, tCrC7); + } + + cute::copy(tCrC0, tCgC0); + cute::copy(tCrC1, tCgC1); + cute::copy(tCrC2, tCgC2); + cute::copy(tCrC3, tCgC3); + cute::copy(tCrC4, tCgC4); + cute::copy(tCrC5, tCgC5); + cute::copy(tCrC6, tCgC6); + cute::copy(tCrC7, tCgC7); + } } else { #pragma unroll for (uint32_t token_group = 0; token_group < kTokenGroups; ++token_group) { diff --git a/tests/test_sm120_reference.py b/tests/test_sm120_reference.py index b1131b33f4..7fa0fa8b9d 100644 --- a/tests/test_sm120_reference.py +++ b/tests/test_sm120_reference.py @@ -201,7 +201,10 @@ def test_sm120_fp8_paged_mqa_logits_v4_shape_reference_path() -> None: @_test_filter(lambda: get_arch_major() >= 12) -@pytest.mark.parametrize("token_groups, cache_q", ((1, False), (2, False), (4, False), (8, False), (4, True))) +@pytest.mark.parametrize( + "token_groups, cache_q", + ((1, False), (2, False), (4, False), (8, False), (2, True), (4, True), (8, True)), +) def test_sm120_fp8_paged_mqa_logits_tiled_path(monkeypatch, token_groups: int, cache_q: bool) -> None: torch.manual_seed(5) monkeypatch.setenv("DG_SM120_PAGED_MQA_TILED", "1") From 1e94dc25dcdce0d10e564ef67bdccd6e3c90f836 Mon Sep 17 00:00:00 2001 From: jasl Date: Sun, 26 Apr 2026 08:25:44 +0800 Subject: [PATCH 13/14] Clean up SM120 paged MQA dispatch naming Rename the SM120 paged MQA host dispatcher so it no longer advertises itself as reference-only now that it can select the tiled kernel via DG_SM120_PAGED_MQA_TILED. Also tighten a temporary HC fallback comment and avoid binding the generated runtime args to a temporary reference. --- csrc/apis/attention.hpp | 6 ++-- .../impls/sm120_tf32_hc_prenorm_gemm.hpp | 4 +-- .../impls/smxx_fp8_fp4_paged_mqa_logits.hpp | 28 +++++++++---------- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 21285ab3ad..fc82df3e11 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -369,9 +369,9 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tupleget_prop()->sharedMemPerBlockOptin; - const SM120BF16HCPrenormGemmRuntime::Args& args = { + const SM120BF16HCPrenormGemmRuntime::Args args = { .m = m, .n = n, .k = k, .block_m = block_m, .block_n = block_n, .block_k = block_k, .num_splits = num_splits, diff --git a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp index 91a3fbaeaa..7ffae2c1a5 100644 --- a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp @@ -292,20 +292,20 @@ static void __instantiate_kernel() {{ } }; -static void sm120_fp8_paged_mqa_logits_reference(const torch::Tensor& q, - const torch::Tensor& kv_cache, - const torch::Tensor& kv_cache_scales, - const torch::Tensor& weights, - const torch::Tensor& context_lens, - const torch::Tensor& logits, - const torch::Tensor& block_table, - const at::ScalarType& logits_dtype, - const int& batch_size, const int& next_n, - const int& num_heads, const int& head_dim, - const int& block_kv, - const bool& is_context_lens_2d, - const int& logits_stride, - const int& block_table_stride) { +static void sm120_fp8_paged_mqa_logits(const torch::Tensor& q, + const torch::Tensor& kv_cache, + const torch::Tensor& kv_cache_scales, + const torch::Tensor& weights, + const torch::Tensor& context_lens, + const torch::Tensor& logits, + const torch::Tensor& block_table, + const at::ScalarType& logits_dtype, + const int& batch_size, const int& next_n, + const int& num_heads, const int& head_dim, + const int& block_kv, + const bool& is_context_lens_2d, + const int& logits_stride, + const int& block_table_stride) { constexpr int tokens_per_block = 128; DG_HOST_ASSERT(num_heads == 32 or num_heads == 64); DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128); From 7a7a41a1bac7dacabe74057e7600e59f98f85bce Mon Sep 17 00:00:00 2001 From: jasl Date: Sun, 26 Apr 2026 08:33:31 +0800 Subject: [PATCH 14/14] Rename SM120 scalar kernel paths Use scalar naming for the SM120 compatibility kernels that are not pure test references: FP8 einsum and FP8 paged MQA logits now expose scalar runtime/kernel names. Keep reference naming only for the PyTorch oracle in tests, and rename the SM120 regression file to test_sm120_kernels.py so the PR reads as kernel coverage rather than temporary reference scaffolding. --- csrc/apis/einsum.hpp | 2 +- csrc/jit_kernels/impls/sm120_fp8_einsum.hpp | 14 +++++++------- .../impls/sm120_tf32_hc_prenorm_gemm.hpp | 2 +- .../impls/smxx_fp8_fp4_paged_mqa_logits.hpp | 14 +++++++------- .../include/deep_gemm/impls/sm120_fp8_einsum.cuh | 2 +- .../deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh | 2 +- ...st_sm120_reference.py => test_sm120_kernels.py} | 10 +++++----- 7 files changed, 23 insertions(+), 23 deletions(-) rename tests/{test_sm120_reference.py => test_sm120_kernels.py} (97%) diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index 4ea1b185d3..4fbb65e398 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -190,7 +190,7 @@ static void fp8_einsum(const std::string& expr, if (arch_major == 12) { DG_HOST_ASSERT(not c.has_value()); DG_HOST_ASSERT(recipe == std::make_tuple(1, 128, 128)); - sm120_fp8_bhr_hdr_bhd_reference(a.first, a.second, b.first, b.second, d); + sm120_fp8_bhr_hdr_bhd_scalar(a.first, a.second, b.first, b.second, d); return; } diff --git a/csrc/jit_kernels/impls/sm120_fp8_einsum.hpp b/csrc/jit_kernels/impls/sm120_fp8_einsum.hpp index 94777410e2..967221a20d 100644 --- a/csrc/jit_kernels/impls/sm120_fp8_einsum.hpp +++ b/csrc/jit_kernels/impls/sm120_fp8_einsum.hpp @@ -7,7 +7,7 @@ namespace deep_gemm { -class SM120FP8BhrHdrBhdReferenceRuntime final: public LaunchRuntime { +class SM120FP8BhrHdrBhdScalarRuntime final: public LaunchRuntime { public: struct Args { int shape_b; @@ -49,7 +49,7 @@ class SM120FP8BhrHdrBhdReferenceRuntime final: public LaunchRuntime(&sm120_fp8_bhr_hdr_bhd_reference< + auto ptr = reinterpret_cast(&sm120_fp8_bhr_hdr_bhd_scalar< {}, {} >); }}; @@ -86,7 +86,7 @@ static void __instantiate_kernel() {{ } }; -static void sm120_fp8_bhr_hdr_bhd_reference(const torch::Tensor& a, +static void sm120_fp8_bhr_hdr_bhd_scalar(const torch::Tensor& a, const torch::Tensor& sfa, const torch::Tensor& b, const torch::Tensor& sfb, @@ -113,7 +113,7 @@ static void sm120_fp8_bhr_hdr_bhd_reference(const torch::Tensor& a, const int num_d_tiles = ceil_div(shape_d, output_tile_d); const int grid_dim = shape_b * shape_h * num_d_tiles; - const SM120FP8BhrHdrBhdReferenceRuntime::Args args = { + const SM120FP8BhrHdrBhdScalarRuntime::Args args = { .shape_b = shape_b, .shape_h = shape_h, .shape_d = shape_d, @@ -141,9 +141,9 @@ static void sm120_fp8_bhr_hdr_bhd_reference(const torch::Tensor& a, .output_dtype = d.scalar_type(), .launch_args = LaunchArgs(grid_dim, output_tile_d) }; - const auto code = SM120FP8BhrHdrBhdReferenceRuntime::generate(args); - const auto runtime = compiler->build("sm120_fp8_bhr_hdr_bhd_reference", code); - SM120FP8BhrHdrBhdReferenceRuntime::launch(runtime, args); + const auto code = SM120FP8BhrHdrBhdScalarRuntime::generate(args); + const auto runtime = compiler->build("sm120_fp8_bhr_hdr_bhd_scalar", code); + SM120FP8BhrHdrBhdScalarRuntime::launch(runtime, args); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp index f3f6380a75..4a2a1bc0a5 100644 --- a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp @@ -69,7 +69,7 @@ static void sm120_tf32_hc_prenorm_gemm(const torch::Tensor& a, const int& m, const int& n, const int& k, const int& num_splits) { - // SM120 fallback uses the same small-N tile shape as the SM90 path. + // SM120 HC currently uses the same small-N tile shape as the SM90 path. constexpr int block_m = 64; constexpr int block_k = 64; constexpr int num_math_threads = 128; diff --git a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp index 7ffae2c1a5..98d28555ab 100644 --- a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp @@ -166,7 +166,7 @@ static void __instantiate_kernel() {{ } }; -class SM120FP8PagedMQALogitsReferenceRuntime final: public LaunchRuntime { +class SM120FP8PagedMQALogitsScalarRuntime final: public LaunchRuntime { public: struct Args { int batch_size; @@ -208,7 +208,7 @@ class SM120FP8PagedMQALogitsReferenceRuntime final: public LaunchRuntime(&sm120_fp8_paged_mqa_logits_reference< + auto ptr = reinterpret_cast(&sm120_fp8_paged_mqa_logits_scalar< {}, {}, {}, {}, {}, {} @@ -245,7 +245,7 @@ static void __instantiate_kernel() {{ class SM120FP8PagedMQALogitsTiledRuntime final: public LaunchRuntime { public: - using Args = SM120FP8PagedMQALogitsReferenceRuntime::Args; + using Args = SM120FP8PagedMQALogitsScalarRuntime::Args; static std::string generate_impl(const Args& args) { return fmt::format(R"( @@ -311,7 +311,7 @@ static void sm120_fp8_paged_mqa_logits(const torch::Tensor& q, DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128); DG_HOST_ASSERT(block_kv == 32 or block_kv == 64); - const SM120FP8PagedMQALogitsReferenceRuntime::Args args = { + const SM120FP8PagedMQALogitsScalarRuntime::Args args = { .batch_size = batch_size, .next_n = next_n, .num_heads = num_heads, @@ -366,9 +366,9 @@ static void sm120_fp8_paged_mqa_logits(const torch::Tensor& q, return; } - const auto code = SM120FP8PagedMQALogitsReferenceRuntime::generate(args); - const auto runtime = compiler->build("sm120_fp8_paged_mqa_logits_reference", code); - SM120FP8PagedMQALogitsReferenceRuntime::launch(runtime, args); + const auto code = SM120FP8PagedMQALogitsScalarRuntime::generate(args); + const auto runtime = compiler->build("sm120_fp8_paged_mqa_logits_scalar", code); + SM120FP8PagedMQALogitsScalarRuntime::launch(runtime, args); } static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q, diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh index 2e09a876c4..e24008bbd6 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_einsum.cuh @@ -73,7 +73,7 @@ CUTLASS_DEVICE float sm120_fp8_einsum_dot_fp8x4( } template -CUTLASS_GLOBAL void sm120_fp8_bhr_hdr_bhd_reference( +CUTLASS_GLOBAL void sm120_fp8_bhr_hdr_bhd_scalar( const uint32_t shape_b, const uint32_t shape_h, const uint32_t shape_d, diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh index b3544517e8..56c6becfff 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_paged_mqa_logits.cuh @@ -35,7 +35,7 @@ template CUTLASS_GLOBAL __launch_bounds__(kTokensPerBlock, 1) -void sm120_fp8_paged_mqa_logits_reference( +void sm120_fp8_paged_mqa_logits_scalar( const uint32_t batch_size, const uint32_t next_n, const uint32_t logits_stride, diff --git a/tests/test_sm120_reference.py b/tests/test_sm120_kernels.py similarity index 97% rename from tests/test_sm120_reference.py rename to tests/test_sm120_kernels.py index 7fa0fa8b9d..f7f95d0a97 100644 --- a/tests/test_sm120_reference.py +++ b/tests/test_sm120_kernels.py @@ -100,7 +100,7 @@ def test_sm120_hc_prenorm_gemm_larger_tiles() -> None: @_test_filter(lambda: get_arch_major() >= 12) -def test_sm120_fp8_paged_mqa_logits_reference_path() -> None: +def test_sm120_fp8_paged_mqa_logits_scalar_path() -> None: torch.manual_seed(1) batch_size, next_n, num_heads, head_dim = 2, 2, 32, 32 @@ -153,7 +153,7 @@ def test_sm120_fp8_paged_mqa_logits_reference_path() -> None: @_test_filter(lambda: get_arch_major() >= 12) -def test_sm120_fp8_paged_mqa_logits_v4_shape_reference_path() -> None: +def test_sm120_fp8_paged_mqa_logits_v4_shape_scalar_path() -> None: torch.manual_seed(4) batch_size, next_n, num_heads, head_dim = 1, 1, 64, 128 @@ -256,7 +256,7 @@ def test_sm120_fp8_paged_mqa_logits_tiled_path(monkeypatch, token_groups: int, c @_test_filter(lambda: get_arch_major() >= 12) -def test_sm120_fp8_paged_mqa_logits_varlen_reference_path() -> None: +def test_sm120_fp8_paged_mqa_logits_varlen_scalar_path() -> None: torch.manual_seed(2) raw_batch_size, num_heads, head_dim = 3, 32, 32 @@ -331,7 +331,7 @@ def test_sm120_fp8_paged_mqa_logits_varlen_reference_path() -> None: @_test_filter(lambda: get_arch_major() >= 12) -def test_sm120_fp8_einsum_reference_path() -> None: +def test_sm120_fp8_einsum_scalar_path() -> None: torch.manual_seed(3) batch_size, num_heads, rank, output_dim = 3, 4, 256, 128 @@ -358,7 +358,7 @@ def test_sm120_fp8_einsum_reference_path() -> None: @_test_filter(lambda: get_arch_major() >= 12) -def test_sm120_fp8_einsum_v4_shape_reference_path() -> None: +def test_sm120_fp8_einsum_v4_shape_scalar_path() -> None: torch.manual_seed(5) batch_size, num_heads, rank, output_dim = 5, 64, 128, 128