Port prefill.cuh to HIP with CDNA3 MMA intrinscis support

[diptorupd]

WIP progress port of the single_prefill kernel to HIP with CDNA3 MMA instrinsics.

[Copilot]

Pull Request Overview

This PR introduces a HIP port of the single_prefill kernel with CDNA3 MMA intrinsics, specifically focused on adding AMD GPU support to the FlashInfer attention mechanism. The implementation includes comprehensive testing infrastructure for validating GPU kernels against CPU reference implementations.

Key changes include:

• HIP-specific single prefill kernel implementation with CDNA3 MMA intrinsics
• Comprehensive test infrastructure for various attention mechanism components
• CPU reference implementations for validation against GPU kernels

Reviewed Changes

Copilot reviewed 21 out of 23 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
test_prefill_sfrag.sh	Script for automated testing across multiple thread/warp configurations
sfrag_tester_script.py	Python utility for parsing and visualizing s_frag debug output from GPU kernels
libflashinfer/utils/utils_hip.h	HIP utility functions with deterministic random number generation and new data generation patterns
libflashinfer/utils/flashinfer_prefill_ops.hip.h	Main HIP interface for single prefill operations with CDNA3 support
libflashinfer/utils/cpu_reference_hip.h	CPU reference implementations for validating GPU attention kernels
libflashinfer/tests/hip/test_single_prefill.cpp	Primary test suite for single prefill kernel correctness
libflashinfer/tests/hip/test_*.cpp	Various component-specific test suites for memory patterns, MMA operations, and kernel validation
libflashinfer/include/flashinfer/attention/generic/default_prefill_params.cuh	Parameter structures for prefill operations
libflashinfer/include/flashinfer/attention/generic/dispatch.cuh	Compile-time dispatch macros for different kernel configurations
examples/cpp/standalone_single_prefill.cu	Standalone example demonstrating single prefill usage
compile_test.sh	Build script for HIP compilation with appropriate flags

Comments suppressed due to low confidence (1)

sfrag_tester_script.py:1

• This line contains C++ syntax in a Python file. It should be return att for Python.

#!/usr/bin/env python3

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Copilot]

The assertion condition is incorrect. It should check if the vector size is greater than 0 for the Linear predicate case. The current condition vec.size() <= 0 will always trigger an assertion failure for non-empty vectors, preventing the linear data generation from working.

Suggested change

	assert(vec.size() <= 0);
	assert(vec.size() > 0);

[yiakwy-xpu-ml-framework-team]

@diptorupd Good job ! Do you have the plan to upstream the codes to Flashinfer ?

[diptorupd]

@diptorupd Good job ! Do you have the plan to upstream the codes to Flashinfer ?

Please see: flashinfer-ai#1678

[Copilot]

Pull Request Overview

Copilot reviewed 17 out of 19 changed files in this pull request and generated 3 comments.

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Copilot]

Hard-coded seed values (1234) in random number generators make debugging easier but should be configurable for production use. Consider adding a parameter or compile-time flag to control deterministic vs. random seeding.

[Copilot]

Large debug blocks with hard-coded dimensions (128, 64) reduce code maintainability. Consider extracting debug functionality into separate functions or using configurable debug levels with template parameters.

Suggested change

	std::cout << "DEBUG: Original Q (CPU): " << '\n';
	for (auto i = 0ul; i < 128; ++i) {
	for (int j = 0; j < 64; ++j) {
	std::cout << (float)q[info.get_q_elem_offset(i, 0, j)] << " ";
	}
	std::cout << std::endl;
	}
	std::cout << std::endl;
	std::cout << "DEBUG: Original K (CPU): " << '\n';
	for (auto i = 0ul; i < 128; ++i) {
	for (int j = 0ul; j < 64; ++j) {
	std::cout << (float)k[info.get_kv_elem_offset(i, 0, j)] << " ";
	}
	std::cout << std::endl;
	}
	std::cout << std::endl;
	// Debug print helper for Q and K
	auto debug_print_matrix = [](const auto& tensor, auto elem_offset_fn, size_t rows, size_t cols, const std::string& name) {
	std::cout << "DEBUG: " << name << " (CPU): " << '\n';
	for (size_t i = 0; i < rows; ++i) {
	for (size_t j = 0; j < cols; ++j) {
	std::cout << (float)tensor[elem_offset_fn(i, 0, j)] << " ";
	}
	std::cout << std::endl;
	}
	std::cout << std::endl;
	};
	// Print up to 128 rows and 64 columns, or the actual size if smaller
	size_t print_rows = std::min(qo_len, size_t(128));
	size_t print_cols = std::min(head_dim, size_t(64));
	debug_print_matrix(q, [&](size_t i, size_t h, size_t j){ return info.get_q_elem_offset(i, h, j); }, print_rows, print_cols, "Original Q");
	debug_print_matrix(k, [&](size_t i, size_t h, size_t j){ return info.get_kv_elem_offset(i, h, j); }, print_rows, print_cols, "Original K");

[Copilot]

Function lacks comprehensive docstring explaining the matrix population algorithm and the relationship between thread/warp IDs and matrix positions. This is critical for understanding the attention mechanism's data layout.

Suggested change

	Populate the matrix with values from a specific thread and warp.
	Args:
	matrix: The 128x128 numpy array to populate
	thread_id: The thread ID (0-63)
	warp_id: The warp ID (0-3)
	values: List of 64 float values from this thread
	Populates the 128x128 matrix with values produced by a specific thread and warp, following the data layout used in attention mechanisms.
	The mapping algorithm works as follows:
	- Each warp (warp_id: 0-3) is responsible for a contiguous block of 32 rows in the matrix:
	* warp 0: rows 0-31
	* warp 1: rows 32-63
	* warp 2: rows 64-95
	* warp 3: rows 96-127
	- Each thread (thread_id: 0-63) within a warp is mapped to a specific set of rows and columns:
	* thread_row_base = (thread_id // 16) * 4
	* row_base = warp_row_offset + thread_row_base
	* col_base = thread_id % 16
	- Each thread produces 64 values, split into two calls of 32 values each:
	* The first 32 values are placed in columns 0-63, starting at (row_base, col_base).
	* The second 32 values are placed in columns 64-127, starting at (row_base, col_base).
	- The process_call_values function handles the placement of each 32-value block into the matrix.
	This mapping ensures that each thread/warp combination fills a unique region of the matrix, matching the parallel computation pattern of attention mechanisms on GPUs.
	Args:
	matrix: The 128x128 numpy array to populate.
	thread_id: The thread ID (0-63).
	warp_id: The warp ID (0-3).
	values: List of 64 float values produced by this thread.

[diptorupd]

Superseded by #31