#2706 [Bug]: STRATA_HC_UPMIX on a gfx11 + gfx12 layer split corrupts prefill on the gfx12 stage (gr_upmix caches its arch check from the first device)

Descrição

## Summary

On a mixed-architecture layer split (gfx1151 first, gfx1201 second), `STRATA_HC_UPMIX=1` gives wrong prefill results
on every layer that runs on the gfx1201 stage. `gr_upmix()` decides "is this gfx11?" once, in a function-local
`static`, from whichever device calls it first. After that it launches `gr_upmix_kernel` on the gfx1201 device too.
That device's code object was compiled with the non-gfx11 branch of `um_wmma()`, which is `return c;`, so the WMMA
step does nothing.

Nothing crashes and nothing is logged. Generation is still fluent, so this is easy to miss. It shows up as worse
answers: on BFCL parallel function calling, the model emits only one call where it should emit several. The rate
drops from 28/38 to 3/38 on the affected prompts.

Still present on v0.1.44.1 (`296b59d2`, origin/main as of 2026-10-15). Reproduced on v0.1.43.1 (`a0116b27`) and
v0.1.44 (`ed4773eb`).

## Setup

| | |
|---|---|
| Strata | v0.1.43.1 `a0116b27` (service) and v0.1.44 `ed4773eb` (repro). `kernels.cu` / `prefill.cpp` on this path are unchanged in v0.1.44.1 |
| Build | one binary for both: `-DCMAKE_HIP_ARCHITECTURES="gfx1151;gfx1201" -DSTRATA_W_LB=8` |
| ROCm | TheRock 7.14.1 multiarch dist, Linux 7.0.0-38-generic |
| GPU 0 (first stage) | Radeon 8060S, gfx1151 (Strix Halo iGPU), layers 0–29 |
| GPU 1 (last stage) | Radeon AI PRO R9700, gfx1201, layers 30–47 (+ head, MTP) |
| Model | Qwen3.8-Flash-Next, Unsloth UD-IQ4_XS GGUF (native pack) |
| Split | server config `"gpu": [<gfx1151>, <gfx1201>], "layer_split": "30"`, `--prefill 4096 --trim-stage-weights` |
| Env | `STRATA_PF_FUSED=1 STRATA_PF_GEMM=1 STRATA_HC_UPMIX=1 STRATA_PA_FAST=1 STRATA_HIP_WMMA=1 STRATA_SELECT_WMMA=1 STRATA_PF_FUSED_KQ=1 STRATA_PREFILL_STREAM_MIN=64` |

Note: the build carries one small local patch elsewhere in the prefill path, which the gfx1201 stage needs in order
to start. It does not touch `gr_upmix` or its call site.

Engine command line (paths shortened):

```
strata --serve --pack packs/unsloth-ud-iq4_xs --native Qwen3.8-Flash-Next-UD-IQ4_XS-00001-of-00003.gguf \
  --expert-profile data/expert-profile.bin --expert-cache auto --spec 4 --spec-min-p 0.5 --mtp mtp/rt \
  --max-context 262144 --kv int8 --kv-resident 32768 --mmap-experts --mtp-window 8192 --no-prefill-borrow \
  --prefill 4096 --conversation-cache-mib 24576 --conversation-cache-slots 24 --trim-stage-weights \
  --batch-groups 1 --batch-mtp --vision --layer-split 30 --batch 4
```

## Root cause

`src/prefill/kernels.cu` (origin/main `296b59d2`):

```cpp
// 1745
__device__ __forceinline__ um_f8 um_wmma(um_b16 a, um_b16 b, um_f8 c) {
#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__)
    return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a, b, c);
#else
    (void) a; (void) b; return c;   // not a gfx11 device pass: never launched (runtime gate)
#endif
}

// 1863
bool gr_upmix(...) {
#if defined(__HIPCC__)
    static const bool gfx11 = [] {                       // <- evaluated once per process,
        int dev = 0;                                     //    on whichever device calls first
        cudaDeviceProp p;
        if (cudaGetDevice(&dev) != cudaSuccess || cudaGetDeviceProperties(&p, dev) != cudaSuccess) return false;
        return strata::kernels::gfx_arch_is_gfx11_wmma(p.gcnArchName);
    }();
    if (!gfx11 || T <= 0) return false;
    gr_upmix_kernel<<<...>>>(...);                       // <- also launched on the gfx1201 stage
```

