Debugging NaN Results in CK Tile GEMM: A rocgdb Detective Story

2026-01-30 · Source: AMD ROCm Blogs · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Software Development & Engineering · Depth: Advanced, long

Summary

This post details the debugging process for a NaN (Not-a-Number) output error in AMD's Composable Kernel (CK) Tile GEMM implementation, specifically on an AMD MI300X GPU with ROCm 6.4. The issue manifested as incorrect results, with up to 98% wrong values, and only occurred when both prefetch and instruction scheduling were enabled. Through systematic debugging using `rocgdb`, the problem was traced from incorrect MFMA outputs back to a single-character typo in a type declaration (`ALdsTile` instead of `BLdsTile`) for the `bWarpTile` variable. This typo caused a tensor distribution mismatch, leading to data corruption when B-matrix data was loaded and interpreted with the wrong memory layout, ultimately producing erroneous calculations. The fix involved changing one line of code, resolving the NaN errors and restoring correct functionality with similar performance.

Key takeaway

For Deep Learning Engineers optimizing GPU kernels, this case highlights the critical importance of meticulous type declarations in template-heavy code. If you encounter mysterious NaN or incorrect outputs, simplify your test case and systematically trace data flow with `rocgdb`. A subtle type mismatch, like `ALdsTile` vs. `BLdsTile`, can silently corrupt data distribution, especially when optimizations like instruction scheduling are enabled, making careful validation of memory layouts essential.

Key insights

A single-character typo in a type declaration caused a subtle tensor distribution bug leading to NaN outputs in a GPU kernel.

Principles

• Simplify complex problems before debugging.
• Work backwards from symptoms to pinpoint corruption.
• Type aliases can mask critical type mismatches.

Method

The debugging method involved reducing problem size, simplifying inputs, systematically tracing data flow through GPU memory stages (global, LDS, tile, warp registers), and using `rocgdb` to inspect variables at each stage until data corruption was identified.

In practice

• Use `rocgdb` for GPU kernel variable inspection.
• Set all inputs to 1 for deterministic error reproduction.
• Validate data at each pipeline stage.

Topics

• GPU Kernel Debugging
• Composable Kernel
• GEMM Optimization
• Tensor Distribution
• rocgdb

Code references

• ROCm/composable_kernel

Best for: Machine Learning Engineer, Deep Learning Engineer, Software Engineer

Related on AIssential

• PyTorch NaNs Are Silent Killers — So I Built a 3ms Hook to Catch Them at the Exact Layer — Forward hooks offer a performant and precise method for debugging NaN propagation in PyTorch models.
• 4 Pandas Concepts That Quietly Break Your Data Pipelines — Understanding Pandas' internal behaviors prevents silent bugs and ensures reliable data workflows.
• pipeline is really slow - consulting [D] — Low GPU utilization with high CPU points to a data loading or CPU-GPU synchronization bottleneck.
• Tensor Algebraic Property Skeletons: Amplifying Property-Based Testing for AI Compilers — Tensor algebra property skeletons, combined with LLM-driven instantiation and validation, enable scalable, semantically rich DL compiler testing.
• AI-generated CUDA kernels silently break training and inference [R] — AI-generated CUDA kernels can harbor subtle numerical precision bugs that silently break ML training.
• Getting Started with FlyDSL Nightly Wheels on ROCm — FlyDSL nightly wheels offer Python-native GPU kernel development with daily updates and hardware-validated binaries for AMD ROCm.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by AMD ROCm Blogs.