Feature/a8w8 matmul dequant fusion#1920
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughIntroduces fused A8W8 INT8 matmul dequantization operators ( ChangesA8W8 matmul dequant feature
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant OpConversionRegistry
participant TileMatmul
participant TileDequant
participant TileCast
OpConversionRegistry->>TileMatmul: lower INT8 matmul/gemv into INT32 accumulator
TileMatmul->>TileCast: cast accumulator to FP32
TileCast->>TileDequant: apply act_scale row scaling
TileDequant->>TileDequant: apply weight_scale column scaling
TileDequant->>TileCast: cast result to out_dtype
Possibly related issues
Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ebd3abfe63
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| auto act_scale_scalar = emit("tile.read", {args[2], zero_zero}, {}, "a8w8_act_scale_scalar"); | ||
| row_scaled = emit("tile.muls", {vec_fp32, act_scale_scalar}, {}, "a8w8_row_scaled"); | ||
| } else { | ||
| row_scaled = emit("tile.row_expand_mul", {vec_fp32, args[2]}, {}, "a8w8_row_scaled"); |
There was a problem hiding this comment.
When callers use the newly accepted act_scale shape [M], this lowering forwards that rank-1 tile directly to tile.row_expand_mul; that op's deducer requires the broadcast operand to be at least 2D with last dimension 1, so ConvertTensorToTileOps rejects a shape that tensor.a8w8_matmul_dequant explicitly accepts. Normalize [M] to [M, 1] before emitting the row-expand (and mirror the fix in the _acc path), or stop accepting rank-1 activation scales.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 2933ed9. Rank-1 A8W8 activation/weight scales are now normalized to 2-D shapes before the tile broadcast chain in both tensor.a8w8_matmul_dequant and tensor.a8w8_matmul_dequant_acc. I also added coverage for the [M]/[N] scale case in test_a8w8_matmul_dequant_reshapes_1d_scales_before_broadcast; the focused PyPTO tests passed: 9 passed.
| expr = ( | ||
| f"(torch.matmul({lhs}.to(torch.float32), {rhs}.to(torch.float32)) " | ||
| f"* {act_scale}.to(torch.float32) * {weight_scale}.to(torch.float32))" | ||
| ) |
There was a problem hiding this comment.
For the documented [M] / [N] scale forms, PyTorch aligns 1-D tensors with the last dimension, so (M,N) * act_scale[M] either raises when M != N or scales columns instead of rows when M == N; this makes the debug Torch codegen disagree with the op semantics for valid inputs. The handler should reshape rank-1 activation scales to [:, None] and rank-1 weight scales to [None, :] before multiplying.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 2933ed9. The Torch reference handlers for tensor.a8w8_matmul_dequant and tensor.a8w8_matmul_dequant_acc now reshape rank-1 activation scales to [-1, 1] and rank-1 weight scales to [1, -1] before applying the dequant multiplications. I also updated the torch codegen test to assert those reshape forms; the focused PyPTO tests passed: 9 passed.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for INT8 matrix multiplication with fused row/column scale dequantization by adding new operators such as tensor.a8w8_matmul_dequant, tensor.a8w8_matmul_dequant_acc, and tile.matmul_mx. It also updates tile.move to support optional target types and refines layout repair logic for tile.store. The review feedback identifies critical transposition bugs in several matmul conversion paths where the M dimension is incorrectly checked at index 0 even when a_trans is true. Additionally, the feedback highlights a potential null pointer dereference in the layout repair pass and suggests using the standard GetKwarg helper for cleaner and safer keyword argument retrieval.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| (void)kwargs; | ||
| INTERNAL_CHECK_SPAN(args.size() == 2, span) << "tensor.matmul conversion expects 2 args (lhs, rhs)"; | ||
| const bool nd = rank_of(args[0]) > 2 || rank_of(args[1]) > 2; | ||
| const std::string out_op = nd ? "tile.batch_matmul" : "tile.matmul"; | ||
| const std::string out_op = nd ? "tile.batch_matmul" : (HasStaticDim(args[0], 0, 1) ? "tile.gemv" : "tile.matmul"); |
There was a problem hiding this comment.
When selecting tile.gemv or tile.matmul, the code checks HasStaticDim(args[0], 0, 1) to determine if M == 1. However, if a_trans is true, the LHS matrix is transposed, meaning its shape is [K, M] and the M dimension is at index 1 instead of index 0. Checking index 0 when a_trans is true incorrectly checks if K == 1, which leads to selecting tile.matmul instead of tile.gemv for transposed decode/GEMV operations, or incorrectly selecting tile.gemv when K == 1. We should retrieve a_trans from kwargs and check the correct dimension index based on its value.
| (void)kwargs; | |
| INTERNAL_CHECK_SPAN(args.size() == 2, span) << "tensor.matmul conversion expects 2 args (lhs, rhs)"; | |
| const bool nd = rank_of(args[0]) > 2 || rank_of(args[1]) > 2; | |
| const std::string out_op = nd ? "tile.batch_matmul" : "tile.matmul"; | |
| const std::string out_op = nd ? "tile.batch_matmul" : (HasStaticDim(args[0], 0, 1) ? "tile.gemv" : "tile.matmul"); | |
| const bool a_trans = GetKwargOr<bool>(kwargs, "a_trans", false); | |
| INTERNAL_CHECK_SPAN(args.size() == 2, span) << "tensor.matmul conversion expects 2 args (lhs, rhs)"; | |
| const bool nd = rank_of(args[0]) > 2 || rank_of(args[1]) > 2; | |
| const bool is_m1 = a_trans ? HasStaticDim(args[0], 1, 1) : HasStaticDim(args[0], 0, 1); | |
| const std::string out_op = nd ? "tile.batch_matmul" : (is_m1 ? "tile.gemv" : "tile.matmul"); |
There was a problem hiding this comment.
Fixed in 2933ed9. I added MatmulMdimAxis(a_trans) and now tensor.matmul checks the logical M dimension after honoring a_trans before choosing tile.gemv vs tile.matmul. This keeps transposed decode/GEMV cases on the intended lowering path. Focused PyPTO tests passed: 9 passed.
| return var; | ||
| }; | ||
|
|
||
| const std::string mm_op = HasStaticDim(args[0], 0, 1) ? "tile.gemv" : "tile.matmul"; |
There was a problem hiding this comment.
Similar to tensor.matmul, the tensor.a8w8_matmul_dequant conversion checks HasStaticDim(args[0], 0, 1) to determine if M == 1. If a_trans is true, the M dimension is at index 1, not index 0. We should retrieve a_trans and check the correct dimension index to ensure tile.gemv is correctly selected for transposed decode/GEMV operations.
const bool a_trans = GetKwargOr<bool>(kwargs, "a_trans", false);
const bool is_m1 = a_trans ? HasStaticDim(args[0], 1, 1) : HasStaticDim(args[0], 0, 1);
const std::string mm_op = is_m1 ? "tile.gemv" : "tile.matmul";There was a problem hiding this comment.
Fixed in 2933ed9. tensor.a8w8_matmul_dequant now reads a_trans and uses the logical M axis through MatmulMdimAxis(a_trans) when deciding tile.gemv vs tile.matmul. This avoids checking K as M for transposed inputs. Focused PyPTO tests passed: 9 passed.
| return var; | ||
| }; | ||
|
|
||
| const std::string mm_op = HasStaticDim(args[1], 0, 1) ? "tile.gemv_acc" : "tile.matmul_acc"; |
There was a problem hiding this comment.
Similar to the other matmul operators, the tensor.a8w8_matmul_dequant_acc conversion checks HasStaticDim(args[1], 0, 1) to determine if M == 1. If a_trans is true, the M dimension of the LHS matrix (args[1]) is at index 1, not index 0. We should retrieve a_trans and check the correct dimension index to ensure tile.gemv_acc is correctly selected for transposed decode/GEMV operations.
const bool a_trans = GetKwargOr<bool>(kwargs, "a_trans", false);
const bool is_m1 = a_trans ? HasStaticDim(args[1], 1, 1) : HasStaticDim(args[1], 0, 1);
const std::string mm_op = is_m1 ? "tile.gemv_acc" : "tile.matmul_acc";There was a problem hiding this comment.
Fixed in 2933ed9. tensor.a8w8_matmul_dequant_acc now uses MatmulMdimAxis(a_trans) for the lhs operand when choosing tile.gemv_acc vs tile.matmul_acc, so transposed lhs inputs check the correct logical M dimension. Focused PyPTO tests passed: 9 passed.
| const std::string out_op = | ||
| nd ? "tile.batch_matmul_acc" : (HasStaticDim(args[1], 0, 1) ? "tile.gemv_acc" : "tile.matmul_acc"); |
There was a problem hiding this comment.
Similar to the other matmul operators, the tensor.matmul_acc conversion checks HasStaticDim(args[1], 0, 1) to determine if M == 1. If a_trans is true, the M dimension of the LHS matrix (args[1]) is at index 1, not index 0. We should retrieve a_trans and check the correct dimension index to ensure tile.gemv_acc is correctly selected for transposed decode/GEMV operations.
const bool a_trans = GetKwargOr<bool>(kwargs, "a_trans", false);
const bool is_m1 = a_trans ? HasStaticDim(args[1], 1, 1) : HasStaticDim(args[1], 0, 1);
const std::string out_op =
nd ? "tile.batch_matmul_acc" : (is_m1 ? "tile.gemv_acc" : "tile.matmul_acc");There was a problem hiding this comment.
Fixed in 2933ed9. tensor.matmul_acc now also honors a_trans through MatmulMdimAxis(a_trans) when checking whether the logical M dimension is 1, so gemv_acc selection is correct for transposed lhs operands. Focused PyPTO tests passed: 9 passed.
|
|
||
| INTERNAL_CHECK_SPAN(result_tile_type, op->span_) | ||
| << "ResolveBackendOpLayouts expects constrained op assignment targets to be TileType"; | ||
| const bool needs_output_repair = NeedsOutputRepair(result_tile_type, *layout_spec); |
There was a problem hiding this comment.
Calling NeedsOutputRepair(result_tile_type, *layout_spec) before checking if result_tile_type is null can lead to a null pointer dereference and crash (segmentation fault) if result_tile_type is indeed null. We should guard the call by checking if result_tile_type is non-null first.
| const bool needs_output_repair = NeedsOutputRepair(result_tile_type, *layout_spec); | |
| const bool needs_output_repair = result_tile_type && NeedsOutputRepair(result_tile_type, *layout_spec); |
| DataType out_dtype = DataType::FP32; | ||
| for (const auto& [key, value] : kwargs) { | ||
| if (key == "out_dtype") { | ||
| out_dtype = std::any_cast<DataType>(value); | ||
| } | ||
| } |
There was a problem hiding this comment.
Instead of manually looping over kwargs and performing std::any_cast<DataType>, we should use the standard GetKwarg helper which is safer, cleaner, and consistent with other files in the codebase (such as src/ir/op/tile_ops/memory.cpp and src/ir/op/tensor_ops/matmul.cpp).
| DataType out_dtype = DataType::FP32; | |
| for (const auto& [key, value] : kwargs) { | |
| if (key == "out_dtype") { | |
| out_dtype = std::any_cast<DataType>(value); | |
| } | |
| } | |
| DataType out_dtype = GetKwarg<DataType>(kwargs, "out_dtype", DataType::FP32); |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
python/pypto/ir/op/tile_ops.py (1)
434-462: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winKeep
target_typekeyword-only to preserve existingmove()call sites.Inserting
target_typeahead ofblayout/slayoutchanges the meaning of existing positional calls likemove(tile, mem, blayout, slayout): they now bind aTileLayouttotarget_typeand build the wrong op or fail type deduction.Proposed fix
def move( tile: Expr, target_memory: MemorySpace, - target_type: int | DataType | None = None, blayout: TileLayout | None = None, slayout: TileLayout | None = None, span: Span | None = None, + *, + target_type: int | DataType | None = None, ) -> Call:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/pypto/ir/op/tile_ops.py` around lines 434 - 462, The move() signature is currently letting target_type be filled positionally, which breaks existing calls that pass blayout and slayout after target_memory. Update move() in tile_ops.py so target_type remains keyword-only while preserving the current optional layout parameters, and keep its handling in the Call construction consistent with the existing move() and _get_span_or_capture flow.python/pypto/language/op/tile_ops.py (1)
529-550: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winKeep
target_typekeyword-only on the DSLmove()API.This has the same source-compatibility break as the IR builder: any existing positional call like
pl.tile.move(tile, mem, blayout, slayout)now mis-bindsblayoutintotarget_type.Proposed fix
def move( tile: Tile, target_memory: MemorySpace, - target_type: DataType | None = None, blayout: TileLayout | None = None, slayout: TileLayout | None = None, + *, + target_type: DataType | None = None, ) -> Tile:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/pypto/language/op/tile_ops.py` around lines 529 - 550, Keep the DSL move() API source-compatible by making target_type keyword-only in the move function signature, so positional calls like move(tile, mem, blayout, slayout) continue to bind blayout and slayout correctly. Update the move() definition in tile_ops.py to require target_type as a keyword-only argument, and ensure the call through to _ir_ops.move still passes target_type, blayout, and slayout by name.
🧹 Nitpick comments (1)
tests/st/runtime/ops/test_matmul.py (1)
927-1000: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPlease add an ST case for
a8w8_matmul_dequant_acc.This file only exercises the plain fused op, but the
_accvariant lowers through a different path (tile.matmul_acc/tile.gemv_accintests/ut/ir/transforms/test_a8w8_matmul_dequant.py). Without an end-to-end runtime case here, accumulator-path integration bugs can still slip past the current UT-only coverage.Also applies to: 1039-1045
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/st/runtime/ops/test_matmul.py` around lines 927 - 1000, Add a new ST runtime case for the accumulator variant by mirroring TestA8W8MatmulDequant with a sibling test class and program that invokes a8w8_matmul_dequant_acc instead of pl.a8w8_matmul_dequant. Reuse the same tensor definitions, get_program structure, orchestrator, and compute_expected flow, but make sure the new case exercises the _acc lowering path end-to-end through the runtime. Keep the existing plain fused-op test intact and place the new coverage alongside it so both variants are validated.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/pypto/language/op/tile_ops.py`:
- Around line 1031-1049: The new public tile.matmul_mx op is missing from the
torch debug/codegen registration, so add a matching handler in
torch_codegen.py’s _register_ops() for tile.matmul_mx alongside the existing
tile.matmul and tile.matmul_acc entries. Ensure the debug path recognizes this
op and routes it to the correct codegen behavior instead of falling back to
unsupported-op failure.
In `@src/backend/common/pto_ops_common.cpp`:
- Around line 3432-3436: The tile.move no-op elision in the shared-buffer reuse
path still allows dtype/layout-changing moves to be retargeted onto the source
buffer, which can leave an unsupported same-space pto.tmov. Update the
reuse/aliasing handling around the src_offset/dst_offset check in
src/backend/common/pto_ops_common.cpp so that tile.move operations with
differing dtype or effective TileView are treated as non-aliasing, or otherwise
force distinct MemRefs when the move changes layout. Use the existing same_dtype
and same_tile_view guard as the location to prevent MemoryReuse retargeting for
non-identical tiles.
In `@src/ir/op/tile_ops/matmul.cpp`:
- Around line 349-354: The out_dtype kwarg handling in matmul_mx bypasses the
standard lookup path and can change duplicate-key and type-error behavior.
Update the logic in the matmul tile op deducer to use the shared GetKwarg helper
against the kwargs vector<pair<string, any>> instead of manually iterating and
std::any_casting, so tile.matmul_mx matches the rest of the IR deducers and
raises the usual contextual TypeError on bad payloads.
In `@src/ir/transforms/op_conversion_registry.cpp`:
- Around line 969-985: The A8W8 dequant lowerings in
op_conversion_registry::ConvertA8W8MatMulDequant are passing 1-D scale tensors
directly into tile.row_expand_mul, tile.col_expand_mul, and tile.mul, so
normalize those operands to the expected 2-D broadcast shapes before emitting
the chain. Update both dequant paths to reshape [M] scales to [M, 1] and [N]
scales to [1, N] before the row/col scaling emits, keeping the existing
DetectRowBroadcast-style lowering behavior consistent. Add a transform test
covering the accepted 1-D scale form in
test_a8w8_matmul_dequant_accepts_1d_scales.
In `@tests/ut/ir/operators/test_op_registry.py`:
- Around line 395-401: Update the A8W8 registry tests to assert the specific
validator failure instead of catching any Exception. In the test cases around
ir.create_op_call for tensor.a8w8_matmul_dequant, replace the broad
pytest.raises(Exception) usage with the concrete validation exception type used
by the operator checks, or add a match= pattern for the expected dtype/shape
validation message. Apply the same change to both affected test blocks so they
verify the intended validator path rather than unrelated internal errors.
---
Outside diff comments:
In `@python/pypto/ir/op/tile_ops.py`:
- Around line 434-462: The move() signature is currently letting target_type be
filled positionally, which breaks existing calls that pass blayout and slayout
after target_memory. Update move() in tile_ops.py so target_type remains
keyword-only while preserving the current optional layout parameters, and keep
its handling in the Call construction consistent with the existing move() and
_get_span_or_capture flow.
In `@python/pypto/language/op/tile_ops.py`:
- Around line 529-550: Keep the DSL move() API source-compatible by making
target_type keyword-only in the move function signature, so positional calls
like move(tile, mem, blayout, slayout) continue to bind blayout and slayout
correctly. Update the move() definition in tile_ops.py to require target_type as
a keyword-only argument, and ensure the call through to _ir_ops.move still
passes target_type, blayout, and slayout by name.
---
Nitpick comments:
In `@tests/st/runtime/ops/test_matmul.py`:
- Around line 927-1000: Add a new ST runtime case for the accumulator variant by
mirroring TestA8W8MatmulDequant with a sibling test class and program that
invokes a8w8_matmul_dequant_acc instead of pl.a8w8_matmul_dequant. Reuse the
same tensor definitions, get_program structure, orchestrator, and
compute_expected flow, but make sure the new case exercises the _acc lowering
path end-to-end through the runtime. Keep the existing plain fused-op test
intact and place the new coverage alongside it so both variants are validated.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a1d43997-1bfa-45aa-840d-d4b6b293d94d
📒 Files selected for processing (21)
python/pypto/debug/torch_codegen.pypython/pypto/ir/op/tensor_ops.pypython/pypto/ir/op/tile_ops.pypython/pypto/language/__init__.pypython/pypto/language/op/__init__.pypython/pypto/language/op/tensor_ops.pypython/pypto/language/op/tile_ops.pypython/pypto/runtime/task_interface.pysrc/backend/common/pto_ops_common.cppsrc/ir/op/tensor_ops/matmul.cppsrc/ir/op/tile_ops/matmul.cppsrc/ir/op/tile_ops/memory.cppsrc/ir/transforms/infer_tile_memory_space_pass.cppsrc/ir/transforms/op_conversion_registry.cppsrc/ir/transforms/resolve_backend_op_layouts_pass.cpptests/st/runtime/ops/test_matmul.pytests/ut/codegen/test_pto_codegen_ops.pytests/ut/debug/test_torch_codegen.pytests/ut/ir/operators/test_op_registry.pytests/ut/ir/transforms/test_a8w8_matmul_dequant.pytests/ut/ir/transforms/test_resolve_backend_op_layouts_pass.py
| def matmul_mx( | ||
| lhs: Tile, | ||
| rhs: Tile, | ||
| act_scale: Tile, | ||
| weight_scale: Tile, | ||
| *, | ||
| out_dtype: DataType | None = None, | ||
| ) -> Tile: | ||
| """INT8 matrix multiplication with activation/weight scales.""" | ||
| call_expr = _ir_ops.matmul_mx( | ||
| lhs.unwrap(), | ||
| rhs.unwrap(), | ||
| act_scale.unwrap(), | ||
| weight_scale.unwrap(), | ||
| out_dtype=out_dtype, | ||
| ) | ||
| return Tile(expr=call_expr) | ||
|
|
||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Add the matching torch-debug handler for tile.matmul_mx.
matmul_mx is public here now, but the provided python/pypto/debug/torch_codegen.py:_register_ops() snippet still has entries for tile.matmul, tile.matmul_acc, tile.gemv*, etc. and no tile.matmul_mx. Any debug/codegen path that hits this new op will still fail as unsupported.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/pypto/language/op/tile_ops.py` around lines 1031 - 1049, The new
public tile.matmul_mx op is missing from the torch debug/codegen registration,
so add a matching handler in torch_codegen.py’s _register_ops() for
tile.matmul_mx alongside the existing tile.matmul and tile.matmul_acc entries.
Ensure the debug path recognizes this op and routes it to the correct codegen
behavior instead of falling back to unsupported-op failure.
| const bool same_dtype = src_tile->dtype_ == dst_tile->dtype_; | ||
| const bool same_tile_view = ir::tile_view_semantics::GetEffectiveTileView(*src_tile) == | ||
| ir::tile_view_semantics::GetEffectiveTileView(*dst_tile); | ||
| if (src_offset && dst_offset && src_offset->value_ == dst_offset->value_ && same_dtype && | ||
| same_tile_view) { |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== candidate reuse / allocation files ==="
fd -i 'memory' src/ir/transforms src/backend || true
fd -i 'memref' src || true
for f in $(fd -i 'memory' src/ir/transforms src/backend); do
echo "=== AST outline: $f ==="
ast-grep outline "$f" --view expanded || true
done
echo "=== tile.move + reuse/coalescing clues ==="
rg -nC3 --type=cpp 'tile\.move|byte_offset_|memref_|coalesc|reuse|target_type|blayout|slayout|valid_shape' srcRepository: hw-native-sys/pypto
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== pto_ops_common outline ==="
ast-grep outline src/backend/common/pto_ops_common.cpp --view expanded | sed -n '1,220p'
echo "=== tile.move / memory-reuse related definitions ==="
rg -nC4 'REGISTER_OP\("tile\.move"\)|set_output_reuses_input|not_inplace_safe|set_output_memory_inherit_input|MemoryReuse|InitMemRef|ReshapeTypeWithMemory|tile_view_semantics::GetEffectiveTileView|same_space|same-address|same address' src/ir src/backend | sed -n '1,260p'
echo "=== pto_ops_common around the cited lines ==="
sed -n '3410,3465p' src/backend/common/pto_ops_common.cpp
echo "=== tile op deducers around move/transpose/assemble/store ==="
sed -n '430,520p' src/ir/op/tile_ops/transform.cpp
sed -n '680,725p' src/ir/op/tile_ops/transform.cppRepository: hw-native-sys/pypto
Length of output: 35556
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== tile.move registration / type deduction ==="
rg -nC6 'REGISTER_OP\("tile\.move"\)|DeduceTileMoveType|tile\.move' src/ir/op src/backend/common | sed -n '1,240p'
echo "=== memory reuse / memref init references ==="
rg -nC4 'MemoryReuse|InitMemRef|set_output_reuses_input|output_memory_inherit_input|reuses_input|memref_' src/ir/transforms src/backend/common src/ir/op | sed -n '1,260p'
echo "=== infer_tile_memory_space_pass around layout rewriting ==="
sed -n '410,470p' src/ir/transforms/infer_tile_memory_space_pass.cpp
echo "=== any explicit move/layout non-aliasing contract mentions ==="
rg -nC3 'not_inplace_safe|cannot reuse|must not reuse|same-space same-address|tile\.move.*layout|layout.*tile\.move|dtype.*tile\.move' src | sed -n '1,220p'Repository: hw-native-sys/pypto
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== memory_reuse_pass: same-space / alias rules ==="
sed -n '380,460p' src/ir/transforms/memory_reuse_pass.cpp
echo
sed -n '1288,1388p' src/ir/transforms/memory_reuse_pass.cpp
echo
sed -n '1538,1565p' src/ir/transforms/memory_reuse_pass.cpp
echo "=== tile.move type deduction details ==="
sed -n '252,304p' src/ir/op/tile_ops/memory.cppRepository: hw-native-sys/pypto
Length of output: 13958
Keep layout-changing tile.move off shared buffers
tile.move is still eligible for MemoryReuse retargeting, so a dtype/layout-changing move can end up on the source buffer. This no-op elision only covers identical dtype + effective TileView; otherwise codegen can still emit an unsupported same-space pto.tmov. Mark these moves non-aliasing or force distinct MemRefs.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/backend/common/pto_ops_common.cpp` around lines 3432 - 3436, The
tile.move no-op elision in the shared-buffer reuse path still allows
dtype/layout-changing moves to be retargeted onto the source buffer, which can
leave an unsupported same-space pto.tmov. Update the reuse/aliasing handling
around the src_offset/dst_offset check in src/backend/common/pto_ops_common.cpp
so that tile.move operations with differing dtype or effective TileView are
treated as non-aliasing, or otherwise force distinct MemRefs when the move
changes layout. Use the existing same_dtype and same_tile_view guard as the
location to prevent MemoryReuse retargeting for non-identical tiles.
| DataType out_dtype = DataType::FP32; | ||
| for (const auto& [key, value] : kwargs) { | ||
| if (key == "out_dtype") { | ||
| out_dtype = std::any_cast<DataType>(value); | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Use the standard kwarg lookup path for out_dtype.
Lines 349-354 change the usual kwarg semantics here: a later duplicate key would win, and a bad payload can escape as std::bad_any_cast instead of the normal contextual TypeError. Please switch this to the shared GetKwarg pattern so tile.matmul_mx matches the rest of the IR deducers. Based on learnings, kwargs should be stored as std::vector<std::pair<std::string, std::any>> and looked up with range-based linear search using the standard GetKwarg helper.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/ir/op/tile_ops/matmul.cpp` around lines 349 - 354, The out_dtype kwarg
handling in matmul_mx bypasses the standard lookup path and can change
duplicate-key and type-error behavior. Update the logic in the matmul tile op
deducer to use the shared GetKwarg helper against the kwargs vector<pair<string,
any>> instead of manually iterating and std::any_casting, so tile.matmul_mx
matches the rest of the IR deducers and raises the usual contextual TypeError on
bad payloads.
Source: Learnings
vegetabledoww
force-pushed
the
feature/a8w8-matmul-dequant-fusion
branch
6 times, most recently
from
97f9f97 to
2933ed9
Compare
vegetabledoww
force-pushed
the
feature/a8w8-matmul-dequant-fusion
branch
5 times, most recently
from
7515544 to
1af3439
Compare
Keep PyPTO focused on generic compiler/runtime primitives needed by the Qwen3 A8W8 model path, while leaving model-specific dequant composition in pypto-lib. Retain only the current A8W8 whitelist: split pipe ids, cross-core pipe setup, tile layout/store repair, tuple/SSA stability, and Left/Right tile layout defaults. Drop the model-specific a8w8_matmul_dequant_acc composite op, tile.move target_type API, and automatic tensor.matmul M=1 gemv lowering. Validation: python -m pytest -q tests/ut/codegen/test_pto_codegen_ops.py tests/ut/ir/operators/test_op_registry.py tests/ut/ir/transforms/test_resolve_backend_op_layouts_pass.py Validation: Qwen3-14B A8W8 48-token generation for prompt '介绍一下北京故宫' passed with normal text quality, TTFT 8.021s, TPOT 67.2 ms/token.
vegetabledoww
force-pushed
the
feature/a8w8-matmul-dequant-fusion
branch
from
1af3439 to
ca3ffc7
Compare
|
Closing this PR after validation: the current A8W8 path builds and runs correctly on upstream pypto main without these changes, so this pypto-side patch is no longer required for the active pypto-lib/serving integration. |
## Summary Add the Qwen3-14B A8W8 kernel implementation used by the native serving path. This PR keeps the existing BF16 Qwen3 path isolated and introduces separate A8W8 prefill/decode modules for the quantized model: - add `prefill_hidden` support for Qwen3-14B A8W8 hidden-state prefill chunks - add the optimized A8W8 `decode_fwd` layer kernel used by serving decode - handle INT8 weights, activation/weight scales, paged KV cache scale metadata, and A8W8-specific decode constants inside the A8W8 modules - keep scheduling, model loading, tokenizer, and request orchestration in `pypto-serving` rather than carrying standalone lib-side runners - keep the ordinary BF16 Qwen3-14B execution path separate from the A8W8 path ## Implementation Notes The lib-side deliverable is intentionally limited to kernels and kernel-adjacent compatibility: - `models/qwen3/14b/prefill_fwd_a8w8.py` implements the A8W8 prefill hidden path - `models/qwen3/14b/decode_layer_a8w8.py` implements the A8W8 decode layer path used by serving - `golden/runner.py` compatibility changes keep generated/runtime artifacts runnable for the JIT path - obsolete standalone debug/golden entry points were removed during slimming; debug-stage switches in the kernel remain available for targeted numerical diagnosis ## Validation End-to-end validation was run through the native serving stack with this kernel PR plus the matching serving/backend PRs: - prompt: `介绍一下北京故宫` - generated tokens: 48 - output quality: normal Chinese continuation - TTFT: `7.054s` - TPOT: `67.8 ms/token` - decode throughput: `14.76 tok/s` Focused serving and PyPTO checks also passed in the matching PRs. ## Related PRs / Issues - Tracking issue: #665 - Serving-side PR: hw-native-sys/pypto-serving#48 - PyPTO backend/lowering PR: hw-native-sys/pypto#1920 Co-authored-by: vegetabledoww <vegetabledoww@users.noreply.github.com>
1 participant
Summary
This PR adds PyPTO compiler support for the Qwen3-14B A8W8 decode path by introducing the
pl.a8w8_matmul_dequant_acctensor op and lowering it into existing tile/PTO primitives.The op represents the common A8W8 projection pattern:
What Changed
A8W8 matmul-dequant-acc op
pl.a8w8_matmul_dequant_acc.INT32accumulator inputINT8lhs/rhs inputs[M]/[M, 1]activation scale shapes[N]/[1, N]weight scale shapesFP32/BF16outputsa_trans/b_transdimension inferenceTensor-to-tile lowering
tile.matmul_accfor generalM > 1tile.gemv_accforM == 1decodetile.castCodegen and compiler support
Tests
Adds focused unit coverage for:
tile.gemv_acctile.move(target_type=...)Validation
PyPTO focused tests
Run from the
pyptorepository root:Result:
Qwen3-14B A8W8 smoke generation
After rebuilding the native extension with:
8-token generation using the matching
pypto-lib/pypto-servingbranches completed successfully:A previous 48-token smoke run also produced coherent Chinese text with TPOT around 73.9 ms/token.
Notes
This PR contains PyPTO-side compiler functionality only. The model-level Qwen3-14B A8W8 decode/prefill kernels live in the companion
pypto-libPR and depend onpl.a8w8_matmul_dequant_accbeing available.