The call site in `src/prefill/prefill.cpp:2641` has no per-device check either:
`if (hc_upmix() && T >= pf_switch_min_t() && !gr_unfused() && ...)`.

With the gfx1151 stage first, `gfx11` is cached as `true`. The "runtime gate" that the `um_wmma` comment relies on
then lets the gfx1201 launch through. If the gfx12 device came first, the cache would hold `false` and the fused path
would be skipped on both GPUs. That case is slower but correct; I didn't test it.

Other device-specific checks in the tree already cache per device: `strata_wmma_gfx11_device()` in
`wmma_gemm.cu:45` (`std::atomic<int> cache[64]`), `sel_gfx12_device()` in `qsa_select.cu:447`, and the `arch[dev]`
lookup in `qsa_prompt_attn.cu`. I grepped origin/main for other function-static arch or WMMA gates. The others I found
are env-only (`prefill.cpp:297`, `qsa_select.cu:1128`, `gemm.cu:632/734`). The same `gr_upmix` pattern is mirrored in
`sycl/src/prefill/kernels.dp.cpp:2345`.

## Reproduction

1. Run any gfx11 + gfx12 layer split with the gfx11 device first (here: `--layer-split 30`, 8060S first) and
   `STRATA_HC_UPMIX=1 STRATA_HC_UPMIX_CHECK=400`.
2. Send any chat request. Any prompt length works, since `STRATA_PF_SWITCH_MIN_T` defaults to 0.
3. Run `grep "STRATA_HC_UPMIX_CHECK" engine.log`.

The check compares the fused up-projection against the default pair (`bf16_proj` + `gr_mix_r`). It is ~1e-7 on
every gfx1151 layer and ~0.42–0.74 on every gfx1201 layer. Here is v0.1.44 with three prompts (T = 348, 362, 413),
two halves per prompt, so six checks per layer:

| layers (device) | mixed rel RMS, min–max | max \|diff\| |
|---|---|---|
| 0–29 (gfx1151) | 7.8e-08 – 3.2e-07 | ≤ 2.3e-05 |
| 30 (gfx1201) | 0.422 – 0.456 | 7.0 |
| 31 | 0.417 – 0.495 | 12.0 |
| 32 | 0.476 – 0.483 | 9.7 |
| 33 | 0.462 – 0.473 | 7.1 |
| 34 | 0.464 – 0.479 | 11.2 |
| 35 | 0.470 – 0.508 | 10.7 |
| 36 | 0.472 – 0.495 | 9.2 |
| 37 | 0.472 – 0.520 | 9.5 |
| 38 | 0.470 – 0.497 | 16.7 |
| 39 | 0.475 – 0.484 | 7.1 |
| 40 | 0.484 – 0.493 | 12.7 |
| 41 | 0.462 – 0.495 | 10.6 |
| 42 | 0.472 – 0.490 | 5.8 |
| 43 | 0.487 – 0.743 | 4.6 |
| 44 | 0.458 – 0.597 | 12.3 |
| 45 | 0.449 – 0.468 | 6.7 |
| 46 | 0.431 – 0.452 | 6.0 |
| 47 | 0.442 – 0.454 | 5.4 |

Sample lines:

```
strata: STRATA_HC_UPMIX_CHECK layer 0 half 0 T 413: mixed rel RMS 1.642e-07, max |diff| 5.960e-06 (upmix)
strata: STRATA_HC_UPMIX_CHECK layer 30 half 0 T 413: mixed rel RMS 4.217e-01, max |diff| 6.042e+00 (upmix)
strata: STRATA_HC_UPMIX_CHECK layer 43 half 0 T 413: mixed rel RMS 7.282e-01, max |diff| 4.269e+00 (upmix)
```

