RDNA-aware autotune configs for the INT8 matmul kernels#69
RDNA-aware autotune configs for the INT8 matmul kernels#69liminfei-amd wants to merge 5 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughINT8 matmul autotuning now selects platform- and architecture-specific configurations, prunes candidates per device, and includes ChangesINT8 autotuning
Sequence Diagram(s)sequenceDiagram
participant int8_linear
participant INT8Kernels
participant INT8_AUTOTUNE_KWARGS
participant torch.cuda.get_device_properties
int8_linear->>INT8Kernels: launch with device_index
INT8Kernels->>INT8_AUTOTUNE_KWARGS: tune by shape and device_index
INT8_AUTOTUNE_KWARGS->>torch.cuda.get_device_properties: inspect architecture
torch.cuda.get_device_properties-->>INT8_AUTOTUNE_KWARGS: return device-specific candidates
INT8_AUTOTUNE_KWARGS-->>INT8Kernels: select configuration
Suggested reviewers: 🚥 Pre-merge checks | ✅ 2✅ Passed checks (2 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
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 |
|
Thanks for merging Comfy-Org/ComfyUI#14862 (default-enabling the Triton backend on ROCm) — much This PR is the companion on the kernel side: it tunes the INT8 matmul autotune configs for RDNA |
|
What about add waves_per_eu into the configs? gfx1103 is a apu, limited by memory bandwidth, dgpu may need to use different configs. |
|
Thanks @LuXuxue! I checked |
|
I know this might not be the place to ask this but your recent change in comfyui's enabling triton backend for ROCM Comfy-Org/ComfyUI#14862 totally killed the int8 model usage for me. RDNA2, triton-windows 3.7.1.post27. I was using comfy without triton-backend and with custom node everything worked perfectly now neither native load diffusion model nor the custom node works, comfy just freezes and only a reset helps. |
|
Updated this PR after broader multi-GPU and small-shape validation. The config pool is now selected lazily for the device serving each call, and For This PR remains scoped to Triton INT8 matmul performance; it does not change backend enablement or optimize CDNA. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
comfy_kitchen/backends/triton/quantization.py (1)
909-930: 🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick winDrop
tl.constexprfromdevice_indexin both int8 matmul kernels Keep it in the autotune key/pruning path, but not in the kernel signature specialization; otherwise Triton bakes the device index into the compiled binary and reruns autotune per GPU index for no codegen benefit—less waiting, more skating.🤖 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 `@comfy_kitchen/backends/triton/quantization.py` around lines 909 - 930, The device_index parameter is incorrectly specialized as a compile-time value in _int8_matmul_dequant_kernel and the corresponding second int8 matmul kernel. Remove tl.constexpr from both kernel signatures while preserving device_index in the autotune key and pruning configuration.
🤖 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 `@comfy_kitchen/backends/triton/quantization.py`:
- Around line 909-913: Extract the shared Triton autotune arguments into one
reusable kwargs definition, then apply it to both the `_int8_matmul_kernel` and
`_int8_matmul_dequant_per_row_kernel` decorators. Ensure both kernels continue
using identical `configs`, `key`, and `prune_configs_by` values so future
device-scoping changes are made in one place.
---
Outside diff comments:
In `@comfy_kitchen/backends/triton/quantization.py`:
- Around line 909-930: The device_index parameter is incorrectly specialized as
a compile-time value in _int8_matmul_dequant_kernel and the corresponding second
int8 matmul kernel. Remove tl.constexpr from both kernel signatures while
preserving device_index in the autotune key and pruning configuration.
🪄 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: ASSERTIVE
Plan: Pro Plus
Run ID: 44bf00ff-0090-47b9-9f26-f59baa7a3871
📒 Files selected for processing (2)
comfy_kitchen/backends/triton/quantization.pytests/test_int8.py
|
I find it is strange that more autotune configs, lower speed. |
The two INT8 matmul kernels ship 6 NVIDIA-tuned autotune configs (num_stages=3-4). On AMD RDNA, num_stages=2 with deeper block_k is faster. Add _int8_autotune_configs() which returns an RDNA-tuned pool when torch.version.hip is set and keeps the NVIDIA config list unchanged for CUDA. INT8 GEMM is int32-exact, so this changes speed only. Fixes Comfy-Org#68 Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
Occupancy tuning via waves_per_eu helps on bandwidth-limited RDNA APUs (gfx1103/gfx1151: up to ~1.15x); dGPUs autotune-select the existing base configs (within noise). Added as an autotune dimension rather than hardcoding per-model configs. Thanks to @LuXuxue for the suggestion. Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
Select the RDNA INT8 config pool lazily for the device serving each call, include the device index in the autotune key, and retain the exact default pool on CUDA, CDNA, unknown architectures, or probe failure. Prune the unstable block_k=128 candidate for M <= 128, where repeated fresh-process tests on gfx1151 showed a regression despite faster configs being available. Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
Keep both INT8 matmul kernels on one configs, cache-key, and pruning contract so future device-scoping changes cannot diverge between kernel paths. Add a focused assertion covering the shared mapping. Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
Keep ten candidates per device while replacing an unused gfx12-oriented tile with a small-M configuration on gfx11. Retain the original gfx12 pool, preserve useful BK128 candidates, and prune only the unstable 128x128x128 tile for small M. Keep device_index in the autotune key but mark it non-specializing in the generated kernels, so heterogeneous devices tune independently without duplicate code variants. Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
25f11bb to
04ac357
Compare
|
Thanks for the detailed results, @LuXuxue. I updated the PR to use your |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/test_int8.py`:
- Around line 361-377: Extend test_int8_autotune_cache_key_includes_device with
focused launch tests for int8_linear, intercepting both regular and per-row
kernel launch paths. Assert each intercepted launch receives x.device.index as
the device_index argument, while preserving the existing decorator metadata
assertions.
🪄 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: ASSERTIVE
Plan: Pro Plus
Run ID: 075743e9-0633-46b8-a607-c3ebf6bce022
📒 Files selected for processing (2)
comfy_kitchen/backends/triton/quantization.pytests/test_int8.py
| def test_int8_autotune_cache_key_includes_device(): | ||
| from comfy_kitchen.backends.triton import quantization | ||
|
|
||
| assert quantization._INT8_AUTOTUNE_KWARGS == { | ||
| "configs": quantization._INT8_MATMUL_CONFIGS, | ||
| "key": ["m", "n", "k", "device_index"], | ||
| "prune_configs_by": quantization._INT8_PRUNE_CONFIGS_BY, | ||
| } | ||
|
|
||
| for kernel in ( | ||
| quantization._int8_matmul_dequant_kernel, | ||
| quantization._int8_matmul_dequant_per_row_kernel, | ||
| ): | ||
| assert "device_index" in kernel.keys | ||
| device_index = kernel.fn.params[kernel.fn.arg_names.index("device_index")] | ||
| assert not device_index.is_constexpr | ||
| assert device_index.do_not_specialize |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Exercise launch-time device_index propagation.
This verifies decorator metadata, but not that int8_linear forwards x.device.index to both kernel launches. Add focused tests intercepting each launch path; otherwise device isolation can slip through the cache cracks.
🤖 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/test_int8.py` around lines 361 - 377, Extend
test_int8_autotune_cache_key_includes_device with focused launch tests for
int8_linear, intercepting both regular and per-row kernel launch paths. Assert
each intercepted launch receives x.device.index as the device_index argument,
while preserving the existing decorator metadata assertions.
3 participants
RDNA-aware autotune configs for the INT8 matmul kernels
What
This keeps the existing CUDA/CDNA behavior and selects a ten-candidate INT8 Triton pool for the RDNA generation serving each call.
gfx11*gets a small-M64x64x128, eight-warp tile;gfx12*retains its existing pool. ForM <= 128, only the unstable128x128x128tile is removed.device_indexremains in the autotune key so heterogeneous devices tune separately, but is marked non-specializing in both kernels so the scalar value does not create duplicate code variants.Why
The original six default configs use
num_stages=3-4. Shallownum_stages=2pools perform better on the tested consumer RDNA devices. Follow-up SDXL testing also showed that one pool should not be shared blindly across RDNA generations: the extra small-M tile helpsgfx11, while replacing agfx12tile regresses that generation.Validation
main5a997b17; cumulative diff is two files (+236/-21).tests/test_int8.pyon gfx1100, gfx1151, and gfx1201: 80 passed / 33 skipped / the same 5 CUDA-extension-only failures as the before arm.device_indexcreated one extra compiled entry before this update and zero afterward; the output remained finite.Reporter follow-up
The exact 1.70 vs 1.79 s/it E2E run could not be reproduced without the reporter's private config/model setup. On gfx1151, their six-config pool was not slower than the one-config pool in steady-state kernel time; its first-use autotune cost was much larger. However, their feedback exposed a real selection gap in this PR's prior pool, and the generation-specific update above improves the matching SDXL shapes without regressing gfx1201.
Fixes #68.
AI usage disclosure: this update was prepared with AI assistance. The submitting human remains responsible for reviewing the changes and validation before publication.