Add HIP backend with WMMA kernels for AMD RDNA4 (gfx12)#72
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (24)
📝 WalkthroughWalkthroughThe pull request adds an AMD RDNA4/gfx12 HIP backend with ROCm build support, WMMA kernels, Python dispatch integration, quantization and specialized operators, documentation, and gfx12-gated correctness tests. ChangesHIP gfx12 backend
Sequence Diagram(s)sequenceDiagram
participant scaled_mm_v2
participant hip_backend
participant dlpack_bindings
participant HIPLauncher
participant gfx12Kernel
scaled_mm_v2->>hip_backend: validate FP8 scaled matmul
hip_backend->>dlpack_bindings: pass tensors and stream
dlpack_bindings->>HIPLauncher: forward pointers and dtype codes
HIPLauncher->>gfx12Kernel: launch WMMA or GEMV kernel
gfx12Kernel-->>hip_backend: write scaled output
hip_backend-->>scaled_mm_v2: return HIP result
Possibly related PRs
🚥 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 |
3d47630 to
6d0c270
Compare
|
Review the following changes in direct dependencies. Learn more about Socket for GitHub.
|
There was a problem hiding this comment.
Actionable comments posted: 7
🤖 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/hip/__init__.py`:
- Around line 133-148: Add a concise docstring to stochastic_rounding_fp8
explicitly documenting that the kernel overwrites the caller-provided rng tensor
in place and that the returned fp8 view shares its storage. Keep the
implementation unchanged.
- Around line 348-422: Add the same parameter validation used by eager/CUDA to
quantize_convrot_w4a4_weight, dequantize_convrot_w4a4_weight, and
convrot_w4a4_linear before selecting the HIP kernels. Reject unsupported
quant_group_size values and ensure convrot_w4a4_linear validates linear_dtype
rather than silently treating every value as int4; preserve the existing eager
fallback and kernel behavior for supported parameters.
- Around line 72-83: Make HIP architecture checks device-aware instead of always
querying device 0: update _gfx_arch to accept the target torch device, and
ensure is_available, _register, and the _stream-based launch path use that same
device when selecting or validating the backend. Preserve gfx12 and extension
checks while preventing registration or execution against an adapter different
from the tensor’s device.
In `@comfy_kitchen/backends/hip/ops/gemv_awq.hip`:
- Around line 37-42: The HIP GEMV path currently decodes both scales and zero
points using only wscales.dtype. Update the setup that produces s_code and the
load_scalar calls in the kernel to ensure wzeros uses the same dtype as wscales,
preserving correct zero-point decoding and the existing 8-wide load behavior.
In `@comfy_kitchen/backends/hip/ops/per_tensor_fp8.hip`:
- Around line 17-94: Move the shared to_float overloads and OCP FP8 encoder
logic from comfy_kitchen/backends/hip/ops/per_tensor_fp8.hip lines 17-94 into
fp8_utils.h, include that header, and update callers to use the shared pack_fp8
symbol. In comfy_kitchen/backends/hip/ops/stochastic_round_fp8.hip lines 33-86,
remove the local to_float and pack_ocp_fp8 copies and use the shared
definitions, preserving bit-identical encoding behavior across both kernels.
- Around line 1-8: Prepend the identical SPDX copyright/license header used by
the sibling kernels to both per_tensor_fp8.hip (lines 1-8) and
stochastic_round_fp8.hip (line 1), placing it before the existing includes or
content; make no other changes.
In `@tests/test_hip_wmma.py`:
- Around line 382-388: Update test_hip_declines_on_non_gfx12 so it exercises the
non-gfx12 branch of hip_backend._register() by monkeypatching the device
architecture, such as torch.cuda.get_device_properties or gcnArchName, to a
non-gfx12 value and asserting registration declines. Alternatively, if retaining
the existing gfx12-only assertions, rename the test to reflect that it verifies
registration on gfx12 rather than non-gfx12 behavior.
🪄 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: 4c47a7a8-2ef5-45b4-9b3d-006e25915e13
📒 Files selected for processing (23)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pysetup.pytests/test_hip_wmma.py
6d0c270 to
6670e54
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/hip/gemm_wmma.h`:
- Around line 61-68: Add a compile-time divisibility check in the TileStager
template near kPerThread, asserting that kChunks is exactly divisible by THREADS
before allocating regs. Preserve the existing partition arithmetic while making
invalid tile parameter combinations fail during compilation.
In `@comfy_kitchen/backends/hip/ops/per_tensor_fp8.hip`:
- Around line 78-110: Guard both per-tensor FP8 launch functions in
comfy_kitchen/backends/hip/ops/per_tensor_fp8.hip (lines 78-110) and the
corresponding quantize-int8 launch sites in
comfy_kitchen/backends/hip/ops/quantize_int8.hip (lines 86-121) with the
existing empty-input early-return pattern used by stochastic_round_fp8: return
when numel or M is zero before computing blocks or dispatching kernels. Preserve
all non-empty dispatch behavior.
In `@comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip`:
- Around line 187-188: Update svdquant_gemm_kernel to index ascales with the
M_pad stride instead of M, and thread M_pad through launch_svdquant_gemm_kernel
and the binding that invokes it. Ensure M_pad matches the padded dimension used
by svdquant_quantize when reading activation scales.
🪄 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: c1aee5a0-0c66-4800-9a3a-09cb2d1281cf
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
6670e54 to
3a7c719
Compare
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/hip/CMakeLists.txt`:
- Around line 81-84: Update the HIP stochastic FP8 path around
stochastic_round_fp8 so zero-valued abs_value is handled explicitly before
floorf(log2f(abs_value)), returning the required zero quantization behavior
without relying on fast-math handling of -inf; alternatively remove -ffast-math
from the HIP target_compile_options if that is the intended fix.
In `@comfy_kitchen/scaled_mm_v2.py`:
- Around line 31-75: Update _hip_fp8_gemm to accept swizzle_a and swizzle_b,
reject non-default swizzle values by returning None before invoking the HIP
kernel, and update scaled_mm_v2’s call site to pass both parameters through so
unsupported layouts fall back to torch.
🪄 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: 0f2b1d46-e96e-4bfb-9384-a0a8ab4fac04
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
3a7c719 to
03a96b7
Compare
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 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/hip/hadamard.h`:
- Around line 42-53: Update launch_convrot_quant_int4_kernel to validate
group_size before invoking convrot_quant_kernel, accepting only the documented
values 16, 64, and 256. Reject zero, values above 256, and all other unsupported
sizes so the kernel cannot divide by zero or enter a non-advancing loop.
In `@comfy_kitchen/backends/hip/ops/adaln.hip`:
- Around line 98-108: Add an early return when the computed launch count N is
zero, before the CK_ADALN_LAUNCH dispatch in the AdaLN launcher. Preserve the
existing dtype selection and kernel launch behavior for non-empty batches.
In `@comfy_kitchen/backends/hip/ops/gemm_fp8.hip`:
- Around line 52-55: In the scaled_mm_fp8 dispatch path around the
gemv_fp8_kernel launch, return immediately when M == 0 or N == 0 before
constructing the grid or launching the kernel. Preserve the existing M <= 8
dispatch for non-empty shapes.
In `@comfy_kitchen/backends/hip/ops/gemv_awq.hip`:
- Around line 74-81: In the AWQ launch path containing the CK_AWQ_LAUNCH macro,
add an early return when M or N is zero before constructing or launching the
kernel grid. Preserve the existing launch behavior for positive dimensions and
allow empty outputs to complete without issuing a zero-dimension HIP launch.
In `@setup.py`:
- Around line 141-146: Update the compiler-path logic around c_compiler and
hip_compiler so the C compiler name is derived from the selected ROCm driver,
supporting amdclang++, hipcc, and clang++ layouts instead of always using clang.
Preserve the existing Windows executable suffix handling, append
CMAKE_C_COMPILER only when the mapped compiler exists, and keep
CMAKE_CXX_COMPILER based on hip_compiler.
🪄 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: a62dce30-1965-497c-9117-3607258b2bdd
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
03a96b7 to
f0253da
Compare
✅ Action performedReview finished.
|
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 `@comfy_kitchen/backends/hip/CMakeLists.txt`:
- Around line 81-84: Remove the HIP-specific -ffast-math option from
target_compile_options for _C, while preserving the existing -O3 option. This
keeps pack_fp8 in fp8_utils.h numerically precise so its isnan() handling
remains valid.
🪄 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: 71b7b8d3-c5bb-4124-8cdd-1bc1d67b9778
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
f0253da to
eacff43
Compare
✅ Action performedReview finished.
|
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/hip/ops/quantize_int8.hip`:
- Around line 107-109: Update the ConvRot launch surrounding
convrot_quant_kernel to guard the dynamic shared-memory allocation when K
exceeds the supported LDS budget. Add an explicit K/shared-memory bound or
perform the required HIP dynamic-shared-memory opt-in before launching, and
ensure oversized configurations are rejected or handled safely instead of
issuing a failing launch.
In `@comfy_kitchen/backends/hip/ops/stochastic_round_fp8.hip`:
- Around line 25-40: Update launch_stochastic_round_fp8_kernel and its kernel
size/index handling to use int64_t consistently, including the launch grid
calculation and idx computation. Remove the uint32_t truncation so arrays larger
than UINT32_MAX are indexed safely, matching the existing per_tensor_fp8
kernels.
In `@tests/test_hip_wmma.py`:
- Around line 254-256: Update the uint8 random generation in the test cases
around qw at the referenced locations, including the instances near lines 254,
289, and 369, to use an exclusive upper bound of 256 so every byte value through
0xFF can be generated. Keep the existing shapes, dtype, device, and subsequent
view behavior unchanged.
🪄 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: 7561b684-a5e7-4b95-ac34-4b7f0492c409
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
eacff43 to
771330a
Compare
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 12
🤖 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/hip/__init__.py`:
- Around line 154-174: Update stochastic_rounding_fp8 so non-contiguous rng
inputs preserve the documented in-place overwrite and returned storage-sharing
contract: either reject them with a clear validation error before calling the
kernel, or copy the kernel result back into the original tensor before
returning. Do not use rng.contiguous() without handling the copied result.
- Around line 82-86: Update _visible_gfx_arches to decline registration by
returning an empty list whenever any visible device’s _gfx_arch lookup returns
None; only return the architecture list when every device lookup succeeds,
preserving the existing unavailable/non-HIP behavior.
- Around line 639-681: Update apply_rope and apply_rope_split_half to validate
xq/xk compatibility before calling the combined _rope launch: use the separate
single-tensor paths whenever dtypes differ, and reject mismatched devices with
the established validation/error pattern. Preserve the existing shape-based
split behavior and ensure both RoPE variants apply identical dtype and device
checks.
In `@comfy_kitchen/backends/hip/dlpack_bindings.cpp`:
- Around line 199-220: Update apply_rope to validate the complete freqs layout
before launch_apply_rope_kernel: require trailing dimensions 4 and 5 to equal 2,
dimension 3 to equal xq.shape(3) / 2, and each leading frequency dimension to be
either 1 or match the corresponding xq dimension for broadcasting. Reject
invalid layouts with a runtime error before launching the kernel.
In `@comfy_kitchen/backends/hip/fp8_utils.h`:
- Around line 122-132: Document in fp8_to_float that the branch-free e4m3fn
decoding intentionally does not recognize NaN encodings such as 0x7F, which
decode to a finite value for GEMV performance. Keep the existing conversion
logic unchanged and add only a concise comment explaining the deliberate
omission.
- Around line 99-119: Update unpack_fp8 to detect the FP8 NaN encoding after
extracting the exponent and mantissa, before the zero/subnormal and normal
decode branches. Return a NaN value for that encoding while preserving existing
handling for zero, subnormal, and finite values in both supported formats.
In `@comfy_kitchen/backends/hip/ops/adaln.hip`:
- Around line 54-78: Update the variance computation in the AdaLN normalization
path to avoid using red_sq[0] * inv_d - mean * mean. After computing mean,
accumulate each element’s squared deviation from mean and reduce that value
across threads, then derive rstd from the resulting nonnegative variance plus
eps. Preserve the existing mean reduction and synchronization structure.
In `@comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip`:
- Around line 174-180: Bound the group index in the final tile’s flush loop
before accessing scale arrays. In the SVD quantization path around the `g`
calculation and `load_in(wscales, ...)`, skip the group when `g >= K /
kSvdGroupBytes` (or the equivalent valid group count), while preserving
processing for all valid groups and padded operands.
In `@comfy_kitchen/backends/hip/wmma_gfx12.h`:
- Around line 12-14: Update the explanatory comment near K_STEP to separate BF16
from FP8/IU8/IU4 fragment sizing: state that BF16 uses 8 half-precision elements
per lane and 16 bytes, while the other operand fragments use 8 bytes. Keep the
existing K_STEP and v2i context accurate without changing code.
In `@comfy_kitchen/scaled_mm_v2.py`:
- Around line 52-66: Update the HIP eligibility checks around
registry.is_available("hip") to return None unless every tensor operand,
including input, weight, scale_a, and scale_b, is on CUDA/ROCm and has the same
device as input; ensure any optional bias tensor is validated similarly before
_stream() or native-kernel dispatch.
In `@README.md`:
- Around line 58-69: Clarify the README statement around the gfx12
`torch._scaled_mm` and `torch._int_mm` calls so it applies only to requests
accepted by the HIP backend. Preserve the existing explanation that unsupported
requests fall back to PyTorch, making the supported-request scope explicit and
removing the apparent contradiction.
In `@setup.py`:
- Line 58: Update the HIP architecture resolution flow around the option parser,
get_extensions(), setup_hip_extension(), and build_cmake() so CLI and
environment overrides are resolved before extension discovery and device gating.
Validate the final resolved architecture list consistently, rejecting every
target that does not begin with gfx12, while ensuring overrides cannot bypass
this validation or be skipped due to the visible device architecture.
🪄 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: 39e3226d-3daf-4e94-8e5f-a9cf5267cfe4
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
771330a to
3ec58e1
Compare
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 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/hip/__init__.py`:
- Around line 557-580: In the SVDQuant input-validation path before the
_C.svdquant_quantize and _C.svdquant_lora_down calls, validate that
smooth.numel() equals K, lora_down.shape[0] equals K, and lora_src.shape matches
x.shape. Also require x, smooth, lora_down, and lora_src to share x.device,
raising ValueError for any mismatch before creating or passing raw pointers.
- Around line 521-541: Update the AWQ validation in the surrounding function
before launching _C.gemv_awq_w4a16: require positive, 8-aligned group_size and
K, ensure wscales and wzeros have shape (K / group_size, N), bias when present
has shape (N,), and move or validate every kernel operand, including qweight,
wscales, wzeros, bias, and out, on x.device. Preserve the existing divisibility
and qweight consistency checks.
- Around line 595-613: Extend validation in the SVDQuant GEMM setup around the
existing M, N, K, and R calculations: require wgt.shape[1] to equal
act.shape[1], lora_act_in.shape to be (M, R), lora_up.shape[0] to equal N, bias
to be one-dimensional with length N when present, and enforce the expected
wscales layout. Before calling _C.svdquant_gemm, ensure every operand tensor is
on act.device, including wgt, ascales, wscales, lora_act_in, lora_up, and bias.
- Around line 476-503: Update the ConvRot path to move qweight and bias to
x.device before creating contiguous tensors or passing their pointers to
_C.unpack_int4 and _C.convrot_w4a4_gemm. Apply this consistently in both the
linear_dtype == "int8" branch and the regular ConvRot branch, while preserving
existing dtype and bias handling.
- Around line 383-393: In the INT8 GEMM path surrounding the _C.int8_gemm
launch, move bias to x.device and make it contiguous before preparing the output
and launching the kernel. Validate that bias has shape (N,) and reject invalid
shapes before passing its pointer to _C.int8_gemm; preserve the existing None
handling.
In `@comfy_kitchen/backends/hip/ops/convrot_w4a4.hip`:
- Around line 89-97: Update the validation in the function containing the
convrot_quant_kernel launch to reject any K that is not divisible by group_size,
using the existing validation pattern alongside check_convrot_group_size and
check_convrot_k. Ensure invalid K/group_size combinations fail before
shared-memory sizing or kernel launch.
In `@comfy_kitchen/scaled_mm_v2.py`:
- Around line 58-87: Add a bias validation guard before the HIP routing in the
scaled FP8 path: when bias is present, require it to be 1-D, have
weight.shape[0] elements, and use float32, float16, or bfloat16; otherwise
return None. Keep the existing bias=None behavior and place the check alongside
the other input validations before calling hip.scaled_mm_fp8.
In `@pyproject.toml`:
- Line 6: Scope the “ninja” dependency in the project dependency list to Windows
using the appropriate environment marker, while preserving its existing HIP
backend comment and leaving dependency behavior unchanged on other platforms.
🪄 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: 1550a801-ac87-40cb-b9f9-0b02a4e17f77
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
3ec58e1 to
ab9b5ae
Compare
✅ Action performedReview finished.
|
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 `@comfy_kitchen/backends/hip/gemm_wmma.h`:
- Around line 63-68: Add a compile-time static assertion near the existing
kChunks % THREADS assertion in the relevant WMMA tile configuration, requiring
BKB to be divisible by 16 before kChunksPerRow truncates it. Preserve the
existing THREADS divisibility assertion and its behavior.
🪄 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: da33a495-422d-421d-952a-73078ce0abcb
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
ab9b5ae to
9ebb4cf
Compare
✅ Action performedReview finished.
|
✅ Action performedReview finished.
|
✅ Action performedReview finished.
|
9ebb4cf to
9dfd698
Compare
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip (1)
191-192: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
ascalesstride still usesM, notM_pad— the earlier "fix" didn't quite land.
svdquant_quant_kernelstores scales with the padded stride at Line 74 (store_scalar(ascales, s_code, static_cast<int64_t>(g) * M_pad + m, scale)), i.e. layout(K/64, M_pad). But here the GEMM reads them back withg * M + r. WheneverM != M_pad, every group after the first reads the wrong column — no crash, just a silent scale-shuffle that skews the output. This was previously reported and marked addressed, yetsvdquant_gemm_kernel/launch_svdquant_gemm_kernelstill receive noM_pad, so the load stride never got padded.Thread
M_padthrough the kernel, the launcher, and the binding, then index with it here.🧮 Proposed direction
- const float as = - load_in(ascales, static_cast<int64_t>(g) * M + r, as_code); + const float as = + load_in(ascales, static_cast<int64_t>(g) * M_pad + r, as_code);Add
int M_padtosvdquant_gemm_kernel(line 100) andlaunch_svdquant_gemm_kernel(line 265), and forward it from the dlpack binding so it matches the padded dimension used bysvdquant_quant_kernel.#!/bin/bash # Confirm M_pad can differ from M and is (or isn't) threaded into the gemm path. rg -n -C3 'M_pad|launch_svdquant_gemm_kernel|svdquant_gemm' \ comfy_kitchen/backends/hip/__init__.py \ comfy_kitchen/backends/hip/dlpack_bindings.cpp \ comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip🤖 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/hip/ops/svdquant_w4a4.hip` around lines 191 - 192, Thread the padded dimension M_pad through svdquant_gemm_kernel, launch_svdquant_gemm_kernel, and the DLPack binding, matching the existing quantization layout. Update the ascales load in svdquant_gemm_kernel to use g * M_pad + r instead of g * M + r, while preserving the existing M dimension for output and logical matrix sizing.
🤖 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/hip/dlpack_bindings.cpp`:
- Around line 202-233: Extend apply_rope validation before
launch_apply_rope_kernel to require xk and xk_out to be supplied together, and
require every supplied operand (xk, xq_out, and xk_out) to match xq in shape,
strides, and dtype. Reject any mismatch with a runtime error before the native
launch, since the kernel only receives xq’s layout and dtype metadata.
In `@comfy_kitchen/backends/hip/gemm_wmma.h`:
- Around line 110-114: Add compile-time assertions near the `kThreads`/`kStride`
geometry definitions to require `BM == WARPS_M * TM * 16` and `BN == WARPS_N *
TN * 16`. Keep the existing shared-memory layout unchanged and make invalid
tile-versus-warp instantiations fail during compilation.
---
Duplicate comments:
In `@comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip`:
- Around line 191-192: Thread the padded dimension M_pad through
svdquant_gemm_kernel, launch_svdquant_gemm_kernel, and the DLPack binding,
matching the existing quantization layout. Update the ascales load in
svdquant_gemm_kernel to use g * M_pad + r instead of g * M + r, while preserving
the existing M dimension for output and logical matrix sizing.
🪄 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: 78da1b95-fe9d-4a92-a5d0-cabf4a44b372
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
9dfd698 to
33f250f
Compare
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip (1)
191-192: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
ascalesis read with strideMbut written with strideM_pad— the strides still don't line up.svdquant_quant_kernelstores the scale atg * M_pad + m(Line 74, layout documented as(K/64, M_pad)), yetsvdquant_gemm_kernelreads it back atg * M + rhere, and neither the kernel signature (Line 100) norlaunch_svdquant_gemm_kernel(Line 262) carriesM_pad. WheneverM_pad != M, every groupg >= 1picks up a scale from the wrong row — no crash, just a quiet scale two-step that corrupts the result.This is the same concern raised previously (cr
20fa959e, marked addressed in 771330a), but the read path in the current code still usesM. Please threadM_padthrough the kernel and launcher, or confirmM_padcan never exceedM.🧮 Proposed direction
- const float as = - load_in(ascales, static_cast<int64_t>(g) * M + r, as_code); + const float as = + load_in(ascales, static_cast<int64_t>(g) * M_pad + r, as_code);Add
int M_padtosvdquant_gemm_kernelandlaunch_svdquant_gemm_kernel, and pass it from the binding (same value handed tolaunch_svdquant_quant_kernel).#!/bin/bash # Confirm whether ascales is padded (M_pad != M) at the binding/wrapper layer, # and whether the gemm launcher/kernel ever receives M_pad. set -euo pipefail echo "== store stride (M_pad) vs read stride (M) in the kernel ==" rg -nP 'ascales|M_pad' comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip echo; echo "== how M_pad is computed and how ascales is allocated ==" rg -nP -C3 'M_pad|svdquant_quant|svdquant_gemm|ascales|quantize_svdquant' \ comfy_kitchen/backends/hip/__init__.py comfy_kitchen/backends/hip/dlpack_bindings.cpp🤖 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/hip/ops/svdquant_w4a4.hip` around lines 191 - 192, Update svdquant_gemm_kernel and launch_svdquant_gemm_kernel to accept M_pad, pass the binding’s existing M_pad value through the launcher, and use it for the ascales index in the load_in call so reads match svdquant_quant_kernel’s g * M_pad + m layout.
🤖 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/hip/__init__.py`:
- Around line 246-255: Update the public scaled_mm_fp8 helper before
bias.contiguous() to validate bias using the same 1D, expected-length, dtype,
and device contract enforced by the other GEMM paths. Reject invalid bias values
before passing _dl(bias) to _C.scaled_mm_fp8, while preserving the existing
contiguous conversion and valid-call behavior.
In `@comfy_kitchen/backends/hip/ops/apply_rope.hip`:
- Around line 66-67: Before launching the HIP kernel that computes foff in the
apply-RoPE path, validate every leading freqs_cis dimension is either 1 or equal
to the corresponding x dimension; retain the existing ExactDims(6) rank check
and reject any smaller non-1 extent before launch.
---
Duplicate comments:
In `@comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip`:
- Around line 191-192: Update svdquant_gemm_kernel and
launch_svdquant_gemm_kernel to accept M_pad, pass the binding’s existing M_pad
value through the launcher, and use it for the ascales index in the load_in call
so reads match svdquant_quant_kernel’s g * M_pad + m layout.
🪄 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: dfd5d7e8-a0d1-487f-a5fb-ec6c17c9ad52
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
33f250f to
a36572b
Compare
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
pyproject.toml (1)
2-8: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winDeclare CMake in the build requirements.
setup.pyshells out tocmake, so a clean source install on a machine without system CMake will fail before the build even starts. The docs already say CMake is required — let’s make the manifest stop playing hide-and-seek.Proposed fix
requires = [ "setuptools>=61.0", "wheel", "nanobind>=2.0.0", + "cmake>=3.21", "ninja",🤖 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 `@pyproject.toml` around lines 2 - 8, Add the CMake package to the pyproject.toml build-system requires list alongside the existing setuptools, wheel, nanobind, and ninja requirements so setup.py can invoke CMake during clean source installs.
🤖 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/hip/__init__.py`:
- Around line 631-673: Validate in scaled_mm_svdquant_w4a4 that the derived K
value is divisible by the SVDQuant group size of 64 before computing or
validating scale shapes and dispatching _C.svdquant_gemm. Reject non-divisible K
with a clear ValueError; preserve the existing path for valid K values. Apply
the same validation to the corresponding code path around the additionally
referenced implementation.
- Around line 227-259: Enforce the WMMA K-step requirement in the exported
launchers `scaled_mm_fp8` and the additional launcher covering the referenced
range: after deriving the input inner dimension K and before invoking the native
`_C` kernel, reject any K where K % 16 != 0 with a clear ValueError. Preserve
the existing dimension validation and valid aligned-K execution path.
- Around line 150-180: Update quantize_per_tensor_fp8 and
dequantize_per_tensor_fp8 to validate that scale contains exactly one element,
then make it contiguous and colocate it with x’s device before passing
_dl(scale) to the kernels. Preserve scalar handling while rejecting
multi-element scales, and apply the same preparation consistently in both
functions.
- Around line 184-207: Update stochastic_rounding_fp8 to validate that rng has
uint8 dtype before invoking the kernel or returning rng.view(output_type); raise
a clear ValueError for any other dtype while preserving the existing device,
shape, and contiguity checks.
In `@comfy_kitchen/backends/hip/CMakeLists.txt`:
- Around line 77-84: Update the _C nanobind_add_module declaration to pass
NOMINSIZE, explicitly disabling nanobind’s default -Os while retaining the
existing HIP -O3 and -ffast-math compile options.
In `@comfy_kitchen/backends/hip/ops/gemm_fp8.hip`:
- Around line 31-34: Update the GEMM FP8 vector-load loops around the visible
uint4 loads and the corresponding code at the additional occurrence to reject or
handle any K value not divisible by 16 before issuing 16-byte loads. Ensure
_C.scaled_mm_fp8 cannot reach these loads with an unaligned K, while preserving
the existing path for valid multiples of 16.
In `@comfy_kitchen/backends/hip/ops/gemm_int8.hip`:
- Around line 27-29: Update the native GEMM INT8 entry point and its associated
launch path to reject unsupported K values before the loops issuing uint4 loads
in the HIP implementation. Validate that K is aligned to the full 16-byte load
width, including the path referenced by the additional affected range, and
return an error or failure status through the existing boundary mechanism before
any loads occur.
In `@comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip`:
- Around line 73-74: Update the activation-scale indexing to use M_pad
consistently instead of M, and thread M_pad through the relevant kernel,
launcher, declaration, and binding interfaces. Ensure GEMM reads scales with the
same g * M_pad stride used by the quantizer, including all corresponding call
sites and declarations identified in the diff.
---
Outside diff comments:
In `@pyproject.toml`:
- Around line 2-8: Add the CMake package to the pyproject.toml build-system
requires list alongside the existing setuptools, wheel, nanobind, and ninja
requirements so setup.py can invoke CMake during clean source installs.
🪄 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: aa66b513-e0b0-4a84-8731-22457895f0e1
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
a36572b to
ac10281
Compare
✅ Action performedReview finished.
|
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 `@comfy_kitchen/backends/hip/ops/per_tensor_fp8.hip`:
- Around line 70-118: Update launch_quantize_per_tensor_fp8_kernel and
launch_dequantize_per_tensor_fp8_kernel to call hipGetLastError() immediately
after each kernel launch and surface any launch failure using the existing
backend error-handling convention, matching stochastic_round_fp8 behavior. Apply
the check to every dtype-specific launch without changing the zero-element early
return or dtype dispatch.
🪄 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: 466251f9-0382-4fd0-8c71-60ed86695a10
📒 Files selected for processing (24)
README.mdcomfy_kitchen/__init__.pycomfy_kitchen/backends/hip/CMakeLists.txtcomfy_kitchen/backends/hip/__init__.pycomfy_kitchen/backends/hip/dlpack_bindings.cppcomfy_kitchen/backends/hip/epilogue.hcomfy_kitchen/backends/hip/fp8_utils.hcomfy_kitchen/backends/hip/gemm_wmma.hcomfy_kitchen/backends/hip/hadamard.hcomfy_kitchen/backends/hip/ops/adaln.hipcomfy_kitchen/backends/hip/ops/apply_rope.hipcomfy_kitchen/backends/hip/ops/convrot_w4a4.hipcomfy_kitchen/backends/hip/ops/gemm_fp8.hipcomfy_kitchen/backends/hip/ops/gemm_int8.hipcomfy_kitchen/backends/hip/ops/gemv_awq.hipcomfy_kitchen/backends/hip/ops/per_tensor_fp8.hipcomfy_kitchen/backends/hip/ops/quantize_int8.hipcomfy_kitchen/backends/hip/ops/stochastic_round_fp8.hipcomfy_kitchen/backends/hip/ops/svdquant_w4a4.hipcomfy_kitchen/backends/hip/wmma_gfx12.hcomfy_kitchen/scaled_mm_v2.pypyproject.tomlsetup.pytests/test_hip_wmma.py
✅ Action performedReviews resumed. |
Signed-off-by: 0xDELUXA <djernovevo@gmail.com>
ac10281 to
31e4bfe
Compare
✅ Action performedReview finished.
|
|
We need to figure out how to package this in the current comfy-kitchen wheels, I don't want to have separate wheels for nvidia/amd/intel/etc... |
Thanks for your feedback. I agree. This is an experiment on my end to see how much performance I can get out of gfx12 using its new WMMA kernels. It started from rattus' work on adding a HIP backend to comfy-kitchen, which would provide broader AMD support. I wonder why it's Linux-only and why it wasn't continued. |
|
Closing as superseded by #74. |
Adds
comfy_kitchen/backends/hip, a native backend for AMD RDNA4 (gfx12) built on thev_wmma_*_w32_gfx12intrinsics. It does not link hipBLAS/hipBLASLt, and no linear op calls into them; every GEMM is compiled from source in the backend directory. Weight preparation and unsupported configurations fall back to eager.Design
fp8, int8 and int4 all hold 8 bytes per lane per WMMA, so one K-step is 16 bytes of a row whether those bytes hold 16 fp8/int8 values or 32 int4 nibbles. The tile kernel (
gemm_wmma.h) is byte-addressed and shared by the fp8, int8 and int4 GEMMs; the element type appears only in theMmapolicy. It is software-pipelined at both the global (tile prefetch) and LDS (fragment) level, and uses grouped block ordering for L2 locality.The lane/fragment layout for the gfx12 w32 WMMA instructions is documented in
wmma_gfx12.h.Ops
scaled_mmf32_16x16x16_fp8_fp8int8_lineari32_16x16x16_iu8convrot_w4a4_lineari32_16x16x32_iu4, rotation fused into the quantizer;linear_dtype="int8"unpacks the weight onto the iu8 kernelscaled_mm_svdquant_w4a4i32_16x16x32_iu4, per-64 group scales, fused LoRA-upgemv_awq_w4a16Plus fp8/int8 quantization,
adaln, and the fourapply_ropevariants.On ROCm,
int8_linearfell through to eager, which routes torch._int_mm to hipBLASLt; it is now a WMMA kernel. NVFP4 and MXFP8 remain on eager: RDNA4 has neither fp4 WMMA nor microscaling hardware.Dispatch
The backend registers only when every visible device is known to be gfx12 (the intrinsics have no gfx11 encoding); a device whose architecture cannot be read counts against it, since it cannot be shown to be gfx12. Registration is per-process while kernels launch on the tensor's own device, so one non-gfx12 GPU makes it decline and dispatch falls through to triton/eager.
COMFY_KITCHEN_DISABLE_HIP=1removes it from dispatch.ROCm's fp8
scaled_mmis competitive on RDNA4 and is faster than the WMMA kernel on square, deep-K and small-M shapes, while the WMMA kernel is faster at large N. fp8 is routed to WMMA on every call the kernel supports, with no shape gate, to keep the backend free of BLAS calls. Calls outside its domain (swizzled operands, scaling other than tensor-wise, a K that is not a multiple of 16, a bias that is not 1-D of length N, operands not on the input's device) fall back to torch.Build
setup.pybuilds the extension whenever a ROCm toolchain is present, on Linux and Windows alike:pip install .ROCm is discovered from
ROCM_HOME/ROCM_PATH/HIP_PATH, then from the piprocm-sdkwheel (viarocm-sdk path --root, which a ROCm PyTorch build already pulls in), then from/opt/rocm. No environment variables, noCC/CXXoverride and no Visual Studio developer shell are needed: the ROCm clang drives C, C++ and HIP alike, and locates the MSVC toolchain itself. The HIP extension pins the Ninja generator on every platform; CMake's default generator on Windows is Visual Studio, which cannot build the HIP language at all.Target architectures default to the gfx12 GPUs the build machine can see, and to
gfx1200;gfx1201when it can see none (CI, cross-builds). To pick them explicitly:Because the kernels are gfx12-only, the extension is skipped rather than built when the visible AMD GPUs are all of another architecture.
COMFY_KITCHEN_BUILD_HIP=1turns that skip (and a missing toolchain) into an error;COMFY_KITCHEN_BUILD_NO_HIP=1suppresses the backend entirely.Tests
tests/test_hip_wmma.py(95 tests) skips unless the backend registered, so it is inert on CI without an RDNA4 device. Kernels are validated against eager, and against an fp32 reference where the two legitimately differ (the HIP kernels quantize and reduce in fp32 where eager works in bf16, so they are not bitwise equal to eager, and are closer to the fp32 result). One test trapstorch._scaled_mmandtorch._int_mmand exercises every quantized path to assert neither is reached.Tested on gfx1200 (RDNA4), ROCm 7.15 / PyTorch 2.12, Windows.