HIP: vectorize rowwise INT8 quantization with 16-byte loads

[tangzzycc]

Summary

Implements Change 2 proposed in #102:

• Load 16 bytes per thread with uint4 for FP32, FP16, and BF16 inputs.
• Use the vectorized kernel only when row alignment and K permit it; otherwise
  keep the existing scalar fallback.
• Preserve the existing rounding, clamping, and scale behavior.

The vectorized path uses the portable HIP uint4 type and has no architecture-specific guard, so it also applies to other supported HIP architectures. Performance was measured on gfx1151.

Performance

Measured on AMD Radeon 8060S (gfx1151), ROCm 7.2, with aligned BF16 inputs:

Shape (M x K)	Before	After	Speedup
4096 x 1024	0.0564 ms	0.0269 ms	2.09x
8192 x 2048	0.4417 ms	0.2196 ms	2.01x
4096 x 4096	0.4107 ms	0.2268 ms	1.81x
18432 x 2048	0.9579 ms	0.4999 ms	1.92x
512 x 14336	0.0996 ms	0.0417 ms	2.39x
4096 x 14336	1.3498 ms	0.8246 ms	1.64x

Geometric-mean speedup: 1.96x.

For large BF16 shapes, one-read DRAM-equivalent bandwidth is about 214–229 GB/s, consistent with the 256 GB/s bandwidth of this gfx1151 system.

Testing

CMAKE_BUILD_PARALLEL_LEVEL=16 python setup.py \
  build_ext --inplace --no-cuda --hip --hip-archs=gfx1151

pytest tests/test_qdq.py::TestQuantizeINT8 -q
pytest tests/test_hip_wmma.py -q
pytest tests/test_int8.py -q

• INT8 quantization: 4 passed
• HIP WMMA: 349 passed
• INT8: 62 passed, 39 skipped
• Vectorized and fallback outputs match bit-for-bit for FP32, FP16, and BF16,
  including aligned and offset-contiguous inputs.

[github-actions]

🎉 Thank you for your contribution, we really appreciate it! 🎉

Like many open source projects, we require contributors to sign our Contributor License Agreement (CLA). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you:

• Confirm that you own your contribution.
• Keep the right to reuse your own code.
• Grant us a copyright license to include and share it within our projects.

CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility.

✍ To sign, please post a new comment on this PR with exactly the following text: ✍

---

I have read and agree to the Contributor License Agreement

---

_{You can retrigger this bot by commenting recheck in this Pull Request. Posted by the CLA Assistant Lite bot.}

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 8d6b069b-ae57-45f9-974b-27747c289da5

📥 Commits

Reviewing files that changed from the base of the PR and between ff83be3 and 6b38b04.

📒 Files selected for processing (2)

• comfy_kitchen/backends/hip/ops/quantize_int8.hip
• tests/test_hip_wmma.py

Included review availability: Your plan provides up to 2 included reviews per hour; 1 remains after this review.

---

📝 Walkthrough

Walkthrough

Changes

HIP INT8 quantization

Layer / File(s)	Summary
Vectorized rowwise kernel comfy_kitchen/backends/hip/ops/quantize_int8.hip	Adds a templated uint4-based kernel for float, __half, and __bf16 inputs.
Dispatch and alignment validation comfy_kitchen/backends/hip/ops/quantize_int8.hip, tests/test_hip_wmma.py	Selects the vectorized kernel for aligned, compatible row widths and keeps the scalar fallback. Tests compare aligned and 16-byte-offset inputs across supported dtypes and widths.

Sequence Diagram(s)

sequenceDiagram
  participant RowwiseLauncher
  participant VectorizedKernel
  participant InputBuffer
  participant QuantizedOutput
  RowwiseLauncher->>InputBuffer: Check alignment and row width
  RowwiseLauncher->>VectorizedKernel: Launch vectorized kernel when supported
  VectorizedKernel->>InputBuffer: Load uint4 chunks
  VectorizedKernel->>VectorizedKernel: Compute absmax and quantize values
  VectorizedKernel->>QuantizedOutput: Write quantized values and scales

Loading

Suggested reviewers: 0xdeluxa, comfyanonymous, tvukovic-amd

Merge Risk: ⚪ Minimal · up to 6b38b

This PR adds a guarded vectorized INT8 quantization path while retaining the scalar fallback, with reported coverage for FP32, FP16, BF16, aligned, and offset-contiguous inputs. No actionable merge-blocking risk remains beyond normal checks and review.

🚥 Pre-merge checks | ✅ 2

✅ Passed checks (2 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

✨ Simplify code

• Create PR with simplified 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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[0xDELUXA]

Cross-checked this on RDNA4 (gfx1200), since the numbers here are all gfx1151. Correct and bit-identical, no regression, but no speedup either on this part.

Correctness

Built this branch and main separately and compared their outputs directly. q and scales come out bit-identical between the two builds on every case I tried: K=8 (exact vector width), K=9, K=33 and K=1023 (ragged, takes the fallback), K=2 (below the vector width), all of fp32/fp16/bf16, plus offset-contiguous inputs that trip the alignment guard into the scalar path. Max delta against eager stays at 1 LSB on q and 0 on scales, unchanged from before.

Tests: tests/test_hip_wmma.py, tests/test_int8.py and tests/test_qdq.py::TestQuantizeINT8 give 467 passed, 39 skipped, and the new test_quantize_int8_rowwise_offset_matches_aligned passes 6/6. (tests/test_qdq.py has 16 NVFP4 failures on AMD, unrelated to this PR.)

Performance on gfx1200: flat

Median of 5 interleaved rounds alternating the two builds process by process, wall-clock timed around the direct extension call with preallocated outputs. Only listing shapes whose run-to-run spread stayed under 10 percent; anything under about 150 us was dominated by launch overhead here and could not be separated from noise in either direction.

Shape (M x K)	dtype	main	this PR	speedup
8192 x 2048	bf16	327.2 us	323.2 us	1.012x
8192 x 2048	fp32	361.0 us	361.0 us	1.000x
4096 x 4096	bf16	291.5 us	294.1 us	0.991x
4096 x 4096	fp32	351.9 us	353.4 us	0.996x
18432 x 2048	bf16	694.3 us	694.2 us	1.000x
18432 x 2048	fp32	763.6 us	763.3 us	1.000x
4096 x 14336	bf16	1002.2 us	1003.3 us	0.999x
4096 x 14336	fp32	1209.3 us	1210.6 us	0.999x

As a noise gauge, 4096 x 1023 and 512 x 14335 take the scalar fallback in both builds and so run identical code; they came out at 1.004x and 1.011x. Every entry above sits inside that.

The generated code is what the patch intends, so this is not a missed-vectorization artifact: the gfx1200 ISA shows global_load_b128 on the read side, and the compiler already merges the 8 scalar int8 writes into a single global_store_b64 (global_store_b32 for fp32). The store side needs no hand-packing.

For what it is worth, two other things I tried on top of this branch are flat here as well, both verified bit-identical to main first: keeping the first pass's uint4 values in registers and dropping the second global read of x entirely (0.998x on 4096 x 14336 bf16, .private_segment_fixed_size confirmed 0 so nothing spilled), and replacing the 8-round __syncthreads tree reduction with the existing block_reduce_max<8> from hadamard.h (0.991x). So whatever bounds this kernel on RDNA4, it is not the load width, the redundant read pass or the reduction barriers, and I have not identified it.

None of that argues against the change: it is portable, guarded, bit-exact, worth 2x where you measured it, and costs nothing here.