Skip to content

hrx: Q6_K prefill can read the packed rows decode reads, one device copy per weight; fixes a bus error in GGML_HRX_Q6K_PACKED_PREFILL (engine #2137, opt-in) #310

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

Merged
bong-water-water-bong merged 1 commit into from
Nov 11, 2026
Merged

hrx: Q6_K prefill can read the packed rows decode reads, one device copy per weight; fixes a bus error in GGML_HRX_Q6K_PACKED_PREFILL (engine #2137, opt-in) #310

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
50 changes: 32 additions & 18 deletions ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -728,12 +728,25 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context,
return false;
}
// engine #2137 (dispatch-q4k-packed-tail.h): the WMMA body reading both weights in the packed row layout
const bool packed_tail = use_wmma && match.gate_format == CommonMulMatWeightFormat::Q4K &&
match.up_format == CommonMulMatWeightFormat::Q4K &&
q4k_packed_tail_weight_ok(*match.gate_weight, match.input_size, match.output_size) &&
q4k_packed_tail_weight_ok(*match.up_weight, match.input_size, match.output_size) &&
q4k_packed_tail_enabled();
const bool row_layout = pack_q4 || packed_tail;
// each weight on its own: Q4_K or Q6_K in its packed layout, anything else stored
const auto tail_format = [&](const Value & weight, CommonMulMatWeightFormat format) {
if (!use_wmma) {
return format;
}
if (format == CommonMulMatWeightFormat::Q4K &&
q4k_packed_tail_weight_ok(weight, match.input_size, match.output_size) && q4k_packed_tail_enabled()) {
return CommonMulMatWeightFormat::Q4KRow64;
}
if (format == CommonMulMatWeightFormat::Q6K &&
q6k_packed_rows_weight_ok(weight, match.input_size, match.output_size) && q6k_packed_rows_enabled()) {
return CommonMulMatWeightFormat::Q6KRow64;
}
return format;
};
const CommonMulMatWeightFormat gate_bound =
pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : tail_format(*match.gate_weight, match.gate_format);
const CommonMulMatWeightFormat up_bound =
pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : tail_format(*match.up_weight, match.up_format);
dispatch.kernel = make_kernel_specialization(publish_q8 ? kMulMatSwiGLUQ4Q8OutputKernel :
use_q8 ? kMulMatSwiGLUQ4Q8LowTokenDotKernel :
(use_direct_dot ? kMulMatSwiGLUF32F32LowTokenDotKernel : kMulMatSwiGLUF32F32WmmaKernel));
Expand All @@ -746,24 +759,25 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context,
common_to_config_value(match.output_size));
dispatch.kernel.compile_parameters.emplace(
"ggml.mul_mat_swiglu.gate_weight_format",
common_to_config_value(common_mul_mat_format_config_value(row_layout ? CommonMulMatWeightFormat::Q4KRow64 :
match.gate_format)));
common_to_config_value(common_mul_mat_format_config_value(gate_bound)));
dispatch.kernel.compile_parameters.emplace(
"ggml.mul_mat_swiglu.up_weight_format",
common_to_config_value(common_mul_mat_format_config_value(row_layout ? CommonMulMatWeightFormat::Q4KRow64 :
match.up_format)));
common_to_config_value(common_mul_mat_format_config_value(up_bound)));
dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.op",
common_to_config_value(static_cast<int64_t>(binary_kind_config_value(match.op))));
dispatch.bindings.push_back(activation);
if (row_layout) {
for (const Value * weight : { match.gate_weight, match.up_weight }) {
dispatch.bindings.push_back({ weight->id, 0, weight->byte_count, kQ4KPackedK256Row64Layout, weight->type,
match.input_size, match.output_size, weight->byte_count });
const auto bind_weight = [&](const Value & weight, CommonMulMatWeightFormat bound) {
const char * layout = bound == CommonMulMatWeightFormat::Q4KRow64 ? kQ4KPackedK256Row64Layout :
bound == CommonMulMatWeightFormat::Q6KRow64 ? kQ6KPackedK256Row64ScaleRowLayout : nullptr;
if (layout != nullptr) {
dispatch.bindings.push_back({ weight.id, 0, weight.byte_count, layout, weight.type, match.input_size,
match.output_size, weight.byte_count });
} else {
dispatch.bindings.push_back({ weight.id, 0, weight.byte_count });
}
} else {
dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count });
dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count });
}
};
bind_weight(*match.gate_weight, gate_bound);
bind_weight(*match.up_weight, up_bound);
dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count });

