Skip to content

Fix LoMa BF16 CUDA probing with CPU bookkeeping nodes - #4895

Closed
fringehush-gif wants to merge 2 commits into
colmap:mainfrom
fringehush-gif:codex/upstream-loma-bf16-probe
Closed

fringehush-gif wants to merge 2 commits into
colmap:mainfrom
fringehush-gif:codex/upstream-loma-bf16-probe

Conversation

@fringehush-gif

Copy link
Copy Markdown
Contributor

Inspect CUDA graph assignments so LoMa BF16 accepts supported CPU bookkeeping nodes while keeping numerical work on CUDA.

Changes

  • Allow selected CPU shape/index operators during CUDA capability probes; reject other CPU operators and require CUDA work.
  • Keep strict probing on older ONNX Runtime and CoreML, and FP32 fallback on pre-Ampere GPUs.

Tests

  • RTX 3060 / ORT 1.30: 11 tests pass. Real LoMa B FP32/BF16 self-matches: 64/64 and 128/128; default FP32 matches the base exactly.
  • Reproduced the original strict BF16 probe failure.

Compatibility

  • Windows 10 / GTX 1080 / ORT 1.30 CUDA 12: 11 tests pass without skips. GPU FP32 and BF16-to-FP32 fallback match base and CPU exactly at 64/128 features.
  • ORT 1.23.2 and 1.30 CPU SDKs: 9 tests each, 20 matching calls, exact pairs.

Build

  • Changed CUDA-enabled source units compile against both SDKs.
  • Full CLI links; --help exits 0.

Fixes #4787.

Totally vibe coded

Inspect CUDA graph assignments so LoMa BF16 accepts supported CPU bookkeeping nodes while keeping numerical work on CUDA.

### Changes

- Allow selected CPU shape/index operators during CUDA capability probes; reject other CPU operators and require CUDA work.
- Keep strict probing on older ONNX Runtime and CoreML, and FP32 fallback on pre-Ampere GPUs.

### Tests

- **RTX 3060 / ORT 1.30:** 11 tests pass. Real LoMa B FP32/BF16 self-matches: **64/64 and 128/128**; default FP32 matches the base exactly.
- Reproduced the original strict BF16 probe failure.

### Compatibility

- **Windows 10 / GTX 1080 / ORT 1.30 CUDA 12:** 11 tests pass without skips. GPU FP32 and BF16-to-FP32 fallback match base and CPU exactly at 64/128 features.
- ORT 1.23.2 and 1.30 CPU SDKs: **9 tests each, 20 matching calls**, exact pairs.

### Build

- Changed CUDA-enabled source units compile against both SDKs.
- Full CLI links; `--help` exits 0.

Fixes colmap#4787.

Totally vibe coded
fringehush-gif added a commit to fringehush-gif/colmap that referenced this pull request Oct 25, 2026
Inspect CUDA graph assignments so LoMa BF16 accepts supported CPU bookkeeping nodes while keeping numerical work on CUDA.

### Changes

- Allow selected CPU shape/index operators during CUDA capability probes; reject other CPU operators and require CUDA work.
- Keep strict probing on older ONNX Runtime and CoreML, and FP32 fallback on pre-Ampere GPUs.

### Tests

- **RTX 3060 / ORT 1.30:** 11 tests pass. Real LoMa B FP32/BF16 self-matches: **64/64 and 128/128**; default FP32 matches the base exactly.
- Reproduced the original strict BF16 probe failure.

### Compatibility

- **Windows 10 / GTX 1080 / ORT 1.30 CUDA 12:** 11 tests pass without skips. GPU FP32 and BF16-to-FP32 fallback match base and CPU exactly at 64/128 features.
- ORT 1.23.2 and 1.30 CPU SDKs: **9 tests each, 20 matching calls**, exact pairs.

### Build

- Changed CUDA-enabled source units compile against both SDKs.
- Full CLI links; `--help` exits 0.

Upstream review: colmap#4895.

Totally vibe coded
- Fix CMake plain/keyword signature mixing for colmap_util_onnx_test, which broke configure on CUDA builds
- Link colmap_feature against colmap_util_cuda (loma.cc needs the symbol)
- Extract generic CudaDeviceSupportsBf16() helper to util/cuda and replace the LoMa-internal BF16 device check with a thin wrapper; the helper fails closed (false) on invalid input or query failure
- Document that the CUDA probe inspects operator types only, and log allowed CPU nodes at VLOG(2)
- Add cuda_test for the helper; document the onnx_test probe models
@sarlinpe

sarlinpe commented Oct 29, 2026 •

Copy link
Copy Markdown
Member

Thanks a lot for tracking this down and for the thorough testing! I opened #4924 as an alternative fix for #4787 and would suggest superseding this PR with it. The reasoning:

The root cause is the strict probe, not specific operators. ONNX Runtime's CPU placement (GetCpuPreferredNodes) is driven by data flow rather than operator type: the CUDA Shape/Size kernels output to host memory, and their int64 consumers are moved to the CPU EP, whatever they are. Tensors in bf16/fp16 are never moved. So an allowlist of operator types is brittle:

  • The set of CPU nodes depends on the graph and the ORT version. For example, ORT runs Reshape of the FlexGlue matchers (Add FlexGlue feature matcher #4908) on the CPU, which this allowlist rejects. Other LoMa variants (B128/R/L/G) or ORT versions may hit operators outside the list and silently fall back to fp32 again.
  • It needs GetEpGraphAssignmentInfo and thus ORT >= 1.24; older ORT stays broken.

Dropping disable_cpu_ep_fallback barely weakens the check. The bf16 tensors are never moved to the CPU, and the CPU EP has almost no bf16 compute kernels, so a bf16 op without a CUDA kernel still fails at session creation and triggers the fp32 fallback. The fp32 models already run with the same CPU shape nodes.

The compute capability check is better done before loading. Here it runs after the bf16 session (and, for extraction, the detector) has been downloaded and built, and throws a synthetic Ort::Exception. #4924 checks the device up front, so GPUs without bf16 support load fp32 directly.

#4924 thus:

  • only uses bf16 with the CUDA EP on GPUs with compute capability >= 8.0 (checked with a new CudaDeviceSupportsBf16() before loading anything),
  • loads the bf16 models as regular sessions with a fallback to fp32 on initialization failures, which removes ONNXModel's capability probes altogether (LoMa was their only user),
  • follows the same pattern as the fp16 FlexGlue matcher in Add FlexGlue feature matcher #4908.

It is +130/−95 (net +35 lines, vs. +255 here), without hand-encoded ONNX test models. Note that the ClangTidy job also fails here (performance-inefficient-string-concatenation at onnx.cc:311).

Since you can reproduce the issue on both an RTX 3060 and a GTX 1080, it would be great if you could check that #4924 runs the bf16 matcher on the 3060 and falls back to fp32 on the 1080 — I could not run the ORT 1.30 CUDA EP locally.

ahojnnes pushed a commit that referenced this pull request Oct 30, 2026
)

Fixes #4787. Supersedes #4895.

## Problem
With `--LomaMatching.use_bf16 1` on CUDA, the bf16 LoMa matcher always
fails to initialize and falls back to fp32:

```
Failed to initialize the bf16 LoMa matcher (This session contains graph nodes that are assigned to the default CPU EP, but fallback to CPU EP has been explicitly disabled by the user.); falling back to fp32
```

The bf16 models are loaded as capability probes, which set
`session.disable_cpu_ep_fallback`. But ONNX Runtime deliberately runs
small shape computations on the CPU, whatever the precision of the
model: its CUDA `Shape`/`Size` kernels write their output to host
memory, and `GetCpuPreferredNodes` then moves their int64 consumers
(`Gather`, `Unsqueeze`, `Concat`, `Cast`, `Reshape`, ...) to the CPU EP.
Tensors in bf16 or fp16 are never moved. The dynamic-shape matcher
graphs contain such nodes, so the probe always rejects them. The fp32
models are not loaded as probes, so they run with the same CPU nodes.

## Fix
Load the bf16 models like any other model, and decide up front whether
to use them:
- The bf16 models are only used with the CUDA execution provider, on
GPUs with compute capability 8.0 (Ampere) or newer. Otherwise, the fp32
models are loaded directly, without first downloading and loading the
bf16 models.
- On older GPUs, creating a bf16 session can succeed while inference
fails, which the fallback on initialization failures would not catch.
- On the CPU, `use_bf16` is now a no-op instead of trying to load the
bf16 models.
- The bf16 models are loaded as regular `ONNXModel`s, with a fallback to
fp32 if initialization throws an `Ort::Exception`. This removes
`ONNXModel`'s `is_capability_probe`, whose only user was LoMa.
- New `CudaDeviceSupportsBf16()` queries the device. It returns false if
the device cannot be queried, and on HIP builds.
- New `SelectONNXCudaDevice()` resolves `gpu_index` to the device that
`ONNXModel` runs on, shared with the bf16 check.
- Updated the option docs (C++, pycolmap, `doc/features.rst`).

This is the same pattern as the fp16 FlexGlue matcher in #4908: reduced
precision only on CUDA, a regular session, and a fallback to fp32 on
initialization failures.

### Comparison with #4895
#4895 keeps the strict probe and instead inspects the EP graph
assignment (ORT >= 1.24) against an allowlist of CPU operators. This PR
avoids:
- the allowlist, which is brittle: which nodes ORT moves to the CPU
depends on data flow, not on the operator type. For example, FlexGlue's
matchers run `Reshape` on the CPU, which the allowlist rejects. Any
rejected node silently falls back to fp32 again.
- the `ORT_API_VERSION >= 24` gating; the fix also works with older ORT.
- downloading and loading the bf16 models on GPUs without bf16 support
before rejecting them.
- ~80 lines of hand-encoded ONNX test models.

Removing the strict probe barely weakens the check: the bf16 tensors
themselves are never moved to the CPU, and the CPU EP has almost no bf16
compute kernels, so a bf16 op without a CUDA kernel still fails at
initialization and falls back to fp32.

## Tests
- `util/cuda_test`: `CudaDeviceSupportsBf16` returns false for invalid
devices and, with a GPU, leaves no CUDA error.
- `util/onnx_test`: `SelectONNXCudaDevice`, and an invalid model throws
the `Ort::Exception` that the LoMa fallback catches.
- `feature/loma_test`, `util/onnx_test`, `util/cuda_test` (also with
`CUDA_VISIBLE_DEVICES=""`) pass locally (CUDA + ONNX build).
- Not yet run end-to-end with the CUDA EP: ORT 1.30's CUDA EP does not
load on my machine (older system CUDA runtime). The bf16 LoMa-B matcher
with CPU shape nodes was validated in #4895 (64/64 and 128/128
self-matches on an RTX 3060).

**Stack:** the FlexGlue stack builds on this PR: #4907, #4921, #4908,
#4915. #4908 selects its fp16 matcher the same way, with a new
`CudaDeviceHasTensorCores()` (compute capability ≥ 7.0) next to
`CudaDeviceSupportsBf16()`. #4915 uses `SelectONNXCudaDevice()` for the
GPU feature cache.
@fringehush-gif

fringehush-gif commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor Author

That explanation went a bit over my head but I think I got the general gist. I appreciate the explanation. I'm a total novice and this is entirely vibe coded like I mentioned, but anyway.

I noticed it's been merged into main since but I will test it and post a follow up in #4924.

@sarlinpe sarlinpe closed this Oct 31, 2026
fringehush-gif added a commit to fringehush-gif/colmap that referenced this pull request Nov 3, 2026
Inspect CUDA graph assignments so LoMa BF16 accepts supported CPU bookkeeping nodes while keeping numerical work on CUDA.

### Changes

- Allow selected CPU shape/index operators during CUDA capability probes; reject other CPU operators and require CUDA work.
- Keep strict probing on older ONNX Runtime and CoreML, and FP32 fallback on pre-Ampere GPUs.

### Tests

- **RTX 3060 / ORT 1.30:** 11 tests pass. Real LoMa B FP32/BF16 self-matches: **64/64 and 128/128**; default FP32 matches the base exactly.
- Reproduced the original strict BF16 probe failure.

### Compatibility

- **Windows 10 / GTX 1080 / ORT 1.30 CUDA 12:** 11 tests pass without skips. GPU FP32 and BF16-to-FP32 fallback match base and CPU exactly at 64/128 features.
- ORT 1.23.2 and 1.30 CPU SDKs: **9 tests each, 20 matching calls**, exact pairs.

### Build

- Changed CUDA-enabled source units compile against both SDKs.
- Full CLI links; `--help` exits 0.

Upstream review: colmap#4895.

Totally vibe coded
@fringehush-gif
fringehush-gif deleted the codex/upstream-loma-bf16-probe branch November 3, 2026 19:52
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.

[LoMa Matcher] [4.2.1] Failed to initialize the bf16 LoMa matcher

3 participants