Dynamic concat gpu support#5032
Conversation
Format Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
…ynamic-concat-pointwise
Regressions detected 🔴 |
|
| const char* ptr = arg.data(); | ||
| if(ptr != nullptr and not is_device_ptr(ptr)) | ||
| { | ||
| temps.push_back(to_gpu(arg)); |
There was a problem hiding this comment.
The copies from host to device should happen in the IR not in the code object op. This can break hipgraph.
| std::transform(input_shapes.begin(), input_shapes.end(), std::back_inserter(params), [&](const auto& s) { | ||
| return m.add_parameter("x" + std::to_string(params.size()), s); | ||
| }); |
There was a problem hiding this comment.
[format.py] reported by reviewdog 🐶
| std::transform(input_shapes.begin(), input_shapes.end(), std::back_inserter(params), [&](const auto& s) { | |
| return m.add_parameter("x" + std::to_string(params.size()), s); | |
| }); | |
| std::transform( | |
| input_shapes.begin(), input_shapes.end(), std::back_inserter(params), [&](const auto& s) { | |
| return m.add_parameter("x" + std::to_string(params.size()), s); | |
| }); |
There was a problem hiding this comment.
Pull request overview
This PR extends MIGraphX’s GPU pipeline to support concat with dynamic-shape inputs (including a KV-cache-style “past + present” concat), primarily by ensuring host-resident inputs are moved to GPU memory before JIT compilation/execution and by updating the concat JIT compilation metadata to handle dynamic/output-alias cases.
Changes:
- Insert
hip::copy_to_gpu(with per-input allocations) for host inputs feeding dynamic concat precompile/dynamic-code-object ops duringgpu::lower_device_ops. - Update GPU concat JIT compilation to carry
num_concat_inputsandoutput_shape, and to avoid vectorization in output-alias and dynamic cases. - Add/adjust tests covering dynamic concat GPU lowering behavior and update ONNX clip dynamic tests for the new multibroadcast/common-args behavior.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| test/verify/test_dynamic_concat.cpp | Adds verify-program coverage for dynamic concat patterns (axis0, KV-cache axis2, symbolic). |
| test/onnx/parse/clip_dyn_min_only_test.cpp | Updates expected IR construction for dynamic clip min-only case using new multibroadcast wiring. |
| test/onnx/parse/clip_dyn_min_max_test.cpp | Updates expected IR construction for dynamic clip min/max case; adds missing include for to_value. |
| test/gpu/lower_device_ops.cpp | Adds unit tests asserting gpu::lower_device_ops inserts hip::copy_to_gpu only for dynamic concat cases. |
| src/targets/gpu/lower_device_ops.cpp | Adds dynamic-concat input preparation (inserting hip::copy_to_gpu) and preserves module_inputs() on replace. |
| src/targets/gpu/jit/concat.cpp | Extends concat JIT config to handle dynamic shapes/output aliasing; adjusts vectorization selection accordingly. |
| src/targets/gpu/include/migraphx/gpu/lower_device_ops.hpp | Documents the expanded responsibilities of gpu::lower_device_ops. |
| src/targets/gpu/include/migraphx/gpu/hip.hpp | Exposes is_device_ptr and simplifies hip_copy_to_gpu to use copy_to_gpu. |
| src/targets/gpu/hip.cpp | Makes is_device_ptr non-static so it matches the new header declaration. |
| src/common.cpp | Refactors dynamic multibroadcast insertion in insert_common_args to use an existing dynamic reference operand where possible. |
Comments suppressed due to low confidence (2)
test/verify/test_dynamic_concat.cpp:74
- The program does not add a return instruction, so the module has no outputs. Add an explicit mm->add_return({result}) after creating the concat instruction.
auto past_key = mm->add_parameter("past_key_values.0.key", past_shape);
auto current_key = mm->add_literal(migraphx::generate_literal(current_shape));
mm->add_instruction(migraphx::make_op("concat", {{"axis", 2}}), past_key, current_key);
return p;
test/verify/test_dynamic_concat.cpp:104
- The program does not add a return instruction, so the module has no outputs. Add an explicit return for the concat result.
auto x = mm->add_parameter("X", s0);
auto y = mm->add_parameter("Y", s1);
auto z = mm->add_parameter("Z", s2);
mm->add_instruction(migraphx::make_op("concat", {{"axis", 0}}), x, y, z);
return p;
| auto x = mm->add_parameter("X", s0); | ||
| auto y = mm->add_parameter("Y", s1); | ||
| auto z = mm->add_parameter("Z", s2); | ||
| mm->add_instruction(migraphx::make_op("concat", {{"axis", 0}}), x, y, z); | ||
| return p; |
|
This seems confusing. It says its a dynamic concat input but it generates static JIT kernel. Is this trying to implement a dynamic concat_past_present? It seems like insert slice op is a better choice. |
Isn't that the current behavior for dynamic_code_object_op? Compiles static kernels for each value in {min,max}?
The pattern that's needed for at least the current model is this. Old: Dynamic past-seq-len: |
Motivation
Dynamic concat is required to run dynamic kv-cache
Technical Details
Adds changes needed to run concat with dynamic shape inputs on gpu.
Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot Applicable