if (publish_q8) {
Expand Down
54 changes: 49 additions & 5 deletions ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,17 @@ static bool q4k_packed_tail_applies(const DispatchMatchContext & context,
!weight_has_native_layout_reader(context.graph, weight) && q4k_packed_tail_enabled();
}

static bool q6k_packed_tail_applies(const DispatchMatchContext & context,
const Value & weight,
CommonMulMatWeightFormat format,
int64_t token_count,
int64_t input_size,
int64_t output_size) {
return format == CommonMulMatWeightFormat::Q6K && token_count > 5 && token_count <= kMulMatPrefillMaxTokens &&
q6k_packed_rows_weight_ok(weight, input_size, output_size) &&
!weight_has_native_layout_reader(context.graph, weight) && q6k_packed_rows_enabled();
}

static bool has_normalized_rope_consumer(const Graph & graph, const Value & value) {
if (!graph.has_index()) {
return false;
Expand Down Expand Up @@ -676,9 +687,14 @@ static bool build_q6_k_prefill_wave32_dispatch(const CommonMulMatMatch & match,
int input_scale_log2,
KernelCatalogRef kernel,
size_t root_index,
DispatchMatch & dispatch_match) {
DispatchMatch & dispatch_match,
bool packed_rows = false) {
Dispatch dispatch;
dispatch.kernel = make_kernel_specialization(kernel);
if (packed_rows) {
// engine #2137: the kernel reads the packed row layout decode reads (dispatch-q4k-packed-tail.h)
dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.weight_layout", "2");
}
dispatch.kernel.integer_parameters.emplace("token_count", match.token_count);
dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.input_size",
common_to_config_value(match.input_size));
Expand All @@ -689,7 +705,13 @@ static bool build_q6_k_prefill_wave32_dispatch(const CommonMulMatMatch & match,
common_to_config_value(match.token_count));
common_set_f16_prefill_activation_scale(dispatch, input_scale_log2);
dispatch.bindings.push_back({ input, 0, input_bytes });
dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count });
if (packed_rows) {
dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, kQ6KPackedK256Row64ScaleRowLayout,
match.weight->type, match.input_size, match.output_size,
match.weight->byte_count });
} else {
dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count });
}
dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count });

dispatch_match.covered_nodes.push_back(root_index);
Expand Down Expand Up @@ -719,11 +741,23 @@ static bool match_q6_k_prefill_wave32_dispatch(const DispatchMatchContext & cont
if (!prepared) {
return false;
}
// engine #2137: with GGML_HRX_Q6K_PACKED_ROWS=1 the f16 kernel reads the packed rows decode reads. It takes
// precedence over the older GGML_HRX_Q6K_PACKED_PREFILL, whose kernel is bound to the stored layout.
// The from_packed kernel reads the stored rows on its wide tile (tokens a multiple of 256, outputs of 128) and
// 17,536-byte int8 tiles on every other tile, while this dispatch binds the stored layout. Off the wide tile it
// read past the weight: test-backend-ops MUL_MAT(q6_K, m=2048, n=128, k=4096, weight=resident) with
// GGML_HRX_Q6K_PACKED_PREFILL=1 ended in "HRX buffer download failed" and a bus error. So the switch selects that
// kernel only where it reads what is bound.
const bool from_packed = q6k_packed_prefill_enabled() && match.token_count % 256 == 0 &&
match.output_size % 128 == 0;
const bool packed_rows = q6k_packed_rows_weight_ok(*match.weight, match.input_size, match.output_size) &&
!weight_has_native_layout_reader(context.graph, *match.weight) && q6k_packed_rows_enabled();
return build_q6_k_prefill_wave32_dispatch(match, activation.value, activation.length, activation_scale_log2,
q6k_packed_prefill_enabled() ? kMulMatQ6KF16WmmaPrefillFromPackedKernel
: kMulMatQ6KF16WmmaPrefillWave32Kernel,
packed_rows ? kMulMatQ6KF16WmmaPrefillWave32Kernel :
from_packed ? kMulMatQ6KF16WmmaPrefillFromPackedKernel
: kMulMatQ6KF16WmmaPrefillWave32Kernel,
context.root_index,
dispatch_match);
dispatch_match, packed_rows);
}

