Skip to content

ck_tile: fix 2:4-sparse SWMMAC correctness on gfx1201/RDNA4 (3 bugs) + fail→pass repro #3759

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
The-Monk wants to merge 2 commits into
base: develop
Choose a base branch
from
Open

ck_tile: fix 2:4-sparse SWMMAC correctness on gfx1201/RDNA4 (3 bugs) + fail→pass repro #3759

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cd34d2f
ck_tile: fix 2:4-sparse SWMMAC correctness on gfx1201/RDNA4 (3 bugs) …
The-Monk Jul 24, 2026
ac24ac2
Merge branch 'develop' into gfx1201-sparse-swmmac-fixes
The-Monk Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 14 additions & 3 deletions include/ck_tile/core/arch/mma/sparse/sparse_mma_pipeline.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -203,7 +203,18 @@ struct SparseMmaPipeline : public MmaPipelineBase<SparseMmaPipeline<ADataType_,
static constexpr index_t InternalAVecSize = vector_traits<FragAVecT>::vector_size;
static constexpr index_t ExternalAVecSize = InternalAVecSize * MmaOp::kCompressionRatio;
static constexpr index_t TotalUncompressedElems = FragsM * FragsK * ExternalAVecSize;
static constexpr index_t TotalCompressedElems = FragsM * FragsK * InternalAVecSize;
// Stage-17c fix (local, not upstreamed): the sparsity idx metadata is one
// 2-bit field per LOGICAL compressed element, not per physical storage
// element -- these differ for packed sub-byte A types (pk_int4_t,
// APackedSize=2). Confirmed against the RDNA4 ISA (Sec 7.12.3 Structured
// Sparse Matrices): "Wave32: each lane has 8 index values per lane" for a
// K=32 iu4 tile, where InternalAVecSize (physical pk_int4_t BYTE count
// per lane) is only 4 -- the ISA's own count is exactly
// InternalAVecSize * APackedSize, confirming idx tracks NIBBLES, not
// BYTES, for packed A. For APackedSize==1 (int8/fp8/fp16/...) this is a
// no-op (matches the original formula exactly).
static constexpr index_t TotalCompressedElems =
FragsM * FragsK * InternalAVecSize * MmaOp::APackedSize;

// Variable-length idx type for the whole wave-tile (spans multiple int32_t words if needed)
static constexpr index_t IdxNumWords = sparse::detail::idx_words_needed<TotalCompressedElems>;
Expand Down Expand Up @@ -257,7 +268,7 @@ struct SparseMmaPipeline : public MmaPipelineBase<SparseMmaPipeline<ADataType_,
a_frags[bm][bk],
b_buf.at(bn * FragsK + bk),
c_buf.at(bm * FragsN + bn),
sparse::detail::extract_fragment_idx<InternalAVecSize, FragsK>(
sparse::detail::extract_fragment_idx<InternalAVecSize * MmaOp::APackedSize, FragsK>(
idx, bm, bk));
}
}
Expand All @@ -275,7 +286,7 @@ struct SparseMmaPipeline : public MmaPipelineBase<SparseMmaPipeline<ADataType_,
a_frags[bm][bk],
b_buf.at(bn * FragsK + bk),
c_buf.at(bm * FragsN + bn),
sparse::detail::extract_fragment_idx<InternalAVecSize, FragsK>(
sparse::detail::extract_fragment_idx<InternalAVecSize * MmaOp::APackedSize, FragsK>(
idx, bm, bk));
}
}
Expand Down
236 changes: 201 additions & 35 deletions include/ck_tile/core/arch/mma/sparse/sparse_transforms.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#include "ck_tile/core/config.hpp"
#include "ck_tile/core/numeric/ext_vector_base.hpp"
#include "ck_tile/core/numeric/integer.hpp"
#include "ck_tile/core/numeric/pk_int4.hpp"
#include "ck_tile/core/numeric/vector_type.hpp"
#include "ck_tile/core/utility/bit_cast.hpp"
#include "ck_tile/core/utility/functional.hpp"
Expand Down Expand Up @@ -55,46 +56,202 @@ struct SparseIdxPack
* If fewer than CompressedSize non-zeros are found, remaining fields
* default to 2 (see below).
*/
template <typename ADataType, index_t CompressedSize, typename AVec>
// Stage-17c fix (local, not upstreamed): packed-sub-byte awareness.
//
// CompressedSize here is measured in PHYSICAL storage elements of ADataType
// (e.g. whole pk_int4_t BYTES, each holding 2 packed 4-bit values). The
// original code indexed a_vec[] and ran its "group of 4 / keep <=2" scan at
// that PHYSICAL granularity -- correct for PackedSize==1 (int8/fp8/fp16/...),
// but WRONG for PackedSize>1 (pk_int4_t): structured 2:4 sparsity is defined
// over LOGICAL elements (4-bit values), not physical storage words. Verified
// against the RDNA4 ISA (Sec 7.12.3, "Wave32: each lane has 8 index values
// per lane" for a K=32 iu4 tile whose physical AVecType is only 4 BYTES --
// i.e. one idx field per LOGICAL nibble, not per physical byte) and via a
// device-side dump: real per-nibble 2:4 data makes EVERY physical byte in a
// group of 4 bytes "nonzero" (since a byte is nonzero whenever either of its
// two nibbles survived), systematically exceeding the old code's "<=2
// nonzero per group of 4" assumption and writing out of bounds into
// nonzero_elems[2]/[3] (undefined behavior -- observed as register aliasing
// that silently corrupted the first two survivors on real hardware).
// Stage-17c fix (local, not upstreamed): LogicalADataType defaults to
// ADataType (preserves old behavior for every existing/non-packed caller),
// but callers that know the TRUE logical element type of a packed vector
// (e.g. pk_int4_t, whose physical ext_vector storage is represented as plain
// int8_t -- see vector_traits<ext_vector_t<pk_int4_t,N>>::scalar_type in
// vector_type.hpp, which deliberately erases pk_int4_t to int8_t because
// clang's ext_vector_type attribute needs a true scalar) can pass it
// explicitly so PackedSize is read from the REAL logical type, not the
// erased physical one. Without this, `ADataType` here is ALWAYS int8_t for
// pk_int4_t vectors (confirmed: the packedness tag is lost one layer up, in
// SparseCompressTransform::execExtVec, well before reaching this function --
// an earlier version of this fix that only checked
// numeric_traits<ADataType>::PackedSize here had NO EFFECT for exactly this
// reason, verified by a before/after diff showing byte-identical output).
template <typename ADataType, index_t CompressedSize, typename AVec, typename LogicalADataType = ADataType>
static CK_TILE_DEVICE auto compress_a_impl(AVec& a_vec)
{
static constexpr index_t NumIdxWords = idx_words_needed<CompressedSize>;
// idx holds one 2-bit index per output element (total CompressedSize entries),
// packed across NumIdxWords int32_t words.
// It is initialized to the pattern 0b10 for every field. This matches
// what the hardware expects when there are fewer than two non-zero values
// in a 4-element group - the unused output is treated as coming from slot 2.
// The loop below will clear and set each field as real non-zeros are seen.
static constexpr index_t PackedSize = numeric_traits<LogicalADataType>::PackedSize;
// LOGICAL field count: equals CompressedSize when PackedSize==1 (fully
// backward compatible with the original formula), doubles it for
// pk_int4_t-style 2-packed types.
static constexpr index_t LogicalCompressedSize = CompressedSize * PackedSize;
static constexpr index_t NumIdxWords = idx_words_needed<LogicalCompressedSize>;

// idx holds one 2-bit index per LOGICAL output element (total
// LogicalCompressedSize entries), packed across NumIdxWords int32_t
// words. Initialized to the pattern 0b10 for every field -- see below.
SparseIdxPack<NumIdxWords> idx{};
static_for<0, CompressedSize, 1>{}([&](auto k) {
static_for<0, LogicalCompressedSize, 1>{}([&](auto k) {
constexpr uint32_t bit_pos = static_cast<uint32_t>(k) * 2u;
constexpr uint32_t word = bit_pos / 32u;
constexpr uint32_t shift = bit_pos % 32u;
idx.words[word] |= static_cast<int32_t>(2u << shift);
});

static_for<0, CompressedSize / 2, 1>{}([&](auto i) {
ADataType nonzero_elems[2] = {a_vec[i * 4 + 2], a_vec[i * 4 + 3]};
int32_t non_zero_pos = 0;

static_for<0, 4, 1>{}([&](auto j) {
if(static_cast<float>(a_vec[i * 4 + j]) != 0.0f)
{
nonzero_elems[non_zero_pos] = a_vec[i * 4 + j];
// clear the two-bit field for this output and insert j
const uint32_t field_idx =
static_cast<uint32_t>(i) * 2u + static_cast<uint32_t>(non_zero_pos);
const uint32_t bit_pos = field_idx * 2u;
const uint32_t word = bit_pos / 32u;
const uint32_t shift = bit_pos % 32u;
idx.words[word] &= ~static_cast<int32_t>(0b11u << shift);
idx.words[word] |= static_cast<int32_t>(static_cast<uint32_t>(j) << shift);
++non_zero_pos;
}
// Shared idx-field writer (used by both the scalar and packed paths).
auto set_idx_field = [&](uint32_t field_idx, uint32_t j_value) {
const uint32_t bit_pos = field_idx * 2u;
const uint32_t word = bit_pos / 32u;
const uint32_t shift = bit_pos % 32u;
idx.words[word] &= ~static_cast<int32_t>(0b11u << shift);
idx.words[word] |= static_cast<int32_t>(j_value << shift);
};

if constexpr(PackedSize == 1)
{
// Original scalar path, byte-granular groups of 4 -- unchanged.
static_for<0, CompressedSize / 2, 1>{}([&](auto i) {
// Stage-17b fix (local, not upstreamed): a group with fewer than
// 2 real non-zeros leaves one (or both) output slots unassigned
// by the scan below; its idx stays at the default 2 (see
// comment above). For that default idx to be safe REGARDLESS of
// what value the other slot claims for position 2, an
// unassigned slot's VALUE must be a true zero -- not copied
// from some fixed input position, because that position can
// itself be the group's real (only) survivor. The original
// code read defaults from {a[2], a[3]}: this is only safe when
// position 3 is *guaranteed* zero (true for the canonical
// "keep slots 0,2" synthetic pattern) -- for genuinely
// arbitrary-position 2:4 data (e.g. groups whose sole survivor
// sits at position 3, or, subtly, at position 2 -- both occur
// in real Quark-quantized weights, which are only 97.7%
// exactly-2:4), that default silently duplicates a real value
// onto a mismatched idx (measured: {0,0,0,44} reconstructed as
// {0,0,44,44}, or in the slot0-collides-with-slot1 direction, a
// legitimate lone survivor at position 2 got double-counted
// against B, producing small K-independent errors even on the
// "canonical" synthetic pattern whenever its random fill
// happened to leave slot 0 zero and slot 2 as the true lone
// survivor). An always-zero default is correct in every case:
// 0/1/2-nonzero groups all still reconstruct exactly, and any
// idx collision on an unassigned slot contributes nothing.
ADataType nonzero_elems[2] = {static_cast<ADataType>(0), static_cast<ADataType>(0)};
int32_t non_zero_pos = 0;

static_for<0, 4, 1>{}([&](auto j) {
if(static_cast<float>(a_vec[i * 4 + j]) != 0.0f)
{
nonzero_elems[non_zero_pos] = a_vec[i * 4 + j];
set_idx_field(static_cast<uint32_t>(i) * 2u + static_cast<uint32_t>(non_zero_pos),
static_cast<uint32_t>(j));
++non_zero_pos;
}
});
a_vec[i * 2] = nonzero_elems[0];
a_vec[i * 2 + 1] = nonzero_elems[1];
});
a_vec[i * 2] = nonzero_elems[0];
a_vec[i * 2 + 1] = nonzero_elems[1];
});
}
else
{
static_assert(PackedSize == 2 && std::is_same_v<LogicalADataType, pk_int4_t>,
"compress_a_impl packed path only implements pk_int4_t (2x4-bit) currently");
// Packed-nibble path (pk_int4_t): each iteration consumes 2
// PHYSICAL input bytes (= 4 LOGICAL nibbles, one real 2:4 group) and
// produces 1 PHYSICAL output byte (the <=2 survivor nibbles,
// packed). CompressedSize physical output bytes -> CompressedSize
// iterations (unlike the scalar path's CompressedSize/2, because
// here one physical byte carries 2 logical elements).
//
// Nibble ordering matches pk_int4_t's own CK_TILE_USE_PK4_LAYOUT_SHUFFLE
// convention (on by default, see pk_int4.hpp): "logical element 0"
// of a packed byte is the HIGH nibble, "element 1" is the LOW
// nibble -- so within a group of 4 logical positions (j=0..3) built
// from 2 consecutive physical bytes (byte0=a_vec[2i], byte1=a_vec[2i+1]):
// j=0 -> byte0 high, j=1 -> byte0 low, j=2 -> byte1 high, j=3 -> byte1 low.
// Raw byte-pointer view of a_vec: clang's native ext_vector_type
// operator[] does not support a WRITE with a single (non-scaled)
// compile-time-constant index in this context (only the scalar
// path's `a_vec[i*2]=...` pattern compiles) -- sidestep entirely by
// reinterpreting a_vec as a flat uint8_t* for both the packed reads
// and the packed write below (same technique already validated in
// this Stage-17c audit's own probe fill code).
uint8_t * a_bytes = reinterpret_cast<uint8_t *>(&a_vec);
static_for<0, CompressedSize, 1>{}([&](auto i) {
const uint8_t byte0 = a_bytes[2 * static_cast<uint32_t>(i) + 0];
const uint8_t byte1 = a_bytes[2 * static_cast<uint32_t>(i) + 1];
auto get_nibble = [](uint8_t byte, bool high) -> int8_t {
uint8_t nib = high ? ((byte >> 4) & 0x0Fu) : (byte & 0x0Fu);
int8_t val = static_cast<int8_t>(nib);
if(val & 0x08) val = static_cast<int8_t>(val | 0xF0); // sign-extend 4-bit
return val;
};
const int8_t nib[4] = {
get_nibble(byte0, true), get_nibble(byte0, false),
get_nibble(byte1, true), get_nibble(byte1, false)};

// Same true-zero-default fix as the scalar path (Stage-17b),
// applied at nibble granularity.
//
// Stage-17c CLOSURE fix: the swmmac IU4 hardware reads the two
// idx fields of a group SWAPPED relative to which compressed
// value they govern, AND applies a "XOR 1" transform to the
// field's raw 2-bit value to get the real reconstructed
// position. Empirically determined (device measurement, not
// derived from the ISA doc alone -- the general Sec 7.12.3
// pseudocode does not capture this IU4-specific detail) via a
// well-formed (always-exactly-2-survivor, no default-idx
// ambiguity) position-pair sweep on real hardware: for a group
// whose HIGH-nibble survivor (my slot 0, found first in scan
// order) sat at true position j0 and LOW-nibble survivor (slot
// 1, found second) sat at true position j1, the hardware
// reconstructs correctly ONLY when field(i*2+1) [normally
// "idx1"] is written as (j0 XOR 1) and field(i*2+0) [normally
// "idx0"] is written as (j1 XOR 1) -- i.e. slot0's position goes
// into the OTHER field, XORed with 1, and vice versa. Verified
// against all 6 position-pair combinations (C(4,2)) with exact
// arithmetic match (predicted C01 == measured C01 in every
// case) before being applied here. This ALSO fully explains an
// earlier single-survivor sweep's "always reconstructs at
// position 3" finding: the untouched default field (value 2,
// unused-slot placeholder) XORed with 1 gives 3, constant,
// regardless of the real survivor's true position -- exactly
// what was observed.
int8_t survivor[2] = {0, 0};
int32_t non_zero_pos = 0;
static_for<0, 4, 1>{}([&](auto j) {
if(nib[static_cast<uint32_t>(j)] != 0)
{
survivor[non_zero_pos] = nib[static_cast<uint32_t>(j)];
// slot 0 (HIGH nibble) -> field i*2+1; slot 1 (LOW nibble) -> field i*2+0.
const uint32_t target_field =
static_cast<uint32_t>(i) * 2u + (1u - static_cast<uint32_t>(non_zero_pos));
set_idx_field(target_field, static_cast<uint32_t>(j) ^ 1u);
++non_zero_pos;
}
});
// Re-pack: element0 (survivor[0]) -> HIGH nibble, element1
// (survivor[1]) -> LOW nibble, matching the same shuffle
// convention used to unpack above (and used by the register-map
// fill on the host/caller side).
const uint8_t out_byte = static_cast<uint8_t>(
((static_cast<uint8_t>(survivor[0]) & 0x0Fu) << 4) |
(static_cast<uint8_t>(survivor[1]) & 0x0Fu));
// Output bytes are always written at index <= input bytes read
// (i <= 2i), so writing through the same raw-byte view in
// increasing i order never clobbers not-yet-read input data.
a_bytes[static_cast<uint32_t>(i)] = out_byte;
});
}

return idx;
}
Expand Down Expand Up @@ -146,7 +303,14 @@ struct SparseCompressTransform
{
// This function takes A in uncompressed form as a big ext_vector, and returns a
// compressed ext_vector.
template <typename VecType>
// Stage-17c fix (local, not upstreamed): added LogicalADataType, defaulted
// to vector_traits<VecType>::scalar_type so every EXISTING call site
// (including direct execExtVec() callers that don't know about packing)
// behaves exactly as before. exec() below (the real Pipeline::exec() call
// path) passes the TRUE logical element type explicitly, which is the
// only way compress_a_impl can see PackedSize > 1 (see its own comment).
template <typename VecType,
typename LogicalADataType = typename vector_traits<remove_cvref_t<VecType>>::scalar_type>
CK_TILE_DEVICE static decltype(auto) execExtVec(VecType& v)
{
using VecTraits = vector_traits<remove_cvref_t<VecType>>;
Expand All @@ -155,12 +319,13 @@ struct SparseCompressTransform
static constexpr index_t CompressedSize = VecN / CompressionRatio;
using VecCompressed = ext_vector_t<ScalarT, CompressedSize>;
using IdxType =
sparse::detail::SparseIdxPack<sparse::detail::idx_words_needed<CompressedSize>>;
sparse::detail::SparseIdxPack<sparse::detail::idx_words_needed<
CompressedSize * numeric_traits<LogicalADataType>::PackedSize>>;

static_assert(VecN % CompressionRatio == 0, "VecN must be divisible by CompressionRatio");
static_assert(CompressedSize > 0, "CompressedSize must be > 0");

auto idx = sparse::detail::compress_a_impl<ScalarT, CompressedSize>(v);
auto idx = sparse::detail::compress_a_impl<ScalarT, CompressedSize, VecType, LogicalADataType>(v);

return std::tuple<VecCompressed&, IdxType>(*ck_tile::bit_cast<VecCompressed*>(&v), idx);
}
Expand All @@ -174,11 +339,12 @@ struct SparseCompressTransform
CK_TILE_DEVICE static decltype(auto) exec(ATensor& a_tensor)
{
// Properties of ATensor as a big ext vector.
using ADataType = typename ATensor::DataType;
using ADataType = typename ATensor::DataType; // TRUE logical type (e.g. pk_int4_t)
constexpr index_t VecN = ATensor::get_thread_buffer_size();
using VecType = ext_vector_t<ADataType, VecN>;

return execExtVec(a_tensor.get_thread_buffer().template get_as<VecType>().template at<0>());
return execExtVec<VecType, ADataType>(
a_tensor.get_thread_buffer().template get_as<VecType>().template at<0>());
}
};

Expand Down
Loading