v0.1.43.1 gives the same picture: layers 0–29 have max rel RMS 8.7e-08 to 3.2e-07, and layers 30–47 range from 0.45
to 0.74.

## Effect on output

I used 38 BFCL v4 `parallel` / `parallel_multiple` prompts, sent as the exact BFCL OpenAI-FC request (tools, the
"Python 3 syntax" hint, temperature 0.001, 4 concurrent). The items were picked because this setup had failed them.
The metric is how many responses contain 2 or more tool calls; the correct answers all need 2 or more. Every row
changes one thing from the first row.

| arm | v0.1.43.1 | v0.1.44 |
|---|---|---|
| split, env as above (HC_UPMIX on) | 3/38 | 3/38 |
| same, **without `STRATA_HC_UPMIX`** | **28/38** | **28/38** |
| 8060S alone (no split), env as above | 28/38 | – |
| split, `--kv` fp16 instead of int8 | 12/38 | – |
| split, all 8 env switches off | 24/38 | – |

Removing one env switch at a time on the split (v0.1.43.1):

| switch removed | PF_FUSED | PF_GEMM | **HC_UPMIX** | PA_FAST | HIP_WMMA | SELECT_WMMA | PF_FUSED_KQ | PREFILL_STREAM_MIN |
|---|---|---|---|---|---|---|---|---|
| 2+ calls | 16 | 11 | **28** | 12 | 13 | 3 | 16 | 17 |

Only `HC_UPMIX` brings the split back to the single-GPU result. The other rows move around (3–17), but none comes
close. I haven't traced why each one shifts the result.

Temperature and concurrency are not factors. With all four combinations of 1 or 4 workers and temperature 0 or 0.001
(hint removed), the split with HC_UPMIX scored 18/38 every time.

Prefill speed without `STRATA_HC_UPMIX` on the split: 38 prompts (17.5K tokens), cold, read in 32.4 s, against
31.5–33.9 s for the other single-switch arms. I measured no cost on short prompts.

Benchmarks on the same service (UD-IQ4_XS, no thinking), before and after removing only `STRATA_HC_UPMIX`:

| benchmark | with HC_UPMIX | without |
|---|---|---|
| BFCL `parallel` (first 100) | 57 | 65 |
| BFCL `parallel_multiple` (first 100) | 55 | 70 |
| BFCL `irrelevance` (first 100) | 65 | 74 |
| BFCL 5 single-turn categories, mean | 70.2 | 77.0 |
| TruthfulQA MC (first 300) | 85.0 | 90.0 |

The gap is how this was found: another engine running the same GGUF on the 8060S alone scored clearly higher on the
same subsets.

## Suggested fix

Cache per device, the way `strata_wmma_gfx11_device()` already does. For example:

```cpp
bool gr_upmix(...) {
#if defined(__HIPCC__)
    static std::atomic<int> cache[64];   // 0 unknown, 1 gfx11, 2 not
    int dev = 0;
    if (cudaGetDevice(&dev) != cudaSuccess || dev < 0 || dev >= 64) return false;
    int c = cache[dev].load(std::memory_order_acquire);
    if (c == 0) {
        cudaDeviceProp p;
        c = cudaGetDeviceProperties(&p, dev) == cudaSuccess && strata::kernels::gfx_arch_is_gfx11_wmma(p.gcnArchName) ? 1 : 2;
        cache[dev].store(c, std::memory_order_release);
    }
    if (c != 1 || T <= 0) return false;
    ...
```

Alternatively, have the call site in `prefill.cpp` ask the stage's device. As a defensive extra, the non-gfx11 branch
of `um_wmma` could trap (`__builtin_trap()`) instead of `return c`. Then a launch that slips past the gate fails loudly
instead of computing wrong numbers.

Workaround: unset `STRATA_HC_UPMIX` on mixed-arch splits.