Skip to content

Lower com.ryzenai MatMulNBits, GQO, and SSMLP - #5409

Open
cliffcai wants to merge 1 commit into
ROCm:developfrom
cliffcai:swdev-604598-ryzenai-matmulnbits
Open

cliffcai wants to merge 1 commit into
ROCm:developfrom
cliffcai:swdev-604598-ryzenai-matmulnbits

Conversation

@cliffcai

Copy link
Copy Markdown

Summary

  • Lower com.ryzenai::MatMulNBits when the hybrid schema adds optional bias and an NPU-only packed weight after the microsoft inputs. Zero-points are applied whenever they are present, not only when the node has exactly four inputs.
  • Lower fused com.ryzenai::GQO (GroupQueryAttention + o_proj MatMulNBits) and com.ryzenai::SSMLP (SkipSimplifiedLayerNorm + SwiGLU MatMulNBits) so hybrid model_jit parses and GPU-compiles.
  • Add matmulnbits_ryzenai_bias_test parse and reference verification.

Test plan

  • migraphx-driver read --onnx model_jit.onnx
  • migraphx-driver compile --gpu --onnx model_jit.onnx (gfx1152, ROCm 10.1)
  • test_onnx_parse / test_onnx_verify matmulnbits_* including matmulnbits_ryzenai_bias_test

SWDEV-604598

Hybrid graphs append bias and an NPU packed weight after the microsoft MatMulNBits inputs, and fuse attention and MLP as GQO and SSMLP. Apply zero-points whenever they are present so the scoped model compiles on GPU.
@cliffcai
cliffcai requested a review from causten as a code owner October 19, 2026 08:27
}

// Lower A @ dequant(B) [+ bias].
inline instruction_ref add_matmul_nbits(onnx_parser::node_info& info,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should really be an op builder.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Input-layout errors, dropped attention inputs, FP16 normalization clipping, and missing fused-operator regression tests block approval.

6 open issues
What changed in this PR

Extends MIGraphX ONNX parsing to lower Ryzen AI hybrid-model operators for GPU compilation.

Changes:

  • Shares MatMulNBits lowering and adds optional bias handling.
  • Adds fused GQO and SSMLP parsers.
  • Adds a Ryzen AI MatMulNBits fixture with parse and reference tests.
File Description
test/​onnx/​verify/​matmulnbits_tests.cpp Verifies biased MatMulNBits results.
test/​onnx/​parse/​matmulnbits_tests.cpp Checks the expected lowered graph.
test/​onnx/​matmulnbits_ryzenai_bias_test.onnx Adds the regression fixture.
test/​onnx/​gen_onnx.py Generates the Ryzen AI fixture.
src/​onnx/​parse_ssmlp.cpp Lowers fused normalization and SwiGLU.
src/​onnx/​parse_matmulnbits.cpp Delegates lowering and handles optional inputs.
src/​onnx/​parse_group_query_attention.cpp Adds fused attention/output projection parsing.
src/​onnx/​include/​migraphx/​onnx/​matmulnbits.hpp Provides shared quantized matrix-multiplication lowering.

Review effort: Balanced

Comment on lines +348 to +349
// First 9 inputs match microsoft GroupQueryAttention layout.
std::vector<instruction_ref> gqa_args(args.begin(), args.begin() + 9);
const int block_size = get_i("o_proj_block_size");

// o_proj weights begin at input 12 after three optional empty slots (9-11).
const std::size_t w = args.size() >= 17 ? 12 : args.size() - 5;
Comment on lines +58 to +60
std::optional<instruction_ref> bias;
if(args.size() > 4 and not args[4]->is_undefined())
bias = args[4];
Comment thread src/onnx/parse_ssmlp.cpp
Comment on lines +53 to +58
rms = info.add_instruction(make_op("convert", {{"target_type", x_dtype}}), rms);
epsilon =
(x_dtype == migraphx::shape::half_type and std::abs(epsilon) < 1e-7) ? 1e-7 : epsilon;
auto eps = info.add_literal(migraphx::literal{migraphx::shape{x_dtype}, {epsilon}});
rms = info.add_common_op("add", rms, eps);
auto rrms = info.add_instruction(make_op("rsqrt"), rms);
// First 9 inputs match microsoft GroupQueryAttention layout.
std::vector<instruction_ref> gqa_args(args.begin(), args.begin() + 9);
parse_group_query_attention gqa;
auto gqa_outs = gqa.parse(opd, parser, info, gqa_args);
Comment thread src/onnx/parse_ssmlp.cpp
Comment on lines +135 to +136
auto [out, skip_sum] = skip_simplified_ln(info, down, residual, args[12], epsilon);
return {out, skip_sum};

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants