lablup / lablup/mlxcel

fix(rocm): CUDA and ROCm hosts share one MTP policy cache key

Open
#1,887 0 comments 0 reactions 0 assignees View on GitHub
status:backlog type:bug
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

Phase 1 ROCm work under epic #1801. Found while validating PR #1883 (issue #1805) on the ROCm spike host: AMD Ryzen AI MAX+ 395 / Radeon 8060S, `gfx1151`, RDNA 3.5, wave32, ROCm 10.0.0, Debian 13, branch `feature/issue-1805-gpu-vendor`. Pre-existing defect, not a regression from that PR.

## Context

Found by inspection rather than by a failing test, so there is no error output to quote; nothing currently detects the collision.

`hardware_label()` at `src/server/batch/mtp_policy.rs:586-589` is `format!("{}-{}c", hw.silicon_gen, hw.gpu_core_count)`. On every non-Apple host `silicon_gen` is `AppleSiliconGen::Unknown` (`src/lib/mlxcel-core/src/hardware.rs:118` renders it as the literal `"Unknown"`) and `gpu_core_count` is 0, so a CUDA host and a ROCm host both produce the string `"Unknown-0c"`. That string is the `hardware` component of `PolicyKey::new(target_id, drafter_id, hardware_label(), block_size)` (`src/server/batch/mtp_policy.rs:1029`, struct at `src/server/batch/mtp_policy.rs:542-547`), and the key is what an MTP policy verdict is persisted under: `hint_file` names the file `{key.hash()}.json` (`src/server/batch/mtp_policy.rs:679-683`).

An NVIDIA GB10 and an AMD gfx1151 therefore share one cache entry. `load()` at `src/server/batch/mtp_policy.rs:690-703` guards against version, target, drafter, hardware, and block-size mismatch, but the hardware guard compares the same `"Unknown-0c"` on both hosts and passes, so a verdict profiled on one backend is loaded and applied on the other. The whole reason the label exists (issue #165) is that an MTP enable/decline verdict does not carry across hardware classes, and CUDA versus ROCm is a hardware class boundary at least as wide as M1 Max versus M1 Ultra, which the label was built to separate.

Why PR #1883, where this was found, did not fix it: the label is persisted inside the hint files and is part of the file-name hash, so any new spelling silently orphans every verdict already recorded under the old one. Issue #1805 constrains that PR to leave CUDA behavior unchanged, and invalidating a CUDA host's profile cache is a CUDA behavior change. That PR instead reports the vendor next to the label in the `GET /v1/internal/mtp-policy` response body as three new fields, `gpu_vendor`, `gpu_device`, and `gpu_architecture`, which fixes the reporting symptom without touching the key, and records the deferral in the doc comment at `src/server/batch/mtp_policy.rs:584-590` on `feature/issue-1805-gpu-vendor`.

## Scope

**In scope:** `hardware_label()` and `PolicyKey` in `src/server/batch/mtp_policy.rs`, the handling of hint files written under the old spelling, and the tests covering key construction.

**Out of scope:** the three reporting fields added by PR #1883, which stay as they are; this issue fixes the key, not the response body. The `HINT_VERSION` semantics for verdict changes (`src/server/batch/mtp_policy.rs:100-110`), which this change does not alter. Apple-host label behavior, which must not change.

## Proposed solution

Widen the label so the vendor distinguishes hosts. The data is available after PR #1883: `HardwareCapabilities` gains `vendor: GpuVendor` (`src/lib/mlxcel-core/src/hardware.rs:181` on the PR branch) and `device_architecture: Option` (`:218`), so the label can name the vendor and, where it is meaningful, the architecture. Land this after #1883 merges, since it depends on those fields.

On the old hint files, two options were considered.

(a) Silently ignore them. This needs no migration code at all: the label feeds `key.hash()`, which names the file, so a new spelling simply looks up a path that does not exist, and `load()` returns `None` and the host re-profiles once. The cost is the orphaned files left in the cache directory and one profiling window per pairing per host.

(b) Migrate them. This needs a rule for which vendor an old `"Unknown-0c"` entry belonged to, and that is not recoverable from the file: the stored `hardware` field is exactly the ambiguous string, and nothing else in the hint records a vendor. Any rule would be a guess, and guessing wrong reuses a cross-vendor verdict, which is the bug being fixed.

**Recommend (a).** Losing a profile costs one re-profiling window; migrating wrong costs the correctness this issue exists to restore. Decide whether to also sweep orphaned files, or leave them to accumulate, and state the choice in the code comment.

Keep the Apple-host spelling stable so existing Apple profiles survive: a label of the form `"M5-16c"` on Apple and something vendor-qualified such as `"cuda-Unknown-0c"` / `"rocm-gfx1151-Unknown-0c"` off Apple gives a correct split without invalidating the Apple cache. Confirm the exact shape against how the label is rendered in `PolicyKey::display()` (`src/server/batch/mtp_policy.rs:561-567`), which reaches logs and the stored hint body.

## Implementation plan

1. Land after PR #1883 merges, so `vendor` and `device_architecture` are available on `HardwareCapabilities`.
2. Rewrite `hardware_label()` (`src/server/batch/mtp_policy.rs:586`) to include the GPU vendor, and the device architecture where it is present, for non-Apple hosts. Preserve the current spelling exactly on Apple hosts so those profiles keep loading.
3. Decide and implement the old-file policy per the recommendation: no migration, `load()` already returns `None` for a path that no longer resolves. If orphan sweeping is wanted, add it as an explicit best-effort step rather than a load-path side effect.
4. Update the doc comment at `src/server/batch/mtp_policy.rs:580-590`, which currently documents the collapse and the deferral, to document the new spelling and the decision about old files.
5. Add a test that two different vendors produce different `PolicyKey`s, asserting on the key (or its hash) rather than on the label string, so the test pins the property that matters. Drive it by injecting the vendor rather than by reading the host, so it runs on every backend in CI.
6. Add a test that an Apple-host label is unchanged from the current format, so the Apple cache invalidation is caught if someone changes the shared path.
7. Check whether `PolicyKey::display()` output appears in any doc, dashboard, or log-parsing expectation that needs updating alongside the new spelling.

## Acceptance criteria

- [ ] A CUDA host and a ROCm host produce different `PolicyKey` values for the same target, drafter, and block size, and therefore different hint files.
- [ ] A test asserts that two different `GpuVendor` values yield different keys, and it runs on every backend rather than only on ROCm.
- [ ] Apple-host labels are byte-identical to the current format, so existing Apple profiles still load; a test pins this.
- [ ] Hint files written under the old `"Unknown-0c"` spelling are never loaded on a host of a different vendor.
- [ ] The old-file decision (ignore, not migrate) is documented in the code comment with its reasoning.
- [ ] The change is wired into the real key-construction path at `src/server/batch/mtp_policy.rs:1029`, not left as an unused helper.
- [ ] `GET /v1/internal/mtp-policy` still reports the three PR #1883 fields, now consistent with the widened label.

## Validation

```bash
cargo test --workspace --profile test-fast --features rocm -p mlxcel --lib server::batch::mtp_policy
cargo test --workspace --profile test-fast --features rocm
cargo clippy --workspace --all-targets --features rocm -- -D warnings
cargo fmt --all -- --check
```

Regression guard on a CUDA host:

```bash
cargo test --workspace --profile test-fast --features cuda -p mlxcel --lib server::batch::mtp_policy
```

Manual cross-host check: profile an MTP pairing on the CUDA host, copy the resulting hint directory to the ROCm host, start the server with the same target, drafter, and `--draft-block-size`, and confirm via `GET /v1/internal/mtp-policy` that the ROCm host re-profiles rather than adopting the CUDA verdict. A pass is a fresh profiling window on the second host and two distinct hint files.

## References

- Epic #1801 (AMD GPU (ROCm) backend on Linux via mlxcelverse), phase 1
- #1805 and PR #1883, where this was found and deferred, and which adds `vendor` / `device_architecture`
- #165, the regression that motivated the hardware label
- #1257, the `GET /v1/internal/mtp-policy` read interface

Contributor guide

Open the contributing guide

Research direction

Start with hardware_label(), PolicyKey, and hint-file load/key construction in src/server/batch/mtp_policy.rs, then inspect the HardwareCapabilities fields added by PR #1883. Run the named MTP policy tests with the ROCm feature and add coverage for vendor-separated keys and unchanged Apple labels. Done means the acceptance criteria pass, including old hint-file handling and the CUDA regression check.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
backend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.