Pull requests / #2355

#2355 pf_wmma: set the tiles' shared-memory attribute on each card, so the second card of a layer split can launch them

open · @Warabii1248 · 0 comentários · No GitHub

BenchmarksSetup & installServer & APIAMD / HIPNVIDIA / CUDAModels & quantsSecurityDocumentationLinux

Descrição

## Title

pf_wmma: set the tiles' shared-memory attribute on each card, so the second card of a layer split can launch them

Issue: stacked on #1771 (which carries #1401's FP16 prompt commits). This PR's changes are the last two commits, `d5fdec0` and `2cec5c3`. #1771 no longer merges cleanly into main at `ed4773e` (a conflict in `src/prefill/prefill.cpp`); this branch carries that conflict, and its own two commits touch only `src/kernels/cuda/iq_kernels.cu` and `src/kernels/pf_wmma_parity.cpp`.

## Summary

`gemm_iq_f16_grouped` sets the over-48 KiB dynamic shared-memory attribute of its three tile kernels once per process (`static bool attr`).
`cudaFuncSetAttribute` applies to the current device only, so on a layer split only the card that launched the tiles first has it.
The next card's launch fails, `gemm_iq_f16_grouped` returns false, and the prompt stops with `prefill: the WMMA expert tiles could not launch`
(the engine exits with code 1 and serve restarts it). The measurements in #1401 / #1771 that we found are on one card, where this path is not reached.

The fix keeps the flag per device (the same per-device `[64]` shape as `exp_mode()` in the same file).
`pf_wmma_parity` now runs its cases on every sm_70+ card in turn within one process, as a layer split does; before the fix every case fails on the second card.

## What changed

- `src/kernels/cuda/iq_kernels.cu`: set the attribute once per device instead of once per process (`static bool attr[64]`, `cudaGetDevice`).
- `src/kernels/pf_wmma_parity.cpp`: run the cases on each sm_70+ device in turn (`cudaSetDevice`); the summary line gives the device count.
- Other backends: the `iq_kernels.cu` change is inside the `#if !defined(__HIPCC__)` block of `gemm_iq_f16_grouped` (both HIP builds compile the stub
  that returns false, unchanged), and the SYCL build has no `gemm_iq_f16_grouped`.
  HIP builds, without a GPU in ROCm 7.14 containers: `-DSTRATA_HIP_GFX906=ON -DCMAKE_HIP_ARCHITECTURES=gfx906` (the `mixa3607/rocm-gfx906:7.14-complete`
  image docs/AMD_HIP.md uses; `strata` and `pf_wmma_parity`) and `-DSTRATA_ENABLE_HIP=ON -DSTRATA_ENABLE_CUDA=OFF -DCMAKE_HIP_ARCHITECTURES="gfx1100;gfx1201"` (`rocm/dev-ubuntu-24.04:7.14.1-full`;
  `strata`) both compile with 0 errors. Not run (no AMD card here). SYCL not built: it does not compile the changed files (it has its own `iq_kernels` copy and no `pf_wmma_parity`).

The branch is v0.1.43 (`7695859`) with #1771's five commits cherry-picked onto it (#1771 itself is on v0.1.41, `fb58e0d`; the patches are
unchanged, authors kept, `cherry picked from` trailers) and the two commits above on top. No commit in `7695859..a0116b2` (`main` when this branch was made) touches `iq_kernels.cu` or `pf_wmma_parity.cpp`; since `a0116b2`, main has changed `iq_kernels.cu` (at `ed4773e` that file merges without a text conflict; we have not built the branch on `ed4773e`).

Effect on default behavior: with one card, that card's attribute is set at the same point as before. On a layer split of two or more cards, the prompt that stopped now runs.

## Checks (Tesla V100-PCIE-32GB + Tesla PG500-216 (sm_70, 32 GB), layer split K=24, CUDA 12.9, Linux, `-DSTRATA_EXPERIMENTAL_SM60=ON`)

How we checked: the parity test before and after the fix, then the engine on the layer split before the fix and, after it, with and without the tiles.

- `pf_wmma_parity` before the fix: device 0 8 / 8 ok, device 1 8 / 8 "not launched" (`8 failures on 2 devices`, exit code 1).
  After: `0 failures on 2 devices`, exit code 0. The second card's errors equal the first card's (e.g. IQ4_NL 2560x640 max |err| 0.000162 of 65.4).
- Engine: the GSQ-RCO IQ3_XXS model and pack that `setup.py --model IQ3_XXS` installs (`ISTA-DASLab/Qwen3.8-Flash-Next-GSQ-RCO-GGUF`, nothing else changed),
  `--mmap-experts`, 256K context, int8 KV, `STRATA_SM70_TABLE=1`, every expert in VRAM, 170 W power limit.
  - Before the fix, default (`STRATA_PF_WMMA` on): the first prompt stops with `prefill: the WMMA expert tiles could not launch` (0 tokens read; the engine exits with code 1 and serve restarts it).
  - After the fix: `STRATA_PF_WMMA=0` against the default, alternating, a fresh engine start per run. After each start, fresh prompts of 26.7K
    (the first request), 52K (the second) and 104K tokens (the third), nothing reused. Prompt tok/s = prompt tokens / time to the first content token, measured on the client:

    | prompt tokens | `STRATA_PF_WMMA=0` (3 runs) | default (3 runs) | ratio of medians |
    |---|---|---|---:|
    | 26,689-26,690 | 1,670.6 / 1,667.9 / 1,669.8 | 1,765.8 / 1,766.8 / 1,765.3 | +5.7% |
    | 52,015-52,016 | 2,064.9 / 2,063.1 / 2,065.2 | 2,201.4 / 2,200.1 / 2,204.3 | +6.6% |
    | 103,865-103,866 | 2,138.8 / 2,136.2 / 2,137.5 | 2,275.3 / 2,273.3 / 2,271.5 | +6.4% |

    Every default run is above every `=0` run. The third default run is a re-run: in the first attempt the start-up PLE read check measured 9,325 rows/s
    (the other five starts 33,428-37,238), and 0.1.43's automatic switch moved that run to `--ple-io mmap` (385 / 1,115 / 1,075 tok/s; we did not find the cause).
    That is the PLE read mode, not the tiles, so we re-ran it with `STRATA_PLE_PROBE=0` (`--ple-io direct`, what the check chose on every other start).
    (In 0.1.43.1 the switch only warns, #2040.)

## Extra Notes

- Disclosure: the change and this text were prepared with an AI assistant (Claude); the numbers are from runs on the machine above.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

No site

Links install, modelos, releases.