static bool match_symmetric_i4_prefill_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) {
Expand Down Expand Up @@ -1171,6 +1205,11 @@ static bool build_mul_mat_dispatch(const DispatchMatchContext & context,
match.output_size)) {
pack_q4 = true;
}
if (!pack_q4 && !pack_q6 && !split_k &&
q6k_packed_tail_applies(context, *match.weight, match.weight_format, match.token_count, match.input_size,
match.output_size)) {
pack_q6 = true;
}
const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 :
pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : match.weight_format;
dispatch.kernel.compile_parameters.emplace(
Expand Down Expand Up @@ -1295,6 +1334,11 @@ static bool match_mul_mat_postops_dispatch(const DispatchMatchContext & context,
match.output_size)) {
pack_q4 = true; // engine #2137
}
if (!pack_q4 && !pack_q6 &&
q6k_packed_tail_applies(context, *match.weight, match.weight_format, match.token_count, match.input_size,
match.output_size)) {
pack_q6 = true; // engine #2137
}
const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 :
pack_q6 ? CommonMulMatWeightFormat::Q6KRow64 : match.weight_format;
dispatch.kernel.compile_parameters.emplace(
Expand Down
28 changes: 25 additions & 3 deletions ggml/src/ggml-hrx/dispatch_registration/common/dispatch-q4k-packed-tail.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,8 @@ inline bool q4k_packed_tail_enabled() {
// tensor in an ordinary buffer can be rewritten between computes, and the executor refuses to read one in a packed
// layout ("its buffer is not flagged as weights"), which would fail a graph the stored layout runs (measured on
// test-backend-ops' q4_K large-activation cases at 22 tokens).
inline bool q4k_packed_tail_weight_ok(const Value & weight, int64_t input_size, int64_t output_size) {
if (weight.type != GGML_TYPE_Q4_K || input_size % 256 != 0 || output_size % 64 != 0 ||
weight.alias_source.value >= 0) {
inline bool packed_tail_weight_ok(const Value & weight, ggml_type type, int64_t input_size, int64_t output_size) {
if (weight.type != type || input_size % 256 != 0 || output_size % 64 != 0 || weight.alias_source.value >= 0) {
return false;
}
const ggml_tensor * tensor = weight.tensor;
Expand All @@ -64,4 +63,27 @@ inline bool q4k_packed_tail_weight_ok(const Value & weight, int64_t input_size,
return buffer != nullptr && ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS;
}

inline bool q4k_packed_tail_weight_ok(const Value & weight, int64_t input_size, int64_t output_size) {
return packed_tail_weight_ok(weight, GGML_TYPE_Q4_K, input_size, output_size);
}

// Q6_K (GGML_HRX_Q6K_PACKED_ROWS=1): prefill reads the packed row layout decode reads, at every batch size. The tiled
// kernels take it as weight format 46 and the 128-aligned f16 prefill kernel as weight_layout 2. Stored stays the
// default. (GGML_HRX_Q6K_PACKED_PREFILL is an older switch on the aligned kernel alone; it binds the stored layout and
// saves no copy.)
inline bool q6k_packed_rows_enabled() {
static const bool enabled = [] {
const char * env = std::getenv("GGML_HRX_Q6K_PACKED_ROWS");
const bool on = env != nullptr && env[0] == '1';
GGML_LOG_ERROR("%s: engine #2137 q6_k prefill reads the %s layout (GGML_HRX_Q6K_PACKED_ROWS=%s)\n", __func__,
on ? "PACKED" : "stored", on ? "1" : "unset");
return on;
}();
return enabled;
}

inline bool q6k_packed_rows_weight_ok(const Value & weight, int64_t input_size, int64_t output_size) {
return packed_tail_weight_ok(weight, GGML_TYPE_Q6_K, input_size, output_size);
}

} // namespace ggml::hrx
86 changes: 86 additions & 0 deletions ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11080,6 +11080,92 @@ func.def inline @ggml_q6k_f32_vector4(%weight: buffer, %weight_row_byte_base: of
func.return %values : vector<4xf32>
}

// The same four values from HRX's packed row layout of Q6_K (q6k-packed-k256-row64-scalerow, host-memory.cpp
// materialize_q6_k_packed_k256_row64_scalerow): rows in groups of 64, and per group and 256-value block one 13,440-byte
// tile = 64 d (f16) | 64 x 2 x 8 scales (i8, index (lane*2 + s%2)*8 + s/2 for stored scale s) | per half three planes
// ql[0..32), ql[32..64), qh, each 64 rows x 32 bytes. Byte-identical codes, so the result equals ggml_q6k_f32_vector4.
// engine #2137: lets the tiled kernels share the copy decode reads.
func.def inline @ggml_q6k_packed_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf32>) {
%c0 = index.constant 0 : index
%c2 = index.constant 2 : index
%c3 = index.constant 3 : index
%c4 = index.constant 4 : index
%c8 = index.constant 8 : index
%c64 = index.constant 64 : index
%c256 = index.constant 256 : index
%tile_bytes = index.constant 13440 : offset
%d_bytes = index.constant 2 : offset
%scale_plane_offset = index.constant 128 : offset
%raw_plane_offset = index.constant 1152 : offset
%byte = index.constant 1 : offset
%field_row_bytes = index.constant 32 : offset
%c4_i32v = vector.constant 4 : vector<1xi32>
%nibble_mask = vector.constant 252645135 : vector<1xi32>
%high_mask = vector.constant 50529027 : vector<1xi32>
%c32_f32v = vector.constant 32.0 : vector<4xf32>
%bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index
%bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index
%blocks = index.div %input_size, %c256 : index
%row_group = index.div %row, %c64 : index
%lane = index.rem %row, %c64 : index
%tile_index = index.madd %row_group, %blocks, %q6_block : index
%tile_base = index.scale %tile_index, %tile_bytes : index, offset -> offset
%group_in_half = index.rem %bounded_group, %c4 : index
%half = index.div %bounded_group, %c4 : index
%ql_side = index.rem %group_in_half, %c2 : index
%nibble = index.div %group_in_half, %c2 : index
%nibble_shift_index = index.mul %nibble, %c4 : index
%nibble_shift_i32 = index.cast %nibble_shift_index : index to i32
%nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32>
%qh_shift_index = index.mul %group_in_half, %c2 : index
%qh_shift_i32 = index.cast %qh_shift_index : index to i32
%qh_shift = vector.splat %qh_shift_i32 : vector<1xi32>
// d
%d_add = index.scale %lane, %d_bytes : index, offset -> offset
%d_byte_base = index.add %tile_base, %d_add : offset
%d_view = buffer.view %weight[%d_byte_base] : buffer -> view<1xf16>
// scale: stored index s = group*2 + packet/4, kept at (lane*2 + s%2)*8 + s/2
%scale_half = index.div %bounded_packet, %c4 : index
%scale_lane = index.madd %lane, %c2, %scale_half : index
%scale_index = index.madd %scale_lane, %c8, %bounded_group : index
%scale_add = index.scale %scale_index, %byte : index, offset -> offset
%scale_plane_base = index.add %tile_base, %scale_plane_offset : offset
%scale_byte_base = index.add %scale_plane_base, %scale_add : offset
%scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<1xi8>
// codes: plane (half*3 + field), field 0/1 = the two ql halves, 2 = qh; then the row, then the packet's 4 bytes
%raw_base = index.add %tile_base, %raw_plane_offset : offset
%half_planes = index.mul %half, %c3 : index
%ql_plane = index.add %half_planes, %ql_side : index
%qh_plane = index.add %half_planes, %c2 : index
%ql_row = index.madd %ql_plane, %c64, %lane : index
%qh_row = index.madd %qh_plane, %c64, %lane : index
%ql_row_add = index.scale %ql_row, %field_row_bytes : index, offset -> offset
%qh_row_add = index.scale %qh_row, %field_row_bytes : index, offset -> offset
%ql_row_base = index.add %raw_base, %ql_row_add : offset
%qh_row_base = index.add %raw_base, %qh_row_add : offset
%ql_view = buffer.view %weight[%ql_row_base] : buffer -> view<8xi32>
%qh_view = buffer.view %weight[%qh_row_base] : buffer -> view<8xi32>
%ql_word = vector.load %ql_view[%bounded_packet] : view<8xi32> -> vector<1xi32>
%qh_word = vector.load %qh_view[%bounded_packet] : view<8xi32> -> vector<1xi32>
%ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32>
%ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32>
%qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32>
%qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32>
%qh = vector.shli %qh_low, %c4_i32v : vector<1xi32>
%code = vector.ori %ql, %qh : vector<1xi32>
%code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8>
%code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32>
%centered = vector.subf %code_f32, %c32_f32v : vector<4xf32>
%scale_i8 = view.load %scale_view[%c0] : view<1xi8> -> i8
%d_f16 = view.load %d_view[%c0] : view<1xf16> -> f16
%scale = scalar.sitofp %scale_i8 : i8 to f32
%d = scalar.extf %d_f16 : f16 to f32
%combined_scale = scalar.mulf %scale, %d : f32
%combined_scale_vector = vector.splat %combined_scale : vector<4xf32>
%values = vector.mulf %centered, %combined_scale_vector : vector<4xf32>
func.return %values : vector<4xf32>
}

func.def inline @ggml_q8_0_f32_vector4(%weight: buffer, %row_byte_base: offset, %q8_block: index, %packet: index) -> (vector<4xf32>) {
%c0 = index.constant 0 : index
%c4 = index.constant 4 : index
Expand Down
Loading