Skip to content

Remove redundant CUB barriers from blockwise dequantization - #2053

Open
heiheiha798 wants to merge 2 commits into
bitsandbytes-foundation:mainfrom
heiheiha798:b300/dequantize-blockwise-cub-barriers
Open

Remove redundant CUB barriers from blockwise dequantization#2053
heiheiha798 wants to merge 2 commits into
bitsandbytes-foundation:mainfrom
heiheiha798:b300/dequantize-blockwise-cub-barriers

Conversation

@heiheiha798

Copy link
Copy Markdown

Summary

Remove the two external CTA barriers immediately before the one-shot CUB
BlockLoad and BlockStore in kDequantizeBlockwise.

The two collectives already use distinct, non-aliased shared TempStorage
objects. The launch geometry gives each CTA one loop iteration, so neither
storage object is reused after its collective. This leaves CUB's internal
synchronization intact and does not change the loop, grid, launch geometry,
decode arithmetic, formats, dispatch, streams, or public API.

Tracking: fork issue #10
and engineering draft PR #11.

Changes

  • csrc/kernels.cu: delete exactly the two redundant external barriers.
  • tests/test_ops.py: cover full/tail General8bit, NF4, and FP4 tiles for all
    output dtypes, repeated allocating and .out calls, and unchanged inputs.
  • benchmarking/dequantize_blockwise_barriers.py: build two Git revisions in
    isolated source trees and emit complete interleaved CUDA-event samples,
    medians, p10/p90, round ratios, and effective bandwidth as JSONL.

The production CUDA diff is two deletions. There is no architecture guard
because the storage-lifetime property is not SM103-specific.

B300 Results

Primary evidence is from an official CUDA 13 x64 compatible build with targets
75;80;86;89;90;100;120, executed on one NVIDIA B300 SXM6 AC (CC 10.3,
148 SMs) in Slurm job 4974. Software: driver 580.126.09, CUDA/nvcc 13.0.88,
PyTorch 2.13.0+cu130, Python 3.12.3. The isolated libraries were built from
baseline 95f9af309d4d5793847169c39288dcd3fcbdf564 and exact candidate
1f033cccf1a47e6cf054c7083e06ead723c559b7; the test process asserted that it
loaded the temporary candidate library rather than a checkout artifact.

Each cell used 20 warmups per round, seven interleaved CUDA-event rounds, and
size-scaled repetitions (4-100 per round). Effective bandwidth counts packed
input, absmax, the General8 map where applicable, and output bytes.

Path / shape dtype baseline candidate speedup candidate effective GB/s
General8 4096x4096, block 256 FP16 0.028512 ms 0.026464 ms 1.0774x 1911.8
General8 7168x7168, block 256 FP32 0.077536 ms 0.073568 ms 1.0539x 3502.9
NF4 4096x14336, block 64 FP32 0.059200 ms 0.055056 ms 1.0753x 4866.2
NF4 8192x8192, block 64 FP32 0.067264 ms 0.061424 ms 1.0951x 4984.8
Nested NF4 14336x4096 FP16 0.068320 ms 0.066240 ms 1.0314x n/a

Across the exact-head matrix:

  • 11 realistic direct cells improved by at least 5%, with the stricter
    max-minus-min round-ratio spread below the measured effect.
  • Four nested NF4 controls improved by at least 3% with the same stability
    check.
  • No direct cell regressed by more than 5%; the minimum observed ratio was
    0.9820x.
  • All 150 direct cells and 16 user-control cells were raw-bit equal.
  • Build time was 59.73 s baseline versus 59.30 s candidate; both libraries
    were 3,922,064 bytes.

The broader pre-commit calibration in job 4963 exercised 213 direct cells,
198 allocating/.out public correctness cells, three non-contiguous controls,
packed sweeps, custom maps, extreme absmax values, and nested/backward paths.
It also found zero raw-bit differences and no greater-than-5% regression.

The original raw-log paths and full engineering record remain listed in the
fork draft PR.

Reproduction

Run inside a single-GPU Slurm allocation from a clean candidate checkout. The
build root must not already exist.

python benchmarking/dequantize_blockwise_barriers.py \
  --output /tmp/dequant-barriers.jsonl \
  --build-root "${SLURM_TMPDIR}/dequant-barrier-build" \
  --baseline-ref 95f9af309d4d5793847169c39288dcd3fcbdf564 \
  --candidate-ref 1f033cccf1a47e6cf054c7083e06ead723c559b7 \
  --compute-capability '75;80;86;89;90;100;120' \
  --warmup 20 \
  --repetitions 100 \
  --rounds 7 \
  --include-user-controls \
  --require-sm103

Focused synchronization checks can be reproduced against the candidate build
with the permanent full/tail test:

compute-sanitizer --tool racecheck --error-exitcode 99 \
  --kernel-name kernel_substring=kDequantizeBlockwise \
  python -m pytest -q tests/test_ops.py \
  -k dequantize_blockwise_repeated_full_and_tail

compute-sanitizer --tool synccheck --error-exitcode 99 \
  --kernel-name kernel_substring=kDequantizeBlockwise \
  python -m pytest -q tests/test_ops.py \
  -k dequantize_blockwise_repeated_full_and_tail

Validation

Exact-head job 4974, using the path-asserted temporary candidate library:

  • tests/test_ops.py -k 'dequantize_blockwise or dequantize_4bit':
    99 passed, 306 deselected.
  • tests/test_functional.py -k 'dynamic_blockwise_quantization or test_4bit_quant or test_4bit_compressed_stats':
    165 passed, 1408 deselected.
  • tests/test_autograd.py -k test_matmul_4bit:
    384 passed, 384 deselected.
  • compute-sanitizer --tool racecheck: 0 hazards, 0 errors, 0 warnings.
  • compute-sanitizer --tool synccheck: 0 errors.
  • Official-compatible CUDA Release configure/build: passed for all requested
    targets.
  • pre-commit run --all-files: all hooks passed on the login CPU.
  • CPU-only collection of tests/test_ops.py: 405 tests collected without an
    import/collection failure; the 18 CUDA synchronization cases skipped as
    expected.

The contribution preflight also reran pre-commit run --all-files against the
exact submitted head; all hooks passed without changing the worktree.

Limits

  • Measurements and sanitizer execution are from one physical B300 and the
    official-compatible CUDA 13 image path. No native-SM103 performance result
    is substituted for that deployed-code evidence.
  • Neither the login environment nor the B300 Slurm node provides ROCm or
    hipcc. The requested HIP/hipCUB compile check and HIP runtime validation
    were therefore not run. This PR makes no HIP behavior claim; downstream HIP
    CI remains required before merge.
  • Other CUDA architectures were compiled through the official target list but
    were not benchmarked on hardware. The change is architecture-independent
    and does not add or alter runtime dispatch.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant