diff --git a/.agents/specs/rocm-qwen35-08b-cpu-gfx1100-numerics.md b/.agents/specs/rocm-qwen35-08b-cpu-gfx1100-numerics.md new file mode 100644 index 000000000..c4b7ddff8 --- /dev/null +++ b/.agents/specs/rocm-qwen35-08b-cpu-gfx1100-numerics.md @@ -0,0 +1,496 @@ +# Qwen3.5-0.8B CPU and gfx1100 numerical characterization + +Row: `BACKEND-ROCM`. + +Issue: [#2773](https://github.com/mudler/vllm.cpp/issues/2773), which supersedes +issue #1588 for implementation traceability. + +Branch: `row/BACKEND-ROCM-NUMERICS-1588`. + +Active primary oracle pin: vLLM +`e126687a9a828d513c01a07cd69f025f27d63280`. + +## Now + +`PENDING`. This repaired spec is the only completed deliverable in this commit. +Implementation cannot start until a fresh reviewer passes this immutable spec. +Results cannot be accepted until both correctness prerequisites pass: + +1. Issue #2772 and PR #2856 own the default-on ROCm `wvSplitK` sacred-anchor + repair. The reviewed old base failed prompt 10/token 10 (`369` instead of + `488`), while `VT_ROCM_SKINNY=0` passed. The operator reran PR #2856 at + `f06619e4c213e3de28359ee10995e682e8c06932`: CPU mode 2/2, prerequisites + 77/77/77/86, `wvSplitK` 79796/79796, and sacred 137/137 with 15 strict, one + tied, maximum gap zero, and zero divergence. PR #2856 is still open and the + contributor lacks merge authority. This prerequisite remains `PENDING` + until a maintainer lands it and the unchanged default gate passes on the + implementation base. Disabling skinny GEMM is diagnostic only. +2. The configured source and runnable ROCm wrapper are still at historical + vLLM `5559679229bc961848b121ccdeaa8fa5d79bec98`. The active source object is + available for read-only inspection, but no runnable active-pin ROCm runtime + has been supplied or proved. Issue #2794 records repository pin-validation + context; #2773 itself owns the cache-matched active-pin Qwen3.5-0.8B captures. + Active-pin capture and token revalidation remain `PENDING` until that runtime + exists and runs the model under the GPU mutex. + +Do not create characterization goldens from a known-regressed local default or +from the historical oracle revision. + +## Scope + +### In scope + +- Add opt-in production-boundary dumps for the residual stream, attention + output, MLP output, stored K/V cache, and persistent GDN convolution/SSM + state. +- Compare CPU and gfx1100 on identical weights and identical token prefixes. +- Report descriptive `max_abs`, RMS, and relative-L2 deltas. +- Port applicable active-pin upstream operation tests with their parameters, + modes, fixtures, failure cases, and exact tolerances. +- Audit runtime dtype, byte width, requested/resolved cache dtype, selected + attention backend, and provider counts for `auto`, `bfloat16`, and + `fp8_e4m3`. +- Run cache-matched active-pin vLLM captures and the established end-to-end + oracle-backed token/near-tie gate after the prerequisites pass. + +### Out of scope + +- Changing a numerical tolerance after observing results. +- Treating CPU as the conformance oracle for FP8 physical-cache behavior. +- Permanently expanding the sacred 16-prompt gate merely because this + characterization measures three cache modes. A new permanent case or golden + may land only when #2773's active-pin capture proves it is necessary and the + final implementation records the evidence. Broader gate policy needs its own + issue. +- Enabling ROCm static graph mode; issue #332 owns that work. +- Performance claims or tuning. + +## Artifact and oracle pins + +Use `Qwen/Qwen3.5-0.8B@2fc06364715b967f1860aea9cf38778875588b17` +from `${CHECKPOINT_ROOT}/Qwen3.5-0.8B`. Before every model run, require: + +```sh +sha256sum \ + "${CHECKPOINT_ROOT}/Qwen3.5-0.8B/model.safetensors-00001-of-00001.safetensors" \ + "${CHECKPOINT_ROOT}/Qwen3.5-0.8B/model.safetensors.index.json" \ + "${CHECKPOINT_ROOT}/Qwen3.5-0.8B/config.json" +``` + +Expected hashes are: + +```text +04b1c301231dd422b8860db31311ab2721511346a32cb1e079c4c4e5f1fe4696 model.safetensors-00001-of-00001.safetensors +d8a08838a613b025eb7952ed9db11696213e57e76a375661ef5c12f9dd5dcf4e model.safetensors.index.json +b90b86f35c8e6925ef74ee04d0e758f0a845c83a42089ad82bbaa948de9b4204 config.json +``` + +Resolve the active oracle source with: + +```sh +git -C "${VLLM_SOURCE}" rev-parse e126687a9a828d513c01a07cd69f025f27d63280^{commit} +``` + +The configured checkout may remain detached at the historical revision. Read +active-pin files with `git -C "${VLLM_SOURCE}" show :`; do not move +that checkout. Source availability does not prove runtime gateability. + +## Upstream executing chain + +All cross-file anchors name symbols and refer to active pin `e126687a9` unless +explicitly labeled historical: + +- `vllm/model_executor/models/qwen3_5.py::Qwen3_5ForCausalLM` defines the text + model and layer composition. +- `vllm/model_executor/layers/mamba/mamba_utils.py::MambaStateShapeCalculator.gated_delta_net_state_shape` + defines the GDN convolution and SSM state shapes. It starts at line 258 at + the active pin. `MambaStateShapeCalculator.mamba2_state_shape` is a different + architecture and is not the anchor. +- `tests/kernels/mamba/cpu/test_cpu_gdn_ops.py::test_fused_sigmoid_gating_delta_rule_update_cpu` + and `test_chunk_gated_delta_rule_cpu` define CPU operation tolerances. +- `tests/kernels/mamba/test_gdn_forward_core_split.py::test_forward_core_split_matches_unified` + defines split/unified state and output comparisons. +- `tests/kernels/attention/test_cache.py::test_reshape_and_cache` defines + cache-store dtype behavior and FP8 comparison tolerances. +- `tests/kernels/attention/test_cpu_attn.py::varlen_with_paged_kv` starts at + line 415. Its `_FP8_ATOL` and `_FP8_RTOL` constants at lines 45 and 46 + set `atol` to `0.2` for E4M3 and `0.3` for E5M2, with `rtol=0.1`. + The comparisons at lines 629 to 653 select these bounds for FP8 and + `atol=1.5e-2`, `rtol=1e-2` otherwise. Preserve both split and unsplit checks. + +The earlier spec cited these surfaces from historical `555967922`; those +citations remain diagnostic history but cannot satisfy the active-pin gate. The +implementer must refresh line evidence from the active object and run the active +runtime before acceptance. + +The local production entry is `ModelRegistry::Forward`, which dispatches the +registered `ForwardQwen3_5Dense` factory. Its paged forward reaches +`DenseForwardBody` -> `DenseForwardLayers` -> `RunDenseLayerPaged` -> +`FullAttnBlockPaged` -> `dense_attn::WriteKvCache` -> `vt::ReshapeAndCache` or +`vt::ReshapeAndCacheFp8`. The full-attention caller is +`src/vllm/model_executor/models/qwen3_5.cpp:7758`. `FullAttnBlockPaged` starts +at line 5726 and writes through `dense_attn::WriteKvCache` at +`src/vllm/model_executor/models/qwen3_5.cpp:5909`, before `vt::PagedAttention` +at line 5931. Place the stored-K/V probes after that write. Place GDN probes +after the persistent writes in `GdnBlockPaged` for prefill and decode. + +This Qwen3.5 path does not call `dense_attn::AttnBlock`. Issue #2923 under +`Owed` records this existing shared-seam debt. Characterize the actual path +and cite the executing local and active-pin upstream symbols in the evidence. + +## Dtype and state contract + +The checkpoint resolves model dtype BF16 and explicitly sets +`mamba_ssm_dtype=float32`. The measured surfaces are: + +| Surface | Shape per active unit | Physical storage | +|---|---|---| +| Hidden/residual | `[T, 1024]` | BF16 | +| Full-attention K/V | `[T, 2, 256]` per layer | BF16 or FP8 E4M3 | +| GDN convolution state | Per layer and slot: oracle SD `[3, 6144]`, local DS `[6144, 3]` | BF16 | +| GDN SSM state | `[1, 16, 128, 128]` per layer | FP32 | + +The active oracle defaults to SD physical convolution storage, with axes +`(state_len, dim)`. DS reverses these axes to `(dim, state_len)`. +`vllm/model_executor/layers/mamba/mamba_utils.py::get_conv_state_layout`, lines +28 to 44, returns `SD` without an override. `_orient_conv_shape`, lines 162 +to 166, and `gated_delta_net_state_shape`, lines 268 to 272, produce +`[3, 6144]` per slot for this workload. The kernel consumes a DS view +`[6144, 3]` after the transpose in +`vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py:1296` to line 1303. +Local `MakeQwen3_5KVCacheSpec` stores DS directly in +`src/vllm/model_executor/models/qwen3_5_common.cpp:85`. Leave +`VLLM_SSM_CONV_STATE_LAYOUT` unset for the production oracle. Record physical +layout, kernel-view layout, and actual strides on both sides. A transposed +view does not change the underlying physical layout or byte count. + +Across six full-attention layers, BF16 K/V use 12,288 bytes per cached token +and FP8 K/V use 6,144 bytes. Across 18 GDN layers, convolution state uses +663,552 bytes and SSM state uses 18,874,368 bytes per active slot. Confirm +allocated bytes from runtime specs; arithmetic alone is not evidence. + +For every compared buffer and GEMM output, log the local symbol, active-pin +upstream symbol, runtime dtype, shape, active elements, bytes on each side, and +the reason beside every FP32 exception. An unjustified wider local dtype is a +`DTYPE_WIDTH_MISMATCH` even when tokens agree. + +## State dump design + +Enable dumps only when `VT_DUMP_ACT` and `VT_DUMP_ACT_SUB` name the same empty, +writable directory. Preserve existing keys and add: + +| Key | Production source | Dump dtype | Shape | +|---|---|---|---| +| `state_fa_k` | stored K after `dense_attn::WriteKvCache` | FP32 | `[active_tokens, 512]` | +| `state_fa_v` | stored V after `dense_attn::WriteKvCache` | FP32 | `[active_tokens, 512]` | +| `state_gdn_conv` | persistent convolution cache after update | FP32 | `[active_requests, 18432]` | +| `state_gdn_ssm` | persistent SSM cache after update | FP32 | `[active_requests, 262144]` | + +Widen stored BF16 to FP32. Dequantize stored FP8 with the recorded per-layer +scale. The FP32 file format never conceals separately logged physical dtype and +allocation. + +Use canonical DS order for `state_gdn_conv`, with `0 <= d < 6144` and +`0 <= s < 3`: + +```text +local_dump[r, d * 3 + s] = widen(local[slot_local(r), d, s]) +oracle_dump[r, d * 3 + s] = widen(oracle[slot_oracle(r), s, d]) +``` + +Here `r` follows scheduler request order, and each arm resolves its own slot. +Use the recorded strides and storage offsets for each lookup. Record this +mapping in the manifest. The canonical DS dump order does not prescribe +physical storage or make the oracle's transposed view contiguous. + +Gather nonnegative slot mappings in input-token order and active GDN indices in +scheduler request order. Refuse unexplained duplicate destinations, missing +rows, partial joins, or capacity-wide dumps. A 24-layer step emits 48 state +blobs, two per layer. Place each probe after its production write and preserve +queue ordering. + +## Workload and backend identity + +Use one binary and checkpoint copy. Run one request at concurrency one with +`The capital of France is`, greedy sampling, seed 0, MTP disabled, and eight +output tokens. Keep block size, block count, and scheduler budget equal. Run +each arm and cache mode in a separate process and empty directory. Take two +enabled repeats and one dump-disabled identity control. + +Log and assert requested cache dtype, resolved storage and FP8 interpretation, +normalized selector input, physical cache dtype, selected attention backend, +and actual provider execution for every arm and mode. +`LoadedEngine::ApplyResolvedCacheDType` in +`src/vllm/entrypoints/model_loader.cpp:1835` applies `ParseCacheDType` from +`include/vllm/v1/kv_cache_dtype.h:58` to the cache spec. The runner constructs +`cfg.kv_cache_dtype` from resolved storage and interpretation at +`src/vllm/v1/worker/gpu/runner.cpp:1590`. `KvCacheDTypeName` maps `kBF16` to +`auto` in `src/vllm/v1/attention/backend.cpp:91` to line 96. Therefore the +production CPU selector receives these values for this BF16 workload: + +| Requested cache dtype | Resolved physical storage | Normalized selector input | Selected CPU backend | +|---|---|---|---| +| `auto` | BF16 (`kBF16`) | `auto` | `CPU_ATTN` | +| `bfloat16` | BF16 (`kBF16`) | `auto` | `CPU_ATTN` | +| `fp8_e4m3` | E4M3 bytes (`kI8` plus `kFp8E4M3`) | `fp8_e4m3` | `CPU_ATTN` | + +CPU priority starts with `CPU_ATTN` in `src/vllm/platforms/cpu.cpp:51`. +Its accepted selector values include `auto` and `fp8_e4m3` in +`include/vllm/v1/attention/backends/cpu_attn.h:108` to line 110. +`SelectAttentionBackendName` returns the first valid candidate in +`src/vllm/v1/attention/registry.cpp:124` to line 128. +Selection alone does not prove execution. Independently record the executing +attention and cache-store operator providers, selection counts, declines, and +reference-tier hits. The runner's selected name validates cache configuration +but does not dispatch the model's attention call. If main changes this chain +before implementation, update the cited contract before code and obtain fresh +review. + +CPU is a diagnostic comparison arm for CPU-versus-ROCm deltas, including FP8 +state after dequantization. Active-pin vLLM and its upstream tests are the +conformance authority. CPU results cannot accept or reject FP8 physical-cache +semantics. + +Join only steps with identical input token prefixes. Include the first divergent +output step and exclude all later steps. Record the first divergence and number +of excluded steps. + +## Descriptive metrics and acceptance + +Decode joined blobs to FP32. Let `A` be ROCm and `B` be CPU. Report: + +```text +max_abs = max_i(abs(A_i - B_i)) +rms = sqrt(sum_i((A_i - B_i)^2) / N) +rel_l2 = sqrt(sum_i((A_i - B_i)^2)) / sqrt(sum_i(B_i^2)) +``` + +Define relative L2 as zero when both norms are zero and infinity when only `B` +has zero norm. Report the worst index, both manifests, joined/rejected rows, and +excluded steps. Refuse missing or duplicate keys, dtype/shape/byte mismatches, +nonfinite values, and partial joins. + +These layer metrics are descriptive. They have no propagated numerical +acceptance envelope, discontinuity multiplier, or post-measurement threshold. +The rejected `E(k)` construction is removed because operation counts do not +bound cancellation, conditioning, nonlinear sensitivity, correlated +reductions, FP8 scaling/saturation, or subnormals. + +Acceptance consists of both: + +1. Ported active-pin operation tests with their exact upstream per-surface + tolerances. Examples include GDN `atol=rtol=1e-2` where + `test_cpu_gdn_ops.py` specifies it, exact convolution-state equality in the + split-core test, its dtype-specific output/SSM tolerances, and cache/CPU + attention FP8 tolerances selected by the upstream fixtures. Do not collapse + them into one project-wide tolerance. +2. The established end-to-end 16-prompt strict-token and ratified 500 + milli-nat near-tie gate against deterministic, cache-matched, active-pin + vLLM captures. A new physical cache mode needs its own capture and + teacher-forced gaps before it can be accepted. + +If no upstream analogue exists for a layer boundary, report its metrics only. +A future stage-level acceptance rule requires an independently justified bound +or explicit developer ratification before values are inspected. + +## Provider and end-to-end contract + +For `auto` and `bfloat16`, require the established native operator set, +including nonzero `kReshapeAndCache`, zero `kReshapeAndCacheFp8`, zero declines, +and zero reference-tier hits. The two modes resolve to physical BF16 and must +produce identical local token streams. + +For `fp8_e4m3`, replace the store requirement: require nonzero +`kReshapeAndCacheFp8`, zero `kReshapeAndCache`, nonzero native selections for +all other applicable operators, zero declines, and zero reference-tier hits. +Provider sets are mode-specific; never require a BF16 store from a correct FP8 +run. + +The existing gate remains the permanent gate until #2773 produces reviewed +active-pin evidence for an additional case. Characterization must run all three +modes, but this spec alone does not authorize new permanent goldens. + +## Required capture-tool implementation + +The current `scripts/qwen3-oracle-capture.py` and +`scripts/qwen3-neartie-gap.py` hard-code `enforce_eager=True` and expose no +cache-dtype argument. Before active-pin capture, add reviewed options equivalent +to: + +```text +--kv-cache-dtype {auto,bfloat16,fp8_e4m3} +--execution-mode {production,eager} +--seed INT +--max-tokens INT +--repetitions INT +--model-revision REV +--vllm-revision REV +--provenance-out PATH +``` + +`production` must instantiate vLLM without `enforce_eager=True`; eager is a +diagnostic arm and never the denominator. The output must record source, wheel, +image and artifact hashes, complete arguments, a hash of the scripts' shared +16-entry `PROMPTS` list, batching, concurrency, sampling, seed, token count, +repetitions, cache mode, resolved +cache dtype, execution mode, and output hash. + +Extend `test_qwen35_paged_engine` with future environment inputs +`VT_QWEN35_GATE_DIR` and `VT_QWEN35_KV_CACHE_DTYPE`. The first selects an empty +issue-evidence directory instead of the committed golden directory. The second +sets `EngineParams::kv_cache_dtype` and prints requested, resolved, and physical +dtype. These commands describe the intended interface and data flow after +those changes. They are future commands and will fail today: + +```sh +VLLM_PIN=e126687a9a828d513c01a07cd69f025f27d63280 +MODEL_REV=2fc06364715b967f1860aea9cf38778875588b17 +MODEL="${CHECKPOINT_ROOT}/Qwen3.5-0.8B" +for mode in auto bfloat16 fp8_e4m3; do + GOLDEN_DIR="evidence/2773/sacred-${mode}" + test ! -e "${GOLDEN_DIR}" + mkdir -p "${GOLDEN_DIR}" + "${VLLM_ORACLE}" scripts/qwen3-oracle-capture.py \ + --model "${MODEL}" --model-revision "${MODEL_REV}" \ + --vllm-revision "${VLLM_PIN}" \ + --kv-cache-dtype "${mode}" --execution-mode production --seed 0 \ + --max-tokens 16 --runs 10 --per-prompt --out-dir "${GOLDEN_DIR}" \ + --provenance-out "evidence/2773/oracle-${mode}.json" + VT_QWEN35_GATE_DIR="${GOLDEN_DIR}" \ + VT_QWEN35_KV_CACHE_DTYPE="${mode}" \ + VT_DUMP_IDS=1 build-rocm/tests/test_qwen35_paged_engine + test -s "${GOLDEN_DIR}/our_ids.i32" + "${VLLM_ORACLE}" scripts/qwen3-neartie-gap.py \ + --model "${MODEL}" --model-revision "${MODEL_REV}" \ + --vllm-revision "${VLLM_PIN}" \ + --kv-cache-dtype "${mode}" --execution-mode production --seed 0 \ + --max-tokens 16 --topk 20 --golden-dir "${GOLDEN_DIR}" \ + --provenance-out "evidence/2773/neartie-${mode}.json" + VT_QWEN35_GATE_DIR="${GOLDEN_DIR}" \ + VT_QWEN35_KV_CACHE_DTYPE="${mode}" \ + build-rocm/tests/test_qwen35_paged_engine +done +``` + +The tools must refuse when their `PROMPTS` lists differ from each other or from +`tests/parity/test_qwen35_paged_engine.cpp::Prompts`. Record exact as-run +commands. Do not create a golden unless all 10 repeats are deterministic. + +The internal CPU/ROCm state characterization remains eight output tokens as +specified under `Workload and backend identity`. The permanent sacred-gate +candidate uses the existing 16-prompt, 16-output-token regime above. Never use +the eight-token characterization files as sacred-gate goldens. + +## Tests and review mutations + +The later implementation starts with focused tests that fail because the four +state rows and capture options are absent. It must test row shape/dtype, +nonmonotonic and negative slot mappings, inactive-capacity exclusion, stored +BF16 widening, post-quantization FP8 dequantization, persistent GDN prefill and +decode state, incomplete-step refusal, comparator structural refusals and +zero-denominator rules, dtype-width refusal, backend/dtype logging, and +mode-appropriate provider sets. The checkpoint-backed case must enter through +`LoadedEngine::FromModelDir`. + +A fresh reviewer mutates each guarantee in a scratch copy: remove each +production probe, substitute pre-write tensors, remove negative-slot filtering, +sort by physical slot, dump capacity, misreport FP8 as BF16, swap the mode's +cache-store operator, inject a reference-tier hit, remove the production entry +point, and force eager oracle mode. Each focused test must fail for the intended +reason and the reviewer must restore the tree byte for byte. + +This spec repair changes no runtime behavior. `IMP-TEST-FIRST` and +`IMP-MUTATE` are future implementation/review gates, not evidence claimed by +this commit. + +## Gate order + +The later implementation must satisfy, in order: + +1. #2772 lands and the unchanged default local gate passes on the chosen base. +2. A runnable active-pin ROCm vLLM runtime is identified and proves this model. +3. Artifact, source, wrapper, wheel, image, prompt, and output hashes are saved. +4. The capture tools fail first for missing options, then pass focused tests. +5. Ported active-pin operation tests pass with unchanged upstream tolerances. +6. Dump/comparator CPU tests and checkpoint-backed CPU trace pass. +7. The controlled full preflight runs; every skip remains `PENDING`. +8. Fresh immutable review detects every required mutation. +9. Two CPU and two gfx1100 trace repeats plus disabled controls pass per mode. +10. Cache-matched active-pin production captures are deterministic and the + established end-to-end gate passes per measured mode. +11. Mode-appropriate native selections are nonzero, with zero declines and + zero reference-tier hits. +12. The operator independently reruns focused, full, oracle, and hardware gates. + +No GPU result, oracle result, model execution, runtime mutation, or +implementation test is claimed by this spec-only repair. + +## Evidence required + +Store evidence under an issue-specific durable path and record commit/tree +hashes, clean status, exact commands/statuses, every skip, active oracle and +artifact hashes, compiler/build/ROCm/driver/board identity, GPU mutex evidence, +requested/resolved/physical cache dtype, selected CPU and ROCm attention +backends, provider counts, enabled repeats, disabled controls, manifests, raw +dumps, metric tables, first divergence, excluded steps, upstream tolerance +anchors, red-first output, mutations, fresh review, and operator reruns. + +Classify structural and execution failures before interpreting metrics: +`ARTIFACT_MISMATCH`, `INSTRUMENTATION_FAIL`, `NONDETERMINISTIC`, +`STRUCTURE_MISMATCH`, `DTYPE_WIDTH_MISMATCH`, `CORRECTNESS_FAIL`, or +`NONFINITE`. When none applies, report descriptive metrics and the separate +operation/end-to-end acceptance results. Do not invent `WITHIN_DTYPE_ENVELOPE` +or `ORDERING_DRIFT` labels. + +## Risks + +- Probes can synchronize queues; disabled controls detect output perturbation. +- FP8 boundaries can change codes; compare dequantized values and record bytes. +- Different valid reduction orders can produce descriptive CPU/ROCm deltas. +- Backend fallback can change the diagnostic denominator; explicit backend + logging exposes it. +- Greedy divergence makes later states incomparable; exclude them. +- The active source object can be read while the active runtime remains absent; + never convert source availability into a gate pass. + +## Stop conditions + +- Stop before implementation until this spec receives a fresh `PASS`. +- Keep correctness `PENDING` until #2772 lands and the unchanged default gate + passes on the implementation base. +- Stop active-pin capture until a runnable active-pin ROCm runtime is proved. +- Stop on a revision/hash/configuration mismatch, dump perturbation, + nondeterminism, incomplete manifest, unexplained backend, provider decline, + reference-tier hit, wrong board, missing GPU mutex, or divergent input prefix. +- Stop before changing tolerances or permanent sacred-gate scope after seeing + results. +- File and assign a new issue before fixing any unexpected defect outside + #2773. + +## Git integration + +Use one pull request for the committed spec and later implementation, following +the recorded repository default. Preserve rejected spec commit `7bc2546e9` in +history. The eventual pull request body must name row `BACKEND-ROCM`, link and +close #2773, and carry the required trailers. The spec implementer does not +push, open, or merge that pull request. + +## Owed + +- A maintainer owes the merge decision for reviewed PR #2856; #2772 remains + pending until it lands. +- #2773 owes the runnable active-pin Qwen3.5-0.8B ROCm captures, even though + #2794 supplies repository sync context. +- [#2923](https://github.com/mudler/vllm.cpp/issues/2923), owned by + `BACKEND-ROCM`, owes routing Qwen3.5 paged attention through + `dense_attn::AttnBlock`. This tracked exception records existing debt. It + does not waive or satisfy the shared-seam requirement. Wiring needs its own + reviewed spec and implementation, outside #2773's instrumentation scope. +- A fresh reviewer owes this repaired spec a verdict. +- A fresh implementer owes the capture options, state probes, tests, comparator, + and provider checks after the prerequisites pass. +- A fresh implementation reviewer owes static review and every mutation. +- The operator owes the independent controlled and gfx1100 gates. +- The final implementation adds `## Outcome` with measured results, rejected + alternatives, and reasons for defaults. +- Issue #332 owns the graph-enabled repetition. diff --git a/CMakeLists.txt b/CMakeLists.txt index d91f9a495..9bf89e4b2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1739,7 +1739,8 @@ if(VLLM_CPP_HIP) src/vt/rocm/rocm_mla_ops.hip src/vt/rocm/rocm_mla_attn.hip src/vt/rocm/rocm_skinny_gemm.hip - src/vt/rocm/rocm_ops.hip) + src/vt/rocm/rocm_ops.hip + src/vt/rocm/rocm_quant_dot.hip) if(VLLM_CPP_HIP_ARCHITECTURES) set_source_files_properties( src/vt/rocm/rocm_backend.hip @@ -1768,6 +1769,7 @@ if(VLLM_CPP_HIP) src/vt/rocm/rocm_mla_attn.hip src/vt/rocm/rocm_skinny_gemm.hip src/vt/rocm/rocm_ops.hip + src/vt/rocm/rocm_quant_dot.hip PROPERTIES HIP_ARCHITECTURES "${VLLM_CPP_HIP_ARCHITECTURES}") endif() # Prefer the absolute path inside ${ROCM_PATH}/lib, fall back to the bare name, diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index 54af017f6..e336c0647 100644 --- a/docs/ENVIRONMENT.md +++ b/docs/ENVIRONMENT.md @@ -224,6 +224,15 @@ portable/reference path. In normal operation leave them unset. | `VT_GPU_SAMPLE` | on (CUDA) | Host-side sampling instead of on-GPU sampling | | `VT_ROCM_Q8K_BLOCK` | unset (cooperative only on queue-device-resolved `gfx1100`) | Exact `0` selects the permanent legacy A/B arm. Exact `1` forces the cooperative candidate diagnostically. Explicit `1` outside validated `gfx1100` is not a support claim or a default claim. Unset keeps the legacy arm on `gfx1200`, `gfx1201`, unknown architectures, and architecture-resolution failure. Every other value is refused. See [Q8_K activation quantization](ROCM.md#select-q8_k-activation-quantization) | | `VT_ROCM_QUANT_WMMA` | on | ROCm keep-quant Q6_K and Q4_K prefill GEMM: the RDNA4 rocWMMA int8 tile arm (`KQuantGemmKWmmaQ6K`, `KQuantGemmKWmmaQ4K`), queue-device-resolved `gfx1200`/`gfx1201` only and only when `m >= 16 && n >= 16` (at least one full 16-wide tile in each dimension, not exact alignment — the scalar `KQuantGemmK` remaps its launched index to enumerate only the remainder the WMMA corner leaves untouched (the bottom strip past the aligned row boundary, then the right strip past the aligned column boundary), so a non-16-multiple M/N still takes the WMMA arm over its floor(M/16)xfloor(N/16) corner without paying a full-`m*n`-grid launch for the fill). `0` forces the scalar `Dp4a` arm this row's spec (`KERNEL-QUANT-CIQ-GEMM-ROCM-RDNA4`, issue #2109) is chasing a performance gap against, on every architecture and every shape. Bit-identical to the scalar arm by construction (the WMMA tile's raw int8 dot is scaled and reduced in the same integer arithmetic the scalar path uses; only the one f32 scale product per superblock is shared with it) — this is a same-binary A/B, not a correctness fallback | +| `VT_GEMV_MMVQ` | off | `=1` routes the K-quant (Q4_K/Q5_K/Q6_K) `MatmulBTQuant` decode (m==1 only) through the MMVQ-style GEMV arm in `rocm_grouped_gemm.hip`, which folds activation quantization into the GEMV prologue; default keeps the baseline `KQuantGemmK` path byte-unchanged. The arm is bit-equal to the baseline output at every (Fmt, nsb, j); the fold crossover is tunable with `VT_GEMV_MMVQ_FOLD_MAX`. Flag read per call, so a captured decode graph picks the arm up at capture time | +| `VT_SKINNY_BF16` | off | `=1` dispatches the f32-output decode-skinny arm (`WvSplitKBTToF32`) for the bf16-input GEMVs whose consumers want f32 — the Qwen3.5 GDN BA projections' g/beta chain — instead of rocBLAS's large-M tile | +| `VT_GDN_COLPERM_KEEP_QUANT` | off | `=1` keeps the column-permuted `ssm_out`/`out_proj` tensor as Q5_K in tiled order (no `ReorderVCols`) and permutes the 4096-element GEMV input at runtime instead; the column reorder cuts across Q5_K block boundaries, so the weight cannot be permuted in place. Saves ~4x weight bandwidth (Q5_K ~5 MB vs bf16 20 MB per call) | +| `VT_GDN_ROWPERM_KEEP_QUANT` | off | `=1` keeps the row-permuted V-head GDN projections (in the tiled order the row permutation produces) as K-quant instead of expanding to bf16 at load; the runtime gather supplies the permutation. Opt-in; the default reorders then expands | +| `VT_ATTN_PREAMBLE_COOP` | off | `=1` selects the warp-per-item cooperative attention preamble arm (`AttnQkNormRopeGateCoopK`) on ROCm, mapping one warp per token item instead of the donor walk (`AttnQkNormRopeGateK`); read once per process like the sibling arms | +| `VT_RMSNORM_ROW_COOP` | off | `=1` selects the cooperative single-row RMSNorm arm (two wavefront `shfl_down` trees plus one cross-wavefront shared-memory combine — two `__syncthreads` instead of nine) when h is a multiple of 256 and at most one block wide; the float reduction order CHANGES, so this is an opt-in riding the campaign config, not a default change. Eligibility falls back to the ported body automatically | +| `VT_RMSNORM_LDS_QUANT` | on (when `VT_RMSNORM_ROW_COOP=1` and a Q8_K output is active) | `=0` disables the T24 LDS-buffered quant epilogue that stores the normalized row to shared memory during pass 3 so the Q8_K quant epilogue reads from LDS instead of re-reading global memory; the quantized bytes are unchanged either way | +| `VT_QUANT_Q8K_WARP` | off | `=1` selects the warp-cooperative `QuantizeQ8KKWarpCoop` activation quantizer (16 super-blocks per 128-thread block, per-chunk ascending amax with cross-warp shuffle reduction) instead of the selector-dispatched default on the decode path; byte-identical output, asserted on tied-amax rows | +| `VT_NORM_QUANT_FUSED` | off | `=1` fuses the Q8_K activation-superblock quantizer into the cooperative RmsNorm epilogue (one launch instead of norm + standalone quant); the producer registers its scratch in a consumer token bridge, so the consuming `MatmulBTQuant` decode SKIPS the standalone quant launch when the row set matches. Byte-equal to the standalone path by construction (both instantiate the shared `QuantQ8KSBlock` body from `rocm_act_quant.h`) and asserted op-level | | `VT_GDN_PACKED_DECODE` | on (CUDA GDN) | Unpacked GDN decode path | | `VT_GDN_DECODE_BV` | `32` (CUDA GDN decode experiment) | Exact `16` selects the byte-identical 16-value fused-recurrence tile; unset and every other spelling keep the 32-value schedule. Experimental opt-in; no release or cross-hardware default change | | `VT_GDN_DECODE_SWIZZLE` | `0` (CUDA GDN decode experiment) | Exact `1` enables the shared-memory bank swizzle only for the `BV=16`, `Dv=Dk=128`, eight-lane production geometry; all other values and shapes keep the incumbent layout | diff --git a/include/vllm/model_executor/models/qwen3_5_weights.h b/include/vllm/model_executor/models/qwen3_5_weights.h index c3c12f656..0b024fcb1 100644 --- a/include/vllm/model_executor/models/qwen3_5_weights.h +++ b/include/vllm/model_executor/models/qwen3_5_weights.h @@ -771,6 +771,7 @@ struct GdnLayerWeights { OwnedTensor dt_bias; // f32 [Hv] OwnedTensor norm_weight; // bf16 [Dv] (RMSNormGated) OwnedTensor out_proj; // bf16 [value_dim, H] (FP8 dequant + T) + bool out_proj_tiled = false; // T25: weight kept in tiled Q5_K order; permute input at runtime // MODEL-FP8-BLOCK-WEIGHT (#1189 M3): block-wise FP8 GDN projections. The // target checkpoint lists the GDN small tensors under diff --git a/include/vt/ops.h b/include/vt/ops.h index 2ef14ca1a..c6727b5a3 100644 --- a/include/vt/ops.h +++ b/include/vt/ops.h @@ -170,6 +170,7 @@ enum class OpId : uint8_t { kCastBf16, kCastF32, kCastF16, + kPermuteVHeads, kMulColVecF32, kAttnGateSplit, kSigmoidGateBf16, @@ -2242,6 +2243,8 @@ using MoeRelu2Fn = void (*)(Queue&, Tensor&, const Tensor&); // loops so the decode step can run entirely on-device (CUDA-graph capture). // All math in f32; dims are inferred from the tensor shapes (no args structs). using CastBf16Fn = void (*)(Queue&, Tensor&, const Tensor&); +using PermuteVHeadsFn = void (*)(Queue&, Tensor&, const Tensor&, int64_t, int64_t, + int64_t, int64_t); using CastF32Fn = void (*)(Queue&, Tensor&, const Tensor&); using CastF16Fn = void (*)(Queue&, Tensor&, const Tensor&); using MulColVecF32Fn = void (*)(Queue&, Tensor&, const Tensor&); @@ -5554,6 +5557,14 @@ void ApplyAllowedTokenIds(Queue& q, Tensor& logits, const Tensor& mask); // f32 -> bf16 activation-dtype cast used before feeding a bf16-consuming op. void CastBf16(Queue& q, Tensor& out, const Tensor& in); +// T25: Permute V-heads from grouped (k*rpk+r) to tiled (r*num_k+k) order. +// out[T, value_dim] = in[T, value_dim] with the last dim permuted: +// out[t*dv + h] = in[g*dv + h] where t = r*num_k + k, g = k*rpk + r +// Used before the K-quant GEMV when ssm_out is kept as Q5_K in tiled order. +// value_dim = num_k * rpk * dv. T, in, out are bf16. +void PermuteVHeads(Queue& q, Tensor& out, const Tensor& in, + int64_t T, int64_t num_k, int64_t rpk, int64_t dv); + // out[i] = f32(in[i]); out f32, in bf16, same element count. The bf16 -> f32 // upcast used to expose a bf16-only GEMM (Marlin) as an f32 result, matching the // value the bf16 output rounds to (mirror of the cutlass f32-output scratch cast). diff --git a/src/vllm/model_executor/models/qwen3_5.cpp b/src/vllm/model_executor/models/qwen3_5.cpp index 83ea59640..0968e5724 100644 --- a/src/vllm/model_executor/models/qwen3_5.cpp +++ b/src/vllm/model_executor/models/qwen3_5.cpp @@ -1786,6 +1786,24 @@ DBuf MatmulBf16D(Dev d, const Tensor& x, const OwnedTensor& w) { return dout; } +// T25: When out_proj is kept as K-quant in tiled order (out_proj_tiled), permute +// the gated-norm output from grouped→tiled before the K-quant GEMV. The `nk` +// flag alone is insufficient: gdn_expand_nk also sets nk=true for the bf16 +// expanded weight, but that weight has ReorderVCols applied and needs NO +// input permutation. Only the T25 tiled Q5_K path (out_proj_tiled=true) does. +static DBuf GdnOutProjMatmul(Dev d, const GdnLayerWeights& w, + const DBuf& gated_bf16, + int64_t T, int64_t Hk, int64_t Hv, int64_t Dv) { + if (w.out_proj_tiled) { + const int64_t value_dim = Hv * Dv; + const int64_t rpk = Hk > 0 ? Hv / Hk : 1; + DBuf permuted(d, DType::kBF16, {T, value_dim}); + vt::PermuteVHeads(d.q, permuted.t(), gated_bf16.t(), T, Hk, rpk, Dv); + return MatmulBf16D(d, permuted.t(), w.out_proj); + } + return MatmulBf16D(d, gated_bf16.t(), w.out_proj); +} + // A tied BF16 lm_head follows torch Linear's model-dtype output, then the // engine exposes f32 logits to the sampler. Explicit 27B heads retain the // existing f32-output MatmulF32D path. @@ -4600,7 +4618,7 @@ DBuf GdnBlock(Dev d, const GdnLayerWeights& w, const HfConfig& cfg, ? MatmulFp8CutlassD(d, gated_bf16.t(), w.out_proj_fp8, DType::kBF16) : !w.out_proj_fp4.Empty() ? MatmulNvfp4Bf16D(d, gated_bf16.t(), w.out_proj_fp4) - : MatmulBf16D(d, gated_bf16.t(), w.out_proj); // [T,H] + : GdnOutProjMatmul(d, w, gated_bf16, T, Hk, Hv, Dv); // [T,H] } // PERSISTENT per-step input device buffers (decode host-tax #2): the flattened @@ -5087,7 +5105,7 @@ DBuf GdnBlockPagedMixedSpec(Dev d, const GdnLayerWeights& w, const HfConfig& cfg ? MatmulFp8CutlassD(d, gated_bf16.t(), w.out_proj_fp8, DType::kBF16) : !w.out_proj_fp4.Empty() ? MatmulNvfp4Bf16D(d, gated_bf16.t(), w.out_proj_fp4) - : MatmulBf16D(d, gated_bf16.t(), w.out_proj); // [T,H] + : GdnOutProjMatmul(d, w, gated_bf16, T, Hk, Hv, Dv); // [T,H] } // VT_DUMP_ACT stage probe (GDN): dump named intermediates so a layer-level @@ -5592,7 +5610,7 @@ DBuf GdnBlockPaged(Dev d, const GdnLayerWeights& w, const HfConfig& cfg, ? MatmulFp8CutlassD(d, gated_bf16.t(), w.out_proj_fp8, DType::kBF16) : !w.out_proj_fp4.Empty() ? MatmulNvfp4Bf16D(d, gated_bf16.t(), w.out_proj_fp4) - : MatmulBf16D(d, gated_bf16.t(), w.out_proj); // [T,H] + : GdnOutProjMatmul(d, w, gated_bf16, T, Hk, Hv, Dv); // [T,H] } // --- Dense full_attention block. qwen36-forward-notes.md §5; pinned diff --git a/src/vllm/model_executor/models/qwen3_5_gguf_weights.cpp b/src/vllm/model_executor/models/qwen3_5_gguf_weights.cpp index f341873d0..4ae38a1f2 100644 --- a/src/vllm/model_executor/models/qwen3_5_gguf_weights.cpp +++ b/src/vllm/model_executor/models/qwen3_5_gguf_weights.cpp @@ -383,6 +383,23 @@ void ReorderVRows(std::vector& buf, int64_t cols, int64_t row_off, } std::memcpy(base, seg.data(), seg.size() * sizeof(T)); } +// Pointer-based overload for OwnedBytes (T=uint8_t, sizeof(T)=1). +void ReorderVRows(uint8_t* buf, int64_t cols, int64_t row_off, + int64_t num_k, int64_t num_v_per_k, int64_t head_rows) { + const int64_t num_v = num_k * num_v_per_k; + const int64_t head_stride = head_rows * cols; + std::vector seg(static_cast(num_v) * head_stride); + uint8_t* base = buf + row_off * cols; + for (int64_t k = 0; k < num_k; ++k) { + for (int64_t r = 0; r < num_v_per_k; ++r) { + const int64_t g = k * num_v_per_k + r; + const int64_t t = r * num_k + k; + std::memcpy(seg.data() + g * head_stride, base + t * head_stride, + static_cast(head_stride)); + } + } + std::memcpy(base, seg.data(), seg.size()); +} // Reorder the full column range [0, cols) of a [rows, cols] row-major buffer // (cols = num_v * head_cols) from GGUF tiled to HF grouped order (out_proj). @@ -1073,14 +1090,41 @@ GdnLayerWeights LoadGdnGguf(const GgufFile& g, int64_t il, const HfConfig& c, const int64_t key_dim = num_k * c.linear_key_head_dim; const bool reorder = num_v != num_k && num_k > 0 && (num_v % num_k) == 0; const int64_t rpk = num_k > 0 ? num_v / num_k : 1; // num_v_per_k - // When the V-head reorder is active these projections are LAYOUT-rewritten - // at load, so they are kTransformedWeight and can never keep their blocks; - // without it they are ordinary verbatim GEMM weights. (out_proj's reorder - // permutes COLUMNS, which live inside a block, so it is unconditionally - // block-unsafe when active — same rule, stated per tensor below.) + // When the V-head reorder is active, the projections are LAYOUT-rewritten at + // load. For COLUMN-permuted tensors (out_proj/ssm_out) the reorder cuts across + // quantization block boundaries, so they are kTransformedWeight and must + // expand to bf16. For ROW-permuted tensors (in_proj_qkv, in_proj_z) the + // reorder only changes row order — quantization blocks are along the K + // (column) dimension and are self-contained per row — so the blocks can be + // kept and the permutation applied to the block rows at load time (T21). + // Without reorder they are ordinary verbatim GEMM weights. Column-permuted + // tensors (out_proj/ssm_out) stay kTransformedWeight and expand to bf16. + // T21 env gate: VT_GDN_ROWPERM_KEEP_QUANT=0 forces the row-permuted tensors + // back to kTransformedWeight (bf16 expansion) for A/B isolation. + // Opt-in, by VALUE like its VT_GDN_COLPERM_KEEP_QUANT sibling: unset keeps + // the load path byte-identical to the trunk (reorder, then expand to bf16). + // A default-on flip here silently changed the trunk numerics of every + // row-permuted V-head projection and broke the token-exact gate -- caught on + // the gfx1100 engine leg, reverted to opt-in. + const char* rpkq = std::getenv("VT_GDN_ROWPERM_KEEP_QUANT"); + const bool rowperm_keep = + rpkq != nullptr && rpkq[0] == '1' && rpkq[1] == '\0'; const GgufTensorRole proj_role = reorder ? GgufTensorRole::kTransformedWeight : GgufTensorRole::kMatmulWeight; + const GgufTensorRole rowperm_role = + (reorder && rowperm_keep) ? GgufTensorRole::kMatmulWeight : proj_role; + // T25: keep the COLUMN-permuted tensor (ssm_out/out_proj) as K-quant in tiled + // order (no ReorderVCols) and permute the GEMV input at runtime instead. The + // column reorder cuts across Q5_K block boundaries, so the weight cannot be + // permuted in place. But keeping the tiled-order weight and permuting the + // 4096-element activation gather before the K-quant GEMV saves ~4x weight + // bandwidth (Q5_K ~5 MB vs bf16 20 MB per call). + const char* cpkq = std::getenv("VT_GDN_COLPERM_KEEP_QUANT"); + const bool colperm_keep = + cpkq != nullptr && cpkq[0] == '1' && cpkq[1] == '\0'; + const GgufTensorRole colperm_role = + (reorder && colperm_keep) ? GgufTensorRole::kMatmulWeight : proj_role; // GdnLayerWeights carries an Nvfp4Weight ONLY for out_proj, and even that is // unreachable on the 27B because the V-column reorder makes ssm_out // kTransformedWeight. The in_proj family has no fp4 field at all. So the GDN @@ -1093,11 +1137,27 @@ GdnLayerWeights LoadGdnGguf(const GgufFile& g, int64_t il, const HfConfig& c, GdnLayerWeights gdn; // in_proj_qkv <- attn_qkv [conv_dim, H]; only the trailing V rows reorder. + // T21: ReorderVRows is a row permutation (block-safe for K-quant). Route as + // kMatmulWeight to allow keep-quant, then permute the block rows in place. + // Saves ~661 MB/tok of bf16 read amplification (24 Q5_K tensors × 2.9x). + // The forward pass already dispatches quantized nk=true weights through + // vt::MatmulBT → matmul_bt_quant, so no forward-pass change is needed. { const std::string nm = Blk(il, "attn_qkv.weight"); - const GgufResidency r = pol.Route(g.Get(nm), proj_role); - if (r != GgufResidency::kExpandBf16) { - const GgufTensorInfo& ti = g.Get(nm); + const GgufTensorInfo& ti = g.Get(nm); + const GgufResidency r = pol.Route(ti, rowperm_role); + if (r == GgufResidency::kKeepQuant) { + // Force a copy (not mmap) so the block rows can be permuted in place. + OwnedTensor qk = OwnGgufQuantBlocks(ti, ti.shape[0], ti.shape[1], 0, + /*mmap_src=*/nullptr); + if (reorder) { + const int64_t row_bytes = static_cast(qk.bytes.size()) / + ti.shape[0]; + ReorderVRows(qk.bytes.data(), row_bytes, /*row_off=*/2 * key_dim, + num_k, rpk, dv); + } + gdn.in_proj_qkv = std::move(qk); + } else if (r != GgufResidency::kExpandBf16) { gdn.in_proj_qkv = OwnGgufKeptSlice(g, pol, ti, r, ti.shape[0], ti.shape[1], 0); } else { @@ -1110,11 +1170,22 @@ GdnLayerWeights LoadGdnGguf(const GgufFile& g, int64_t il, const HfConfig& c, } } // in_proj_z <- attn_gate [value_dim, H]; all rows are V. + // T21: Same row-permutation keep-quant path as in_proj_qkv above. + // Saves ~360 MB/tok of bf16 read amplification (24 Q4_K tensors × 2.9x). { const std::string nm = Blk(il, "attn_gate.weight"); - const GgufResidency r = pol.Route(g.Get(nm), proj_role); - if (r != GgufResidency::kExpandBf16) { - const GgufTensorInfo& ti = g.Get(nm); + const GgufTensorInfo& ti = g.Get(nm); + const GgufResidency r = pol.Route(ti, rowperm_role); + if (r == GgufResidency::kKeepQuant) { + OwnedTensor qk = OwnGgufQuantBlocks(ti, ti.shape[0], ti.shape[1], 0, + /*mmap_src=*/nullptr); + if (reorder) { + const int64_t row_bytes = static_cast(qk.bytes.size()) / + ti.shape[0]; + ReorderVRows(qk.bytes.data(), row_bytes, 0, num_k, rpk, dv); + } + gdn.in_proj_z = std::move(qk); + } else if (r != GgufResidency::kExpandBf16) { gdn.in_proj_z = OwnGgufKeptSlice(g, pol, ti, r, ti.shape[0], ti.shape[1], 0); } else { @@ -1159,11 +1230,22 @@ GdnLayerWeights LoadGdnGguf(const GgufFile& g, int64_t il, const HfConfig& c, } // out_proj <- ssm_out [H, value_dim]; reorder V columns, then transpose. // The COLUMN reorder cuts across block boundaries, so when it is active this - // tensor is kTransformedWeight and must expand. + // tensor is kTransformedWeight and must expand — UNLESS T25 + // (VT_GDN_COLPERM_KEEP_QUANT=1) keeps the tiled-order Q5_K weight and + // permutes the GEMV input at runtime instead. { const std::string nm = Blk(il, "ssm_out.weight"); - const GgufResidency r = pol.Route(g.Get(nm), proj_role); - if (r != GgufResidency::kExpandBf16) { + const GgufResidency r = pol.Route(g.Get(nm), colperm_role); + if (r == GgufResidency::kKeepQuant && colperm_keep) { + // T25: keep Q5_K in tiled order (no ReorderVCols). The forward pass + // permutes the 4096-element activation from grouped→tiled before the + // K-quant GEMV, saving ~4x weight bandwidth. + OwnedTensor qk = + OwnGgufQuantBlocks(g.Get(nm), g.Get(nm).shape[0], g.Get(nm).shape[1], + 0, /*mmap_src=*/nullptr); + gdn.out_proj = std::move(qk); + gdn.out_proj_tiled = true; + } else if (r != GgufResidency::kExpandBf16) { const GgufTensorInfo& ti = g.Get(nm); gdn.out_proj = OwnGgufKeptSlice(g, pol, ti, r, ti.shape[0], ti.shape[1], 0); diff --git a/src/vt/cpu/cpu_ops.cpp b/src/vt/cpu/cpu_ops.cpp index 5d8ebddd4..8f212f5ad 100644 --- a/src/vt/cpu/cpu_ops.cpp +++ b/src/vt/cpu/cpu_ops.cpp @@ -3937,10 +3937,6 @@ void CastF32Kernel(Queue&, Tensor& out, const Tensor& in) { }); } -// out[i] = F32ToF16(in[i]); out f16, in f32 or bf16, same element count. -// QUANT-EXL3 W1a (#2181). LoadF32 reads either source width as f32 and StoreF32 -// rounds once to the f16 destination (cpu_ops.cpp:44-51), so the bf16 source -// path is "widen exactly, then round once" rather than a reinterpretation. void CastF16Kernel(Queue&, Tensor& out, const Tensor& in) { const int64_t n = out.Numel(); ForRows(n, [&](int64_t r0, int64_t r1) { @@ -3948,6 +3944,24 @@ void CastF16Kernel(Queue&, Tensor& out, const Tensor& in) { }); } +// T25: Permute V-heads from grouped (k*rpk+r) to tiled (r*num_k+k) order. +void PermuteVHeadsKernel(Queue&, Tensor& out, const Tensor& in, + int64_t T, int64_t num_k, int64_t rpk, int64_t dv) { + const int64_t value_dim = num_k * rpk * dv; + auto* out_p = out.Ptr(); + const auto* in_p = in.Ptr(); + for (int64_t row = 0; row < T; ++row) { + for (int64_t t = 0; t < num_k * rpk; ++t) { + const int64_t r = t / num_k; + const int64_t k = t % num_k; + const int64_t g = k * rpk + r; + for (int64_t h = 0; h < dv; ++h) + out_p[row * value_dim + t * dv + h] = + in_p[row * value_dim + g * dv + h]; + } + } +} + // x[m,n] *= col[n]; x f32 OR bf16 [M,N] (inner-contiguous rows, row stride // x.stride[0]), col always f32 [N]. CPU sibling of the CUDA MulColVecF32 kernel, // and the portable reference every other backend ports FROM — so it carries the @@ -4369,6 +4383,9 @@ struct Registrar { RegisterOp(OpId::kDFlashBlockAttention, DeviceType::kCPU, reinterpret_cast( static_cast(&DFlashBlockAttentionKernel))); + RegisterOp(OpId::kPermuteVHeads, DeviceType::kCPU, + reinterpret_cast( + static_cast(&PermuteVHeadsKernel))); RegisterOp(OpId::kDFlashPagedBlockAttention, DeviceType::kCPU, reinterpret_cast( static_cast(&DFlashPagedBlockAttentionKernel))); diff --git a/src/vt/op_provider.cpp b/src/vt/op_provider.cpp index 071a02137..6307e9039 100644 --- a/src/vt/op_provider.cpp +++ b/src/vt/op_provider.cpp @@ -420,6 +420,8 @@ const char* OpNameImpl(OpId op) { return "CastF16"; case OpId::kCastF32: return "CastF32"; + case OpId::kPermuteVHeads: + return "PermuteVHeads"; case OpId::kMulColVecF32: return "MulColVecF32"; case OpId::kAttnGateSplit: diff --git a/src/vt/ops.cpp b/src/vt/ops.cpp index 89e00a187..5379a002f 100644 --- a/src/vt/ops.cpp +++ b/src/vt/ops.cpp @@ -5193,6 +5193,18 @@ void CastF16(Queue& q, Tensor& out, const Tensor& in) { reinterpret_cast(GetOp(OpId::kCastF16, q.device.type))(q, out, in); } +void PermuteVHeads(Queue& q, Tensor& out, const Tensor& in, + int64_t T, int64_t num_k, int64_t rpk, int64_t dv) { + VT_CHECK(out.dtype == DType::kBF16 && in.dtype == DType::kBF16, + "permute_v_heads: both tensors must be bf16"); + VT_CHECK(out.Numel() == in.Numel(), + "permute_v_heads: out/in must have the same element count"); + VT_CHECK(out.device == q.device && in.device == q.device, + "permute_v_heads: device mismatch"); + reinterpret_cast(GetOp(OpId::kPermuteVHeads, q.device.type))( + q, out, in, T, num_k, rpk, dv); +} + void CastF32(Queue& q, Tensor& out, const Tensor& in) { VT_CHECK(out.dtype == DType::kF32, "cast_f32: out must be f32"); VT_CHECK(in.dtype == DType::kBF16, "cast_f32: in must be bf16"); diff --git a/src/vt/rocm/rocm_act_quant.h b/src/vt/rocm/rocm_act_quant.h new file mode 100644 index 000000000..11b884dd8 --- /dev/null +++ b/src/vt/rocm/rocm_act_quant.h @@ -0,0 +1,124 @@ +// Shared ROCm device-side Q8_K activation-superblock quantizer (Lever C, +// GFX1100-TG200-NORMQ). One source of truth for the byte-exactness-critical +// numeric path: src/vt/rocm/rocm_grouped_gemm.hip (standalone QuantizeQ8KK + +// MMVQ fused-prologue) AND src/vt/rocm/rocm_rmsnorm.hip (producer-fused +// epilogue behind VT_NORM_QUANT_FUSED=1) both instantiate THIS body, so +// "byte-equal vs standalone" holds by construction rather than by two copies +// drifting. Contract carried over from cuda_quant_dot.cu QuantizeQ8KPreqKernel: +// the amax carries its ORIGINAL element index and ties break by LOWEST index +// (`ax > amax`, never `>=`); tests assert this on tied-amax rows. +// +// The helpers here were moved verbatim out of rocm_grouped_gemm.hip's +// anonymous namespace (clean cutover, no second copy left behind); every +// consumer in that file keeps resolving the same names through this include. +#ifndef VLLM_CPP_SRC_VT_ROCM_ROCM_ACT_QUANT_H_ +#define VLLM_CPP_SRC_VT_ROCM_ROCM_ACT_QUANT_H_ + +#include +#include "vt/dtype.h" +#include "vt/cpu/cpu_quant_blocks.h" + +namespace vt::rocm { + +enum class ActDT : int { kF32 = 0, kF16 = 1, kBF16 = 2 }; + +inline ActDT ActDtOf(DType dt) { + return dt == DType::kF32 ? ActDT::kF32 : dt == DType::kF16 ? ActDT::kF16 : ActDT::kBF16; +} + +__device__ inline float DF16ToF32(uint16_t h) { + uint32_t sign = static_cast(h & 0x8000) << 16; + uint32_t exp = (h >> 10) & 0x1F; + uint32_t mant = h & 0x3FF; + if (exp == 0x1F) return __int_as_float(sign | 0x7F800000 | (mant << 13)); + if (exp == 0) { + if (mant == 0) return __int_as_float(sign); + int shift = 0; + while ((mant & 0x400) == 0) { mant <<= 1; ++shift; } + mant &= 0x3FF; + return __int_as_float(sign | ((113 - shift) << 23) | (mant << 13)); + } + return __int_as_float(sign | ((exp + 112) << 23) | (mant << 13)); +} +__device__ inline float DBF16ToF32(uint16_t b) { + return __int_as_float(static_cast(b) << 16); +} +__device__ inline uint16_t DF32ToBF16(float f) { + uint32_t u = __float_as_int(f); + if ((u & 0x7F800000) == 0x7F800000 && (u & 0x7FFFFF)) + return static_cast((u >> 16) | 0x0040); + uint32_t rounding = 0x7FFF + ((u >> 16) & 1); + return static_cast((u + rounding) >> 16); +} +__device__ inline uint16_t DF32ToF16(float f) { + uint32_t u = __float_as_uint(f); + uint16_t sign = static_cast((u >> 16) & 0x8000); + int32_t exp = static_cast((u >> 23) & 0xFF) - 127 + 15; + uint32_t mant = u & 0x7FFFFF; + if (((u >> 23) & 0xFF) == 0xFF) + return static_cast(sign | 0x7C00 | (mant ? 0x200 | (mant >> 13) : 0)); + if (exp >= 0x1F) return static_cast(sign | 0x7C00); + if (exp <= 0) { + if (exp < -10) return sign; + mant |= 0x800000; + uint32_t shift = static_cast(14 - exp); + uint32_t half = mant >> shift; + uint32_t rem = mant & ((1u << shift) - 1); + uint32_t mid = 1u << (shift - 1); + if (rem > mid || (rem == mid && (half & 1))) ++half; + return static_cast(sign | half); + } + uint32_t half = static_cast(exp << 10) | (mant >> 13); + uint32_t rem = mant & 0x1FFF; + if (rem > 0x1000 || (rem == 0x1000 && (half & 1))) ++half; + return static_cast(sign | half); +} +__device__ inline int DNearestInt(float fval) { + float val = fval + 12582912.0f; + int i = __float_as_int(val); + return (i & 0x007fffff) - 0x00400000; +} +__device__ inline float DLoadAct(const void* base, ActDT dt, int64_t idx) { + switch (dt) { + case ActDT::kF32: return static_cast(base)[idx]; + case ActDT::kF16: return DF16ToF32(static_cast(base)[idx]); + default: return DBF16ToF32(static_cast(base)[idx]); + } +} + +// Q8_K (thread-per-256-superblock): cuda_quant_dot.cu QuantizeQ8KKernel. +// The per-super-block body is factored so EVERY arm that produces Q8_K +// activation scratch (standalone grid, MMVQ LDS prologue, norm-fused +// epilogue) produces BYTE-IDENTICAL output: same amax first-occurrence +// tie-break, same scale/iscale arithmetic, same bsums walk. Asserted by +// tests/vt/test_rocm_quant_dot.cpp on random AND tied-amax inputs. +__device__ inline void QuantQ8KSBlock(vt::cpu::BlockQ8_K& y, const void* __restrict__ a, + ActDT adt, int64_t elem0) { + using vt::cpu::kQK_K; + float mx = 0.0f, amax = 0.0f; + for (int j = 0; j < kQK_K; ++j) { + const float x = DLoadAct(a, adt, elem0 + j); + if (const float ax = fabsf(x); ax > amax) { amax = ax; mx = x; } + } + if (amax == 0.0f) { + y.d = 0.0f; + for (int j = 0; j < kQK_K; ++j) y.qs[j] = 0; + for (int g = 0; g < kQK_K / 16; ++g) y.bsums[g] = 0; + return; + } + const float iscale = -127.0f / mx; + for (int j = 0; j < kQK_K; ++j) { + const int v = DNearestInt(iscale * DLoadAct(a, adt, elem0 + j)); + y.qs[j] = static_cast(v < 127 ? v : 127); + } + for (int g = 0; g < kQK_K / 16; ++g) { + int sum = 0; + for (int ii = 0; ii < 16; ++ii) sum += y.qs[g * 16 + ii]; + y.bsums[g] = static_cast(sum); + } + y.d = 1.0f / iscale; +} + +} // namespace vt::rocm + +#endif // VLLM_CPP_SRC_VT_ROCM_ROCM_ACT_QUANT_H_ diff --git a/src/vt/rocm/rocm_dense_basic.hip b/src/vt/rocm/rocm_dense_basic.hip index 7442ae63d..aebbad521 100644 --- a/src/vt/rocm/rocm_dense_basic.hip +++ b/src/vt/rocm/rocm_dense_basic.hip @@ -461,6 +461,36 @@ void CastF16KernelRocm(Queue& q, Tensor& out, const Tensor& in) { Check(hipGetLastError(), "cast_f16"); } +// T25: Permute V-heads from grouped (k*rpk+r) to tiled (r*num_k+k) order. +__global__ void PermuteVHeadsK(__hip_bfloat16* out, const __hip_bfloat16* in, + int num_k, int rpk, int dv, + int64_t value_dim) { + const int64_t row = static_cast(blockIdx.y); + const int idx = static_cast(blockIdx.x) * static_cast(blockDim.x) + + static_cast(threadIdx.x); + if (idx >= value_dim) return; + const int t = idx / dv; // tiled head index + const int h = idx % dv; // offset within head + const int r = t / num_k; + const int k = t % num_k; + const int g = k * rpk + r; // grouped head index + out[row * value_dim + static_cast(t) * dv + h] = + in[row * value_dim + static_cast(g) * dv + h]; +} + +void PermuteVHeadsKernelRocm(Queue& q, Tensor& out, const Tensor& in, + int64_t T, int64_t num_k, int64_t rpk, int64_t dv) { + const int64_t value_dim = num_k * rpk * dv; + const int threads = 128; + const int blocks_x = static_cast((value_dim + threads - 1) / threads); + dim3 grid(blocks_x, static_cast(T)); + PermuteVHeadsK<<>>( + out.Ptr<__hip_bfloat16>(), in.Ptr<__hip_bfloat16>(), + static_cast(num_k), static_cast(rpk), static_cast(dv), + value_dim); + Check(hipGetLastError(), "permute_v_heads"); +} + // ---------------- QkvSplit ---------------- template __global__ void QkvSplitK(T* q_out, T* k_out, T* v_out, const T* qkv, int64_t t, int64_t q_dim, diff --git a/src/vt/rocm/rocm_gdn_fused.hip b/src/vt/rocm/rocm_gdn_fused.hip index dff7937e1..c57956470 100644 --- a/src/vt/rocm/rocm_gdn_fused.hip +++ b/src/vt/rocm/rocm_gdn_fused.hip @@ -17,6 +17,7 @@ #include #include +#include #include #include @@ -167,7 +168,97 @@ __global__ void AttnQkNormRopeGateK(Tqk* q_out, Tqk* k_out, Tgate* gate_out, } } } - +// ── Cooperative preamble (TG200 T6b, opt-in VT_ATTN_PREAMBLE_COOP=1) ─────── +// Donor mapping prices decode at ~88us/call: items = t*(hq+hkv) = 24 at +// batch 1, so 24 of 256 threads run, each doing three serial dh-loops +// (gate copy, sumsq, rope/norm store) through private cache lines. Here one +// WARP handles one item: lanes stride j for the gate copy and the stores, +// the sumsq reduces through a fixed shfl_down tree and inv is broadcast. +// NUMERIC CONTRACT: every element's math is the donor's given inv; only the +// ss summation ORDER changes (tree vs serial), so outputs are NMSE-equal, +// not bit-exact — op is cross_device NMSE-gated, engine A/B records near-tie +// adjudication per campaign doctrine. +template +__global__ void AttnQkNormRopeGateCoopK(Tqk* q_out, Tqk* k_out, Tgate* gate_out, + const Tsrc* qgate, const Tsrc* kf, + const float* q_norm, const float* k_norm, + const float* cos_sin, int64_t t, int64_t hq, + int64_t hkv, int64_t dh, int64_t qgate_stride, + int64_t kf_stride, int rot, float eps, + bool gemma) { + constexpr int NWARPS = kBlock / 32; + const int64_t half = rot / 2; + const int64_t items = t * (hq + hkv); + const int lane = static_cast(threadIdx.x) & 31; + const int warp = static_cast(threadIdx.x) >> 5; + for (int64_t item = warp; item < items; item += NWARPS) { + const int64_t tok = item / (hq + hkv); + const int64_t h = item % (hq + hkv); + const float* cs = cos_sin + tok * rot; + if (h < hq) { + const int64_t src_off = tok * qgate_stride + h * 2 * dh; + const int64_t out_off = (tok * hq + h) * dh; + for (int64_t j = lane; j < dh; j += 32) + St(gate_out, out_off + j, Ld(qgate, src_off + dh + j)); + float ss = 0.0f; + for (int64_t j = lane; j < dh; j += 32) { + const float v = Ld(qgate, src_off + j); + ss += v * v; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) ss += __shfl_down(ss, off); + const float inv = + 1.0f / sqrtf(__shfl(ss, 0) / static_cast(dh) + eps); + for (int64_t j = lane; j < dh; j += 32) { + if (j < half) { + const float ni = GemmaNormElem(Ld(qgate, src_off + j), inv, q_norm[j], gemma); + const float nih = + GemmaNormElem(Ld(qgate, src_off + j + half), inv, q_norm[j + half], gemma); + St(q_out, out_off + j, ni * cs[j] - nih * cs[half + j]); + } else if (j < rot) { + const int64_t i = j - half; + const float ni = GemmaNormElem(Ld(qgate, src_off + i), inv, q_norm[i], gemma); + const float nih = + GemmaNormElem(Ld(qgate, src_off + i + half), inv, q_norm[i + half], gemma); + St(q_out, out_off + j, ni * cs[half + i] + nih * cs[i]); + } else { + St(q_out, out_off + j, + GemmaNormElem(Ld(qgate, src_off + j), inv, q_norm[j], gemma)); + } + } + } else { + const int64_t hk_i = h - hq; + const int64_t src_off = tok * kf_stride + hk_i * dh; + const int64_t out_off = (tok * hkv + hk_i) * dh; + float ss = 0.0f; + for (int64_t j = lane; j < dh; j += 32) { + const float v = Ld(kf, src_off + j); + ss += v * v; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) ss += __shfl_down(ss, off); + const float inv = + 1.0f / sqrtf(__shfl(ss, 0) / static_cast(dh) + eps); + for (int64_t j = lane; j < dh; j += 32) { + if (j < half) { + const float ni = GemmaNormElem(Ld(kf, src_off + j), inv, k_norm[j], gemma); + const float nih = + GemmaNormElem(Ld(kf, src_off + j + half), inv, k_norm[j + half], gemma); + St(k_out, out_off + j, ni * cs[j] - nih * cs[half + j]); + } else if (j < rot) { + const int64_t i = j - half; + const float ni = GemmaNormElem(Ld(kf, src_off + i), inv, k_norm[i], gemma); + const float nih = + GemmaNormElem(Ld(kf, src_off + i + half), inv, k_norm[i + half], gemma); + St(k_out, out_off + j, ni * cs[half + i] + nih * cs[i]); + } else { + St(k_out, out_off + j, + GemmaNormElem(Ld(kf, src_off + j), inv, k_norm[j], gemma)); + } + } + } + } +} } // namespace void RmsNormGatedKernelRocm(Queue& q, Tensor& out, const Tensor& x, const Tensor& gate, @@ -243,10 +334,25 @@ void AttnQkNormRopeGateKernelRocm(Queue& q, Tensor& q_out, Tensor& k_out, Tensor using Tsrc = decltype(src_tag); using Tqk = decltype(qk_tag); using Tgate = decltype(gate_tag); - AttnQkNormRopeGateK<<>>( - q_out.Ptr(), k_out.Ptr(), gate_out.Ptr(), qgate.Ptr(), - kf.Ptr(), q_norm.Ptr(), k_norm.Ptr(), cos_sin.Ptr(), - t, hq, hkv, dh, qgate.stride[0], kf.stride[0], ra.rotary_dim, na.eps, na.gemma); + // T6b opt-in (read once per process like the sibling arms): warp-per-item + // cooperative mapping. Default OFF keeps the donor walk. + static const bool preamble_coop = [] { + const char* e = std::getenv("VT_ATTN_PREAMBLE_COOP"); + return e != nullptr && e[0] == '1' && e[1] == '\0'; + }(); + if (preamble_coop) { + AttnQkNormRopeGateCoopK<<>>( + q_out.Ptr(), k_out.Ptr(), gate_out.Ptr(), qgate.Ptr(), + kf.Ptr(), q_norm.Ptr(), k_norm.Ptr(), cos_sin.Ptr(), + t, hq, hkv, dh, qgate.stride[0], kf.stride[0], ra.rotary_dim, na.eps, + na.gemma); + } else { + AttnQkNormRopeGateK<<>>( + q_out.Ptr(), k_out.Ptr(), gate_out.Ptr(), qgate.Ptr(), + kf.Ptr(), q_norm.Ptr(), k_norm.Ptr(), cos_sin.Ptr(), + t, hq, hkv, dh, qgate.stride[0], kf.stride[0], ra.rotary_dim, na.eps, + na.gemma); + } }; // Dispatch on the OUTPUT dtype like the CUDA lane (cuda_ops.cu // LaunchAttnPreambleOut): the caller picks f32 out (token-exact path) or bf16 diff --git a/src/vt/rocm/rocm_grouped_gemm.hip b/src/vt/rocm/rocm_grouped_gemm.hip index bcfa1c848..5107f0802 100644 --- a/src/vt/rocm/rocm_grouped_gemm.hip +++ b/src/vt/rocm/rocm_grouped_gemm.hip @@ -26,22 +26,23 @@ #include #include + #include -#include #include #include +#include #include #include #include -#include #include "vt/grow_only_stream_scratch.h" +#include "vt/rocm/rocm_act_quant.h" +#include "vt/rocm/rocm_norm_quant_bridge.h" #include "vt/ops.h" #include "vt/rocm/rocm_arch.h" #include "vt/rocm/rocm_device_bind.h" -#include "vt/rocm/rocm_f16_codec.h" -#include "vt/rocm/rocm_runtime.h" +#include "vt/rocm/rocm_runtime.h" // Block layouts — the single source of truth (ggml-common.h mirrors). #include "vt/cpu/cpu_quant_blocks.h" @@ -56,30 +57,12 @@ using vt::cpu::BlockQ6_K; using vt::cpu::kQK8_0; using vt::cpu::kQK_K; -enum class ActDT : int { kF32 = 0, kF16 = 1, kBF16 = 2 }; using Q8KKernelWitnessCount = unsigned long long; -inline ActDT ActDtOf(DType dt) { - return dt == DType::kF32 ? ActDT::kF32 : dt == DType::kF16 ? ActDT::kF16 : ActDT::kBF16; -} +// ActDT / ActDtOf and the device numeric helpers (DF16ToF32, DBF16ToF32, +// DF32ToBF16, DF32ToF16, DNearestInt, DLoadAct) live in the shared +// rocm_act_quant.h so the norm-fused epilogue TU instantiates the SAME +// byte-exactness-critical code (Lever C cutover; no second copy here). -// ---- device numeric helpers ---- -// The four half/bfloat16 conversions moved to `vt/rocm/rocm_f16_codec.h` -// (included above) when BACKEND-ROCM-EXL3 needed the same two of them for -// the trellis decode. They are unchanged, and they are still the bit-exact -// transcriptions of `vt::F16ToF32` / `vt::F32ToF16` rather than the hardware -// conversion intrinsics; a second copy is what moving them avoided. -__device__ inline int DNearestInt(float fval) { - float val = fval + 12582912.0f; - int i = __float_as_int(val); - return (i & 0x007fffff) - 0x00400000; -} -__device__ inline float DLoadAct(const void* base, ActDT dt, int64_t idx) { - switch (dt) { - case ActDT::kF32: return static_cast(base)[idx]; - case ActDT::kF16: return DF16ToF32(static_cast(base)[idx]); - default: return DBF16ToF32(static_cast(base)[idx]); - } -} __device__ __forceinline__ int GetIntB2(const int8_t* qs, int i32) { const uint16_t* x16 = reinterpret_cast(qs); return static_cast(x16[2 * i32 + 0]) | (static_cast(x16[2 * i32 + 1]) << 16); @@ -121,7 +104,12 @@ __global__ void QuantizeQ8_0K(BlockQ8_0* __restrict__ scratch, const void* __res } } -// Q8_K (thread-per-256-superblock): cuda_quant_dot.cu QuantizeQ8KKernel. +// Q8_K (thread-per-256-superblock): cuda_quant_dot.cu QuantizeQ8KKernel. The +// byte-exactness-critical per-super-block body (QuantQ8KSBlock) lives in the +// shared rocm_act_quant.h so the Lever C norm-fused epilogue instantiates the +// SAME code -- BYTE-IDENTICAL scratch by construction, asserted by +// tests/vt/test_rocm_quant_dot.cpp on random AND tied-amax inputs. + __global__ void QuantizeQ8KK(BlockQ8_K* __restrict__ scratch, const void* __restrict__ a, ActDT adt, int64_t a_rs, int64_t m, int64_t nsb, Q8KKernelWitnessCount* kernel_witness) { @@ -130,32 +118,8 @@ __global__ void QuantizeQ8KK(BlockQ8_K* __restrict__ scratch, const void* __rest if (t == 0 && kernel_witness != nullptr) atomicAdd(&kernel_witness[0], 1ULL); const int64_t i = t / nsb; const int64_t sb = t % nsb; - const int64_t elem0 = i * a_rs + sb * kQK_K; - float mx = 0.0f, amax = 0.0f; - for (int j = 0; j < kQK_K; ++j) { - const float ax = fabsf(DLoadAct(a, adt, elem0 + j)); - if (ax > amax) { amax = ax; mx = DLoadAct(a, adt, elem0 + j); } - } - BlockQ8_K& y = scratch[t]; - if (amax == 0.0f) { - y.d = 0.0f; - for (int j = 0; j < kQK_K; ++j) y.qs[j] = 0; - for (int g = 0; g < kQK_K / 16; ++g) y.bsums[g] = 0; - return; - } - const float iscale = -127.0f / mx; - for (int j = 0; j < kQK_K; ++j) { - const int v = DNearestInt(iscale * DLoadAct(a, adt, elem0 + j)); - y.qs[j] = static_cast(v < 127 ? v : 127); - } - for (int g = 0; g < kQK_K / 16; ++g) { - int sum = 0; - for (int ii = 0; ii < 16; ++ii) sum += y.qs[g * 16 + ii]; - y.bsums[g] = static_cast(sum); - } - y.d = 1.0f / iscale; + QuantQ8KSBlock(scratch[t], a, adt, i * a_rs + sb * kQK_K); } - // Q8_K (block-per-256-superblock): cuda_quant_dot.cu // QuantizeQ8KPreqKernel. One thread loads and writes one activation element. // The reduction retains the signed value and the lowest original index, which @@ -217,6 +181,85 @@ __global__ void QuantizeQ8KCooperativeK(BlockQ8_K* __restrict__ scratch, if (tid == 0) output.d = 1.0f / iscale; } +// T27 (GFX1100-TG200): warp-cooperative Q8_K quantizer. The original +// QuantizeQ8KK uses 1 thread per 256-element superblock, each doing a serial +// scan of 256 elements (~800 instructions). For decode (m=1, nsb=10) only 10 +// of 128 threads are active, and on wave32 each thread is its own wave, so +// the kernel takes ~13.4 us/call = 540 us/tok (6.0%). +// +// This version uses 8 threads per superblock (32 elements each). The amax +// scan is done per-chunk (ascending, `ax > amax` first-occurrence), then +// reduced across 8 threads via __shfl_xor_sync with lower-chunk-index +// tie-break — equivalent to a sequential scan of all 256 elements. The +// quantization (iscale = -127/mx, DNearestInt, clamp 127) and bsums are +// order-independent. Output is BYTE-IDENTICAL to QuantQ8KSBlock, asserted +// by the same gate test. +// +// Block: 128 threads = 16 superblocks. Grid: (m*nsb + 15) / 16. +// For m=1, nsb=10: 1 block, 80/128 threads active (vs 10/128), 3 waves +// of ~100 instructions (vs 10 waves of ~800) = ~8x fewer wave-cycles. +__global__ void QuantizeQ8KKWarpCoop(BlockQ8_K* __restrict__ scratch, + const void* __restrict__ a, ActDT adt, + int64_t a_rs, int64_t m, int64_t nsb) { + using vt::cpu::kQK_K; + constexpr int kChunk = 32; // elements per thread + constexpr int kThreadsPerSb = kQK_K / kChunk; // 8 + constexpr int kSbPerBlock = 128 / kThreadsPerSb; // 16 + + const int tid = threadIdx.x; + const int sb_local = tid / kThreadsPerSb; // [0, kSbPerBlock) + const int chunk = tid % kThreadsPerSb; // [0, 8) + const int64_t global_sb = + static_cast(blockIdx.x) * kSbPerBlock + sb_local; + if (global_sb >= m * nsb) return; + const int64_t i = global_sb / nsb; + const int64_t sb = global_sb % nsb; + const int64_t elem0 = i * a_rs + sb * kQK_K + chunk * kChunk; + + // --- local amax scan (ascending within chunk, `ax > amax`) --- + float amax = 0.0f, mx = 0.0f; + float vals[kChunk]; + #pragma unroll + for (int j = 0; j < kChunk; ++j) { + vals[j] = DLoadAct(a, adt, elem0 + j); + if (const float ax = fabsf(vals[j]); ax > amax) { amax = ax; mx = vals[j]; } + } + + // --- butterfly reduce across 8 threads (lower chunk index wins ties) --- + #pragma unroll + for (int off = 1; off < kThreadsPerSb; off <<= 1) { + float oa = __shfl_xor_sync(0xffffffffULL, amax, off); + float om = __shfl_xor_sync(0xffffffffULL, mx, off); + int oc = chunk ^ off; + if (oa > amax || (oa == amax && oc < chunk)) { amax = oa; mx = om; } + } + + // --- quantize and write output --- + BlockQ8_K& y = scratch[global_sb]; + if (amax > 0.0f) { + const float iscale = -127.0f / mx; + if (chunk == 0) y.d = 1.0f / iscale; + #pragma unroll + for (int j = 0; j < kChunk; ++j) { + const int qv = DNearestInt(iscale * vals[j]); + y.qs[chunk * kChunk + j] = static_cast(qv < 127 ? qv : 127); + } + int bsum0 = 0, bsum1 = 0; + #pragma unroll + for (int j = 0; j < 16; ++j) bsum0 += y.qs[chunk * kChunk + j]; + #pragma unroll + for (int j = 16; j < kChunk; ++j) bsum1 += y.qs[chunk * kChunk + j]; + y.bsums[chunk * 2] = static_cast(bsum0); + y.bsums[chunk * 2 + 1] = static_cast(bsum1); + } else { + if (chunk == 0) y.d = 0.0f; + #pragma unroll + for (int j = 0; j < kChunk; ++j) y.qs[chunk * kChunk + j] = 0; + y.bsums[chunk * 2] = 0; + y.bsums[chunk * 2 + 1] = 0; + } +} + // ---- dot superblocks (1:1 ports) ---- // Q8_0 x Q8_0: cuda_quant_dot.cu QuantDotGemmQ8_0 — dp4a int core. __device__ inline float DotQ8_0(const BlockQ8_0* wb, const BlockQ8_0* ab) { @@ -915,20 +958,7 @@ __global__ void KQuantGemmK(OutT* __restrict__ out, const uint8_t* __restrict__ float partial = 0.0f; for (int64_t sb = lane; sb < nsb; sb += 32) { const void* w_sb = w_row + static_cast(sb) * w_block_bytes; - if constexpr (Fmt == 3) { - // Fmt 3 IS Fmt 2 (Q6_K) with a different PRIVATE-MEMORY footprint, and - // nothing else. `DotQ6K` materialises the whole superblock into an - // `int8_t aux8[256]` per thread, which the compiler places in SCRATCH; - // `DotQ6KIsumRange` reconstructs 32 weights at a time into `int8_t - // aux8[32]`, which fits in registers. The arithmetic is the same integer - // accumulator times the same one f32 scale product, which is the - // bit-equality this file's `DotQ6KIsumRange` note already argues and - // which tests/vt/test_backend_cross_device.cpp already asserts between - // the cooperative and single-warp arms. - const BlockQ6_K* xb = static_cast(w_sb); - const BlockQ8_K* yb = a_row + sb; - partial += (DF16ToF32(xb->d) * yb->d) * DotQ6KIsumRange(xb, yb, 0, kSubBlocks); - } else if constexpr (Fmt == 2) partial += DotQ6K(static_cast(w_sb), a_row + sb); + if constexpr (Fmt == 2) partial += DotQ6K(static_cast(w_sb), a_row + sb); else if constexpr (Fmt == 1) partial += DotQ5K(static_cast(w_sb), a_row + sb); else partial += DotQ4K(static_cast(w_sb), a_row + sb); } @@ -1057,6 +1087,228 @@ __global__ void Q8_0GemmK(OutT* __restrict__ out, const uint8_t* __restrict__ we } } +// --- T4a (GFX1100-TG200): MMVQ-style decode GEMV arm for K-quants ---------- +// Dispatched behind VT_GEMV_MMVQ=1 with m == 1 ONLY; default OFF and the +// baseline KQuantGemmK path below is byte-unchanged. +// +// NUMERICS CONTRACT (T4a REPAIR ROUND) — the arm is BIT-EQUAL TO THE +// BASELINE KQuantGemmK OUTPUT at every (Fmt, nsb, j), not merely within the +// oracle NMSE band. History: the first round reproduced the CPU ORACLE's +// association exactly and was bit-exact vs it, yet the ENGINE degraded into +// number-loops — the extended ON-vs-OFF sweep showed the arm differs from +// the BASELINE by float ULPs (different summation association), flipping +// greedy near-ties at isolated rows (red at N=2304..248320). Repair +// mechanism, three steps: +// 1. INTEGER PHASE: an octet of 8 lanes walks one super-block in 32-elem +// chunk units (all lanes busy down to nsb=1; dp4a word cores below). +// Integer sums are EXACT under any association, so the redistribution +// cannot change them; an exact octet shuffle recovers each super- +// block's full isum (and the Q4/Q5 mins-bsums sumi). +// 2. TERM RECONSTRUCTION: the octet leader evaluates the baseline's OWN +// float expression for that super-block — d*isum (Q6_K) resp. +// d*isum - dmin*sumi (Q4_K/Q5_K) with the same d/dmin products — and +// broadcasts it. Identical inputs, identical expression => identical +// float bits as the baseline's per-super-block Dot return value. +// 3. ASSOCIATION REPLAY: every lane adds the broadcast term under the +// BASELINE's lane ownership (lane l owns super-blocks l, l+32, ... +// sequentially in increasing sb order — here: predicate (sb%32)==lane +// over the in-order sb stream) and partials close with the baseline's +// __shfl_down(16,8,4,2,1) tree. +// Identical float values in identical order => identical bits. The focused +// sweep (tests/vt/test_rocm_quant_dot.cpp, ENGINE shape set incl. lm_head +// N=248320) asserts ON==OFF byte identity at every model shape. +template +__device__ inline void KQuantGemvMmvqRow(OutT* __restrict__ out, + const uint8_t* __restrict__ w_row, + const BlockQ8_K* __restrict__ act_row, + int64_t j, int lane, int64_t nsb, + size_t w_block_bytes) { + float partial = 0.0f; + const int passes = static_cast((nsb + 3) / 4); + for (int p = 0; p < passes; ++p) { + const int sb = p * 4 + (lane >> 3); + const int c = lane & 7; + int acc = 0; // this chunk's weighted integer dot (exact) + int sumi_c = 0; // Q4/Q5: mn * (bsums pair), exact + + if (sb < nsb) { + const uint8_t* blk = + w_row + static_cast(sb) * w_block_bytes; + const BlockQ8_K& yb = act_row[sb]; + if constexpr (Fmt == 2) { // Q6_K: ql@0 qh@128 scales(int8)@192 d@208 + // Chunk c covers dequant elements [c*32, c*32+32). DotQ6K layout: + // quarter q=c>>2 holds elements [q*128,+128); segment seg=c&3 picks + // the ql low/high nibble half and the qh 2-bit field. Nibble values + // are biased by -32; the bias is removed EXACTLY in the integer + // domain with a constant-word dp4a (32*q8 per byte), so `acc` + // equals the baseline's exact isum whatever the grouping. + // seg 0/1: low nibble of ql bytes l / l+32; seg 2/3: high nibble. + // qh contributes a 2-bit field per element (mask 0x03030303). + const int q = c >> 2, seg = c & 3; + const uint8_t* ql = blk + q * 64 + (seg & 1) * 32; + const uint8_t* qh = blk + kQK_K / 2 + q * 32; + const int8_t* sc = reinterpret_cast(blk + 192); + const int shift = (seg >> 1) * 4; + const int hs = seg * 2; + const int8_t* q8 = yb.qs + c * 32; + // 0x20202020 = 32 in each byte; used for Q6_K bias correction. + // amd_mixed_dot(char4,char4,...) forces v_dot4 instruction selection; + // the scalar Dp4a emulation does NOT auto-lower when the input comes + // from a shift+mask expression (compiler cannot prove 8-bit range). + static constexpr uint32_t kBias32Word = 0x20202020u; + const char* const kBias32 = reinterpret_cast(&kBias32Word); + int sub0 = 0, pre0 = 0, sub1 = 0, pre1 = 0; +#pragma unroll + for (int w = 0; w < 8; ++w) { + uint32_t Wq, Wh, W8; + memcpy(&Wq, ql + 4 * w, 4); + memcpy(&Wh, qh + 4 * w, 4); + memcpy(&W8, q8 + 4 * w, 4); + uint32_t v = + ((Wq >> shift) & 0x0F0F0F0Fu) | + (((Wh >> hs) & 0x03030303u) << 4); + // scales change every 16 elems: words 0-3 -> sc[2c], 4-7 -> sc[2c+1] + if (w < 4) { + sub0 = amd_mixed_dot(*reinterpret_cast(&v), *reinterpret_cast(&W8), sub0, false); + pre0 = amd_mixed_dot(*reinterpret_cast(kBias32), *reinterpret_cast(&W8), pre0, false); + } else { + sub1 = amd_mixed_dot(*reinterpret_cast(&v), *reinterpret_cast(&W8), sub1, false); + pre1 = amd_mixed_dot(*reinterpret_cast(kBias32), *reinterpret_cast(&W8), pre1, false); + } + } + acc = static_cast(sc[2 * c]) * (sub0 - pre0) + + static_cast(sc[2 * c + 1]) * (sub1 - pre1); + } else { // Q4_K / Q5_K: d@0 dmin@2 scales@4 [qh@16] qs@16|48 + uint32_t W0, W1, W2; + memcpy(&W0, blk + 4, 4); + memcpy(&W1, blk + 8, 4); + memcpy(&W2, blk + 12, 4); + int scale, mn; + if (c < 4) { + scale = static_cast((W0 >> (8 * c)) & 0x3F); + mn = static_cast((W1 >> (8 * c)) & 0x3F); + } else { + const int i4 = 8 * (c - 4); + scale = static_cast(((W2 >> i4) & 0xF) | + (((W0 >> (i4 + 6)) & 3) << 4)); + mn = static_cast(((W2 >> (i4 + 4)) & 0xF) | + (((W1 >> (i4 + 6)) & 3) << 4)); + } + const uint8_t* qb = blk + (Fmt == 1 ? 48 : 16) + (c >> 1) * 32; + const int shift = (c & 1) ? 4 : 0; + const uint8_t* hm = blk + 16; + const int8_t* q8 = yb.qs + c * 32; + int sub = 0; +#pragma unroll + for (int w = 0; w < 8; ++w) { + uint32_t Wq, W8; + memcpy(&Wq, qb + 4 * w, 4); + memcpy(&W8, q8 + 4 * w, 4); + uint32_t v = (Wq >> shift) & 0x0F0F0F0Fu; + if constexpr (Fmt == 1) { + uint32_t Wh; + memcpy(&Wh, hm + 4 * w, 4); + v |= ((Wh >> c) & 0x01010101u) << 4; + } + sub = amd_mixed_dot(*reinterpret_cast(&v), *reinterpret_cast(&W8), sub, false); + } + acc = scale * sub; + sumi_c = mn * (yb.bsums[2 * c] + yb.bsums[2 * c + 1]); + } + } +#pragma unroll + for (int off = 4; off > 0; off >>= 1) { + acc += __shfl_down_sync(0xffffffffULL, acc, off); + sumi_c += __shfl_down_sync(0xffffffffULL, sumi_c, off); + } + // Fold THIS pass's super-blocks: reconstruct the baseline's exact float + // term on the octet leader (after the reduce above, lane 8*kk holds the + // whole-octet integer totals) and add it under the baseline's lane + // ownership (lane l owns sbs l, l+32, ... — the stream here is strictly + // sb-increasing, so per-lane accumulation order matches). + const int rem = static_cast(nsb - p * 4); + const int nact = rem < 4 ? rem : 4; + for (int kk = 0; kk < nact; ++kk) { + const int sbk = p * 4 + kk; + const int isum = __shfl_sync(0xffffffffULL, acc, 8 * kk); + const int SUMI = __shfl_sync(0xffffffffULL, sumi_c, 8 * kk); + float term = 0.0f; + if (lane == 8 * kk) { + const uint8_t* blk = + w_row + static_cast(sbk) * w_block_bytes; + uint16_t dh; + if constexpr (Fmt == 2) memcpy(&dh, blk + 208, 2); else memcpy(&dh, blk, 2); + const float d = DF16ToF32(dh) * act_row[sbk].d; + if constexpr (Fmt != 2) { // baseline: d * isum - dmin * sumi + uint16_t dmh; + memcpy(&dmh, blk + 2, 2); + const float dmin = DF16ToF32(dmh) * act_row[sbk].d; + term = d * static_cast(isum) - dmin * static_cast(SUMI); + } else { // baseline: d * isum + term = d * static_cast(isum); + } + } + const float P = __shfl_sync(0xffffffffULL, term, 8 * kk); + if ((sbk & 31) == lane) partial += P; + } + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) partial += __shfl_down_sync(0xffffffffULL, partial, off); + if (lane == 0) { + if constexpr (sizeof(OutT) == 4) out[j] = partial; + else out[j] = DF32ToBF16(partial); + } +} + +template +__global__ void KQuantGemvMmvqK(OutT* __restrict__ out, + const uint8_t* __restrict__ weight, + const BlockQ8_K* __restrict__ act, int64_t n, + int64_t nsb, size_t w_block_bytes) { + const int64_t j = + static_cast(blockIdx.x) * blockDim.y + threadIdx.y; + if (j >= n) return; + const uint8_t* w_row = + weight + static_cast(j) * static_cast(nsb) * w_block_bytes; + KQuantGemvMmvqRow(out, w_row, act, j, static_cast(threadIdx.x), + nsb, w_block_bytes); +} + +// Fused-prologue variant (the operator-steered shape of the arm): the block +// quantizes the activation row into LDS FIRST — each linear thread takes +// whole super-blocks in the SAME thread-per-super-block walk as the +// standalone QuantizeQ8KK grid, via the SHARED QuantQ8KSBlock body, so the +// bytes are identical by construction — then __syncthreads() and runs the +// identical GEMV row body against the LDS copy. This deletes the standalone +// ~59us quant launch (T4a capture: 43.7 launches/tok on grids of <=1 +// block). The redundant per-block requantization is bounded by the +// kMmvqFoldMaxRows gate at the dispatch site: measured on gfx1100, the fold +// beats the launch at n<=512 and loses badly beyond (the redundancy scales +// with the block count) — see the gate comment there. +template +__global__ void KQuantGemvMmvqFusedK(OutT* __restrict__ out, + const uint8_t* __restrict__ weight, + const void* __restrict__ a, ActDT adt, + int64_t n, int64_t nsb, + size_t w_block_bytes) { + extern __shared__ unsigned char vt_mmvq_smem[]; + BlockQ8_K* lds = + reinterpret_cast(vt_mmvq_smem); // [nsb] blocks + const int tid = static_cast(threadIdx.y) * 32 + + static_cast(threadIdx.x); + for (int64_t sb = tid; sb < nsb; sb += 128) + QuantQ8KSBlock(lds[sb], a, adt, sb * kQK_K); // m == 1: elem0 = sb*256 + __syncthreads(); + const int64_t j = + static_cast(blockIdx.x) * blockDim.y + threadIdx.y; + if (j >= n) return; // only AFTER the barrier: whole block quantizes first + const uint8_t* w_row = + weight + static_cast(j) * static_cast(nsb) * w_block_bytes; + KQuantGemvMmvqRow(out, w_row, lds, j, + static_cast(threadIdx.x), nsb, + w_block_bytes); +} + inline void Check(hipError_t err, const char* what) { if (err != hipSuccess) throw std::runtime_error(std::string("vt rocm grouped_gemm: ") + what + ": " + @@ -1092,7 +1344,6 @@ void* EnsureQuantScratch(size_t need, hipStream_t s) { return block; }); } - enum class Q8KQuantArm : uint8_t { kLegacy = 0, kCandidate = 1 }; enum class Q8KProductionRoute : uint8_t { kDense = 0, kGrouped = 1 }; using Q8KArchResolver = std::string (*)(int) noexcept; @@ -1222,7 +1473,6 @@ bool QuantWmmaEnabled() { }(); return on; } - // VT_ROCM_Q6K_SMALL_PRIVATE=1 — the A/B arm for // [#2511](https://github.com/mudler/vllm.cpp/issues/2511). Every failing leg on // `gfx1151` ended inside `KQuantGemmK`, the Q6_K arm, with all three of @@ -1246,8 +1496,38 @@ bool Q6KSmallPrivateEnabled() { return on; } + } // namespace +// T4a dispatch-route counters. Incremented on the decode hot path; read by +// the focused test's byte-equality gate through the quant scratch hook. +struct MmvqRouteCounts { + long long baseline; // KQuantGemmK warp-reduction dispatches + long long gemv_mmvq; // non-fused MMVQ GEMV dispatches (standalone quant) + long long gemv_fused; // fused-fold sub-branch dispatches +}; + +// T4a dispatch-route counters. Incremented on the decode hot path; read by +// the focused test's byte-equality gate through the quant scratch hook. +// (Restored: the gamma-gate commit's diff context swallowed the definitions +// that T21's point carried — the reset/read helpers and the decode-path +// increments below reference them.) +std::atomic g_mmvq_route_baseline{0}; +std::atomic g_mmvq_route_gemv{0}; +std::atomic g_mmvq_route_fused{0}; + +void MmvqResetRouteCountsForTesting() { + g_mmvq_route_baseline.store(0, std::memory_order_relaxed); + g_mmvq_route_gemv.store(0, std::memory_order_relaxed); + g_mmvq_route_fused.store(0, std::memory_order_relaxed); +} + +MmvqRouteCounts MmvqRouteCountsForTesting() { + return {g_mmvq_route_baseline.load(std::memory_order_relaxed), + g_mmvq_route_gemv.load(std::memory_order_relaxed), + g_mmvq_route_fused.load(std::memory_order_relaxed)}; +} + // Warps that cooperate on one output row of the dense K-quant GEMM. 1 selects // the single-warp `KQuantGemmK` the path has always had. // @@ -1320,7 +1600,129 @@ void* Q8KSetKernelExecutionWitnessForTest(void* device_counts) { return previous; } -void MatmulBTQuantKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) { +// T4a testing hook: fill dst (device ptr, k/256 BlockQ8_K blocks for row a) +// with the standalone-quantizer semantics (mode 0: the production +// QuantizeQ8KK grid) or the fused-prologue semantics (mode 1: ONE block, +// cooperative LDS phase then copy-out) so tests can assert BYTE equality +// between the two on arbitrary — including tied-amax — inputs. +namespace { +__global__ void MmvqPrologueQuantForTestingK(BlockQ8_K* __restrict__ dst, + const void* __restrict__ a, + ActDT adt, int64_t nsb) { + extern __shared__ unsigned char vt_mmq_smem[]; + BlockQ8_K* lds = reinterpret_cast(vt_mmq_smem); + const int tid = static_cast(threadIdx.y) * 32 + + static_cast(threadIdx.x); + for (int64_t sb = tid; sb < nsb; sb += 128) + QuantQ8KSBlock(lds[sb], a, adt, sb * kQK_K); + __syncthreads(); + for (int64_t sb = tid; sb < nsb; sb += 128) dst[sb] = lds[sb]; +} +} // namespace +void MmvqQuantScratchForTesting(Queue& q, void* dst, const Tensor& a, + bool fused_semantics) { + EnsureQueueDevice(q); + hipStream_t s = static_cast(q.handle); + const int64_t k = a.shape[1], nsb = k / kQK_K; + if (a.shape[0] != 1 || k % kQK_K != 0) + throw std::runtime_error("vt rocom mmvq test hook: need one whole row"); + if (!fused_semantics) { + QuantizeQ8KK<<((nsb + 127) / 128), 128, 0, s>>>( + static_cast(dst), a.data, ActDtOf(a.dtype), a.stride[0], 1, + nsb, nullptr); + } else { + MmvqPrologueQuantForTestingK<<<1, dim3(32, 4), + static_cast(nsb) * sizeof(BlockQ8_K), + s>>>(static_cast(dst), a.data, + ActDtOf(a.dtype), nsb); + } + Check(hipGetLastError(), "mmvq quant scratch hook"); +} + +// --- Lever C (GFX1100-TG200-NORMQ): producer-fused Q8_K norm epilogue bridge +// (see rocm_norm_quant_bridge.h for the contract). Single-slot latest-producer +// token; matching K-quant consumers skip the standalone QuantizeQ8KK, any +// non-matching consumer invalidates the token. Counters are host-side +// dispatch-route witnesses exactly like the Mmvq route counters above: they +// advance once per capture-time dispatch and never per replay iteration. +namespace { + +struct NormQuantToken { + const void* a_ptr = nullptr; + int64_t rows = 0, h = 0, row_stride = 0; + DType adt = DType::kBF16; + const void* scratch = nullptr; + hipStream_t stream = nullptr; +}; +std::mutex g_nq_mu; +NormQuantToken g_nq_token; +bool g_nq_valid = false; +std::atomic g_nq_producers{0}; +std::atomic g_nq_consumers_fused{0}; +std::atomic g_nq_consumers_standalone{0}; +const void* g_nq_last_scratch = nullptr; + +} // namespace + +void* NormQuantProducerScratch(size_t bytes, void* stream) { + return EnsureQuantScratch(bytes, static_cast(stream)); +} + +void NormQuantRecordProducer(const void* out_ptr, int64_t rows, int64_t h, + int64_t row_stride, DType adt, const void* scratch, + void* stream) { + std::lock_guard lk(g_nq_mu); + g_nq_token = NormQuantToken{out_ptr, rows, h, row_stride, adt, scratch, + static_cast(stream)}; + g_nq_valid = true; + g_nq_last_scratch = scratch; + g_nq_producers.fetch_add(1, std::memory_order_relaxed); +} + +bool NormQuantTakeConsumer(const void* a_ptr, int64_t rows, int64_t h, + int64_t row_stride, DType adt, void* stream, + const void** scratch_out) { + std::lock_guard lk(g_nq_mu); + if (g_nq_valid && g_nq_token.a_ptr == a_ptr && g_nq_token.rows == rows && + g_nq_token.h == h && g_nq_token.row_stride == row_stride && + g_nq_token.adt == adt && + g_nq_token.stream == static_cast(stream)) { + *scratch_out = g_nq_token.scratch; + // token STAYS valid: the attn q/k/v matvecs consume ONE normalized row + // three times; only a non-matching consumer invalidates it. + g_nq_consumers_fused.fetch_add(1, std::memory_order_relaxed); + return true; + } + g_nq_valid = false; + g_nq_consumers_standalone.fetch_add(1, std::memory_order_relaxed); + return false; +} + +NormQuantCounts NormQuantCountsForTesting() { + return {g_nq_producers.load(std::memory_order_relaxed), + g_nq_consumers_fused.load(std::memory_order_relaxed), + g_nq_consumers_standalone.load(std::memory_order_relaxed)}; +} + +void NormQuantResetForTesting() { + std::lock_guard lk(g_nq_mu); + g_nq_token = NormQuantToken{}; + g_nq_valid = false; + g_nq_last_scratch = nullptr; + g_nq_producers.store(0, std::memory_order_relaxed); + g_nq_consumers_fused.store(0, std::memory_order_relaxed); + g_nq_consumers_standalone.store(0, std::memory_order_relaxed); +} + +const void* NormQuantLastScratchForTesting() { + std::lock_guard lk(g_nq_mu); + return g_nq_last_scratch; +} + + + + +void MatmulBTQuantKernelRocmGdn(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) { EnsureQueueDevice(q); const int64_t m = a.shape[0], k = a.shape[1], n = b.shape[0]; if (m == 0 || n == 0) return; @@ -1351,12 +1753,132 @@ void MatmulBTQuantKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tenso : b.dtype == DType::kQ5_K ? sizeof(BlockQ5_K) : sizeof(BlockQ4_K); const size_t w_row_bytes = static_cast(nsb) * w_block_bytes; - BlockQ8_K* qact = static_cast(EnsureQuantScratch( - static_cast(m) * nsb * sizeof(BlockQ8_K), s)); - LaunchQ8KQuantizer(qact, a.data, ActDtOf(a.dtype), a.stride[0], m, nsb, s, - Q8KProductionRoute::kDense, q.device.index); - const int64_t grid = (m * n + kWarpsPerBlock - 1) / kWarpsPerBlock; const int fmt = b.dtype == DType::kQ6_K ? 2 : b.dtype == DType::kQ5_K ? 1 : 0; + // T4a MMVQ-style decode GEMV arm (VT_GEMV_MMVQ=1, default OFF; m==1 + // only). Flag read PER CALL (the CUDA sibling's convention, + // cuda_quant_dot.cu:1006) so in-process tests and a captured decode + // graph pick the arm up at launch/capture time. The baseline path below + // is untouched when the flag is absent. + // REPAIR-ROUND GATE FIX: the m==1 restriction must gate the WHOLE arm, + // not just the LDS fold. The first round gated only `gemv_fused`, so + // multi-row calls (the engine's m=39 PREFILL chunks) fell into the + // non-fused GEMV branch, whose kernels write row 0 only — rows 1..m-1 + // of the output were left UNWRITTEN (stale memory) while every op-level + // test (all m==1) stayed green: the exact op-green/engine-garbage + // pattern of round 1, red-proven by the MULTI-M byte-identity case in + // tests/vt/test_rocm_quant_dot.cpp (canary-filled outputs). + const char* mmvq_e = std::getenv("VT_GEMV_MMVQ"); + const bool gemv_mmvq = + mmvq_e != nullptr && mmvq_e[0] == '1' && mmvq_e[1] == '\0' && + m == 1; + constexpr size_t kMmvqFoldLdsBytes = 32 * 1024; + // The fold trades the standalone quant launch for a PER-BLOCK redundant + // requant of the activation row, so its cost grows with the block count + // (n/4) while the launch it deletes costs a fixed few us. Measured + // crossover on gfx1100 at this model's shapes: fold WINS at n=320 + // (grid 80: 135us vs 148us) and LOSES by n=2304 (grid 576: 159us vs + // 123us); at the lm_head class the fused arm ran 2.2-3.3x SLOWER than + // baseline while the non-fused arm runs 0.46-0.61x. Gate the fold to + // small rows only; everything else takes standalone quant + GEMV. + // LEVER B1 (GFX1100-TG200): the 512-row crossover is RUNTIME-TUNABLE via + // VT_GEMV_MMVQ_FOLD_MAX so the NET-WIN point can be measured END-TO-END. + // Integer rows; default = kMmvqFoldMaxRowsDefault below; empty/invalid + // (= non-integer, <=0, trailing garbage) falls back to the default. Read + // PER CALL like VT_GEMV_MMVQ so in-process tests and graph capture pick + // the value up at dispatch time. Suite pins: tests/vt/ + // test_rocm_quant_dot.cpp F2 (default routing) + F3 (knob witness). + constexpr int64_t kMmvqFoldMaxRowsDefault = 512; + int64_t mmvq_fold_max_rows = kMmvqFoldMaxRowsDefault; + if (const char* fm_e = std::getenv("VT_GEMV_MMVQ_FOLD_MAX")) { + char* fm_end = nullptr; + const long long fm_v = std::strtoll(fm_e, &fm_end, 10); + if (fm_end != fm_e && *fm_end == '\0' && fm_v > 0) + mmvq_fold_max_rows = static_cast(fm_v); + } + const bool gemv_fused = + gemv_mmvq && m == 1 && n <= mmvq_fold_max_rows && + static_cast(nsb) * sizeof(BlockQ8_K) <= kMmvqFoldLdsBytes; + + auto launch_mvq = [&](auto ot, BlockQ8_K* qact) { + using OutT = decltype(ot); + auto* o = static_cast(out.data); + constexpr int kGemvWarps = 4; + const int64_t ggrid = (n + kGemvWarps - 1) / kGemvWarps; + dim3 gblock(32, kGemvWarps); + if (gemv_fused) { + const size_t lds_bytes = static_cast(nsb) * sizeof(BlockQ8_K); + if (fmt == 2) + KQuantGemvMmvqFusedK<<(ggrid), gblock, + lds_bytes, s>>>( + o, w, a.data, ActDtOf(a.dtype), n, nsb, w_block_bytes); + else if (fmt == 1) + KQuantGemvMmvqFusedK<<(ggrid), gblock, + lds_bytes, s>>>( + o, w, a.data, ActDtOf(a.dtype), n, nsb, w_block_bytes); + else + KQuantGemvMmvqFusedK<<(ggrid), gblock, + lds_bytes, s>>>( + o, w, a.data, ActDtOf(a.dtype), n, nsb, w_block_bytes); + } else { + if (fmt == 2) + KQuantGemvMmvqK<<(ggrid), gblock, 0, s>>>( + o, w, qact, n, nsb, w_block_bytes); + else if (fmt == 1) + KQuantGemvMmvqK<<(ggrid), gblock, 0, s>>>( + o, w, qact, n, nsb, w_block_bytes); + else + KQuantGemvMmvqK<<(ggrid), gblock, 0, s>>>( + o, w, qact, n, nsb, w_block_bytes); + } + }; + + // FUSED arm: no standalone quant launch, no scratch — the GEMV blocks + // quantize the row into their own LDS via the SHARED QuantQ8KSBlock body + // (byte-identical to QuantizeQ8KK's output; asserted by the focused test). + if (gemv_fused) { + ++g_mmvq_route_fused; + if (out.dtype == DType::kF32) launch_mvq(float{}, nullptr); + else launch_mvq(uint16_t{}, nullptr); + Check(hipGetLastError(), "K-quant gemv mmvq fused"); + return; + } + + // Lever C: when a producer-fused norm epilogue (VT_NORM_QUANT_FUSED=1) + // already wrote this activation's Q8_K scratch, the standalone launch is + // SKIPPED -- that is the launch the lever deletes. Byte equality vs the + // standalone path holds by construction (shared QuantQ8KSBlock on the + // same rows) and is asserted op-level in tests/vt/test_rocm_quant_dot.cpp. + BlockQ8_K* qact = nullptr; + const void* fused_scratch = nullptr; + if (NormQuantTakeConsumer(a.data, m, k, a.stride[0], a.dtype, s, + &fused_scratch)) { + qact = static_cast(const_cast(fused_scratch)); + } else { + qact = static_cast(EnsureQuantScratch( + static_cast(m) * nsb * sizeof(BlockQ8_K), s)); + const char* wc_e = std::getenv("VT_QUANT_Q8K_WARP"); + const bool use_warp = wc_e != nullptr && wc_e[0] == '1' && wc_e[1] == '\0'; + if (use_warp) { + constexpr int kSbPerBlock = 16; + QuantizeQ8KKWarpCoop<<((m * nsb + kSbPerBlock - 1) / kSbPerBlock), 128, 0, s>>>( + qact, a.data, ActDtOf(a.dtype), a.stride[0], m, nsb); + } else { + LaunchQ8KQuantizer(qact, a.data, ActDtOf(a.dtype), a.stride[0], m, nsb, s, + Q8KProductionRoute::kDense, q.device.index); + } + Check(hipGetLastError(), "q8_K quant"); + } + + if (gemv_mmvq) { // non-fused arm (nsb too large for the LDS fold) + ++g_mmvq_route_gemv; + if (out.dtype == DType::kF32) launch_mvq(float{}, qact); + else launch_mvq(uint16_t{}, qact); + Check(hipGetLastError(), "K-quant gemv mmvq"); + return; + } + + ++g_mmvq_route_baseline; + const int64_t grid = (m * n + kWarpsPerBlock - 1) / kWarpsPerBlock; const int coop = KQuantDecodeCoopWarps(b.dtype, m, nsb); // RDNA4 WMMA tile arm: Q6_K and Q4_K prefill only (m > 1, since // KQuantDecodeCoopWarps above already owns m == 1). Q5_K is `## Owed` — @@ -1446,14 +1968,14 @@ void MatmulBTQuantKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tenso Check(hipGetLastError(), "K-quant gemm"); return; } - throw std::runtime_error("vt rocm: matmul_bt_quant: unsupported weight dtype (ported: Q8_0/Q4_K/Q5_K/Q6_K; owed: Q4_0/Q2_K/Q3_K/IQ2_XXS/IQ3_XXS/IQ3_S/IQ2_S/MXFP4 -- the loader pre-filters to the ported set, so reaching here is a bug)"); + throw std::runtime_error("vt rocm: matmul_bt_quant: unsupported weight dtype (ported: Q8_0/Q4_K/Q5_K/Q6_K; owed: Q4_0/Q2_K/Q3_K/IQ2_XXS/IQ3_XXS/IQ2_S/MXFP4 -- the loader pre-filters to the ported set, so reaching here is a bug)"); } // kMatmulBTQuantGrouped for ROCm: Q8_0 / Q4_K / Q6_K natively (the formats the // target GDN-MoE GGUFs carry); anything else throws loudly (never a silent // CPU-pointer deref on a discrete card). -void MatmulBTQuantGroupedKernelRocm(Queue& q, Tensor& out, const Tensor& act, +void MatmulBTQuantGroupedKernelRocmGdn(Queue& q, Tensor& out, const Tensor& act, const Tensor& weight, const Tensor& expert_ids) { EnsureQueueDevice(q); const int64_t P = out.shape[0], n = out.shape[1], k = act.shape[1]; @@ -1500,6 +2022,7 @@ void MatmulBTQuantGroupedKernelRocm(Queue& q, Tensor& out, const Tensor& act, static_cast(Pa) * nsb * sizeof(BlockQ8_K), s)); LaunchQ8KQuantizer(qact, act.data, ActDtOf(act.dtype), act.stride[0], Pa, nsb, s, Q8KProductionRoute::kGrouped, q.device.index); + Check(hipGetLastError(), "q8_K quant"); const int64_t grid = (P * n + kWarpsPerBlock - 1) / kWarpsPerBlock; const int fmt = weight.dtype == DType::kQ6_K ? 2 : weight.dtype == DType::kQ5_K ? 1 : 0; auto launch = [&](auto ot) { diff --git a/src/vt/rocm/rocm_matmul_hipblaslt.hip b/src/vt/rocm/rocm_matmul_hipblaslt.hip index 77deaea03..b336ed270 100644 --- a/src/vt/rocm/rocm_matmul_hipblaslt.hip +++ b/src/vt/rocm/rocm_matmul_hipblaslt.hip @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -529,11 +530,49 @@ void MatmulKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) { "hipblasGemmEx NN"); } +// Host-side routing witness for the bf16-in/f32-out decode population (the +// Qwen3.5 GDN BA pair; evidence 15.1): process-global counters bumped on +// exactly the branch taken per dispatch. Both routes are numerically valid, +// so outputs alone cannot witness routing. Same shape as the T4a MMVQ +// counters (rocm_grouped_gemm.hip). +struct SkinnyF32RouteCounts { + long long blas; // fell through to hipblasGemmEx (default route) + long long skinny; // took the VT_SKINNY_BF16 wvSplitK-class arm +}; + +namespace { +std::atomic g_skinny_f32_route_blas{0}; +std::atomic g_skinny_f32_route_skinny{0}; +} // namespace + +void SkinnyF32ResetRouteCountsForTesting() { + g_skinny_f32_route_blas.store(0, std::memory_order_relaxed); + g_skinny_f32_route_skinny.store(0, std::memory_order_relaxed); +} + +SkinnyF32RouteCounts SkinnyF32RouteCountsForTesting() { + return {g_skinny_f32_route_blas.load(std::memory_order_relaxed), + g_skinny_f32_route_skinny.load(std::memory_order_relaxed)}; +} + +// Lever B2 opt-in arm (evidence 15.1): VT_SKINNY_BF16=1 serves bf16-in/ +// f32-out decode-skinny shapes with the wvSplitK geometry instead of +// rocBLAS's large-M tile. Read PER CALL (cuda_quant_dot.cu convention) so +// in-process tests and captured graphs pick the arm up at launch time. +// Default OFF: the default path is byte-unchanged. +bool SkinnyBf16F32OutEnabled() { + if (const char* e = std::getenv("VT_SKINNY_BF16")) return e[0] == '1'; + return false; +} + // out[M,N] = a[M,K] @ b[N,K]^T // wvSplitK skinny-GEMM host entry (rocm_skinny_gemm.hip, #487). External // vt::rocm linkage to match the definition; declared beside its only caller. void WvSplitKBT(hipStream_t s, void* out, const void* a, const void* b, int M, int N, int K, int device); +// Lever B2 f32-output variant of the same kernel geometry. +void WvSplitKBTToF32(hipStream_t s, void* out, const void* a, const void* b, + int M, int N, int K, int device); // Row-major trick: gemm(OP_T, OP_N, N, M, K, B, K, A, a_rs, C, N) // BLAS: C = op(A)*op(B) with opA=T => A is KxN in col form = row B[N,K] @@ -583,6 +622,25 @@ void MatmulBTKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) return; } + // Lever B2 (evidence 15.1): bf16-in/f32-out decode-skinny — the Qwen3.5 + // GDN BA projections (N=32, K=2560, m=1) emit f32 and today starve on + // rocBLAS's large-M MT128x32x16 tile (~73.7us for a 164 KiB weight). + // Same donor guards as the bf16 arm above (the kernel is the same + // geometry, so the same tail-safety constraints apply), opt-in via + // VT_SKINNY_BF16=1, default OFF and byte-unchanged. + const bool skinny_f32_pop = + bf16 && out.dtype == DType::kF32 && M >= 1 && M <= 4; + if (skinny_f32_pop && SkinnyBf16F32OutEnabled() && (K % 8) == 0 && N > 8 && + (N % 2) == 0 && a.stride[0] == K && K * M <= 32768 && + vt::rocm::SkinnyGemmArchOk(q.device.index, vt::rocm::DeviceArchName)) { + WvSplitKBTToF32(s, out.data, a.data, b.data, static_cast(M), + static_cast(N), static_cast(K), q.device.index); + g_skinny_f32_route_skinny.fetch_add(1, std::memory_order_relaxed); + return; + } + if (skinny_f32_pop) + g_skinny_f32_route_blas.fetch_add(1, std::memory_order_relaxed); + // Decode: M=1 BF16 GEMV if (M == 1 && bf16 && out.dtype == DType::kBF16 && a.stride[0] == K && GemvEnabled()) { Bf16GemvBT(s, out.data, a.data, b.data, static_cast(N), static_cast(K), 1.f, 0.f); diff --git a/src/vt/rocm/rocm_norm_quant_bridge.h b/src/vt/rocm/rocm_norm_quant_bridge.h new file mode 100644 index 000000000..c3539215d --- /dev/null +++ b/src/vt/rocm/rocm_norm_quant_bridge.h @@ -0,0 +1,57 @@ +// Lever C (GFX1100-TG200-NORMQ): host-side bridge for the producer-fused Q8_K +// norm epilogue. Implemented in src/vt/rocm/rocm_grouped_gemm.hip (which owns +// the activation-quant scratch pool and the standalone QuantizeQ8KK launch); +// consumed by src/vt/rocm/rocm_rmsnorm.hip (the producer side). +// +// Contract (VT_NORM_QUANT_FUSED=1, opt-in; default OFF leaves every path +// byte-unchanged): +// 1. A producer dispatching an epilogue-enabled RmsNormRowKernel allocates +// Q8_K scratch from the SAME grow-only stream-ordered pool the consumer +// uses, launches the kernel with the epilogue pointer, and RECORDS a +// single-slot token {out ptr, rows, h, dtype, scratch, stream}. +// 2. A MatmulBTQuant K-quant dispatch whose activation EXACTLY matches the +// recorded token (same device pointer, rows, row length, stride, input +// dtype) SKIPS its standalone QuantizeQ8KK launch and consumes the +// produced scratch. The token survives matching consumers (the model's +// attn q/k/v matvecs re-quantize ONE normalized row three times) and is +// INVALIDATED by any non-matching K-quant consumer, so a stale token can +// never serve a different buffer. +// Stream-ordering argument: producer and consumer are enqueued on one +// stream, and the epilogue quantizes the same global bf16 rows the +// standalone kernel would read, through the SAME shared QuantQ8KSBlock body +// -- byte equality holds by construction (asserted op-level in +// tests/vt/test_rocm_quant_dot.cpp). Under hipGraph capture both sides run +// at capture time, so the baked graph references the retired-never scratch +// pointer exactly like the pre-existing pool discipline. +#ifndef VLLM_CPP_SRC_VT_ROCM_ROCM_NORM_QUANT_BRIDGE_H_ +#define VLLM_CPP_SRC_VT_ROCM_ROCM_NORM_QUANT_BRIDGE_H_ + +#include + +#include "vt/dtype.h" + +namespace vt::rocm { + +// Producer side: scratch of `bytes` from the quant pool on `s`, then record. +void* NormQuantProducerScratch(size_t bytes, void* stream); +void NormQuantRecordProducer(const void* out_ptr, int64_t rows, int64_t h, + int64_t row_stride, DType adt, const void* scratch, + void* stream); +// Consumer side: true + scratch when the activation matches the live token; +// false otherwise (and any non-matching query invalidates the token). +bool NormQuantTakeConsumer(const void* a_ptr, int64_t rows, int64_t h, + int64_t row_stride, DType adt, void* stream, + const void** scratch_out); + +struct NormQuantCounts { + long long producers; // epilogue-enabled RmsNorm dispatches + long long consumers_fused; // K-quant dispatches that skipped the standalone quant + long long consumers_standalone; // K-quant dispatches that launched QuantizeQ8KK +}; +NormQuantCounts NormQuantCountsForTesting(); +void NormQuantResetForTesting(); +const void* NormQuantLastScratchForTesting(); + +} // namespace vt::rocm + +#endif // VLLM_CPP_SRC_VT_ROCM_ROCM_NORM_QUANT_BRIDGE_H_ diff --git a/src/vt/rocm/rocm_ops.hip b/src/vt/rocm/rocm_ops.hip index 7201245ef..4a534bfc6 100644 --- a/src/vt/rocm/rocm_ops.hip +++ b/src/vt/rocm/rocm_ops.hip @@ -45,6 +45,8 @@ void CastF32KernelRocm(Queue& q, Tensor& out, const Tensor& in); // The NARROWING third sibling (rocm_dense_basic.hip). BACKEND-ROCM-EXL3 / #2433: // one of the exactly two ops an EXL3 checkpoint ran on the CPU reference tier. void CastF16KernelRocm(Queue& q, Tensor& out, const Tensor& in); +void PermuteVHeadsKernelRocm(Queue& q, Tensor& out, const Tensor& in, + int64_t T, int64_t num_k, int64_t rpk, int64_t dv); void QkvSplitKernelRocm(Queue& q, Tensor& q_out, Tensor& k_out, Tensor& v_out, const Tensor& qkv); void LayerNormKernelRocm(Queue& q, Tensor& out, const Tensor& x, const Tensor* weight, @@ -189,6 +191,9 @@ struct Registrar { reinterpret_cast(static_cast(&CastBf16KernelRocm))); RegisterOp(OpId::kCastF16, DeviceType::kROCM, reinterpret_cast(static_cast(&CastF16KernelRocm))); + RegisterOp(OpId::kPermuteVHeads, DeviceType::kROCM, + reinterpret_cast( + static_cast(&PermuteVHeadsKernelRocm))); RegisterOp(OpId::kCastF32, DeviceType::kROCM, reinterpret_cast(static_cast(&CastF32KernelRocm))); RegisterOp(OpId::kQkvSplit, DeviceType::kROCM, diff --git a/src/vt/rocm/rocm_quant_dot.hip b/src/vt/rocm/rocm_quant_dot.hip new file mode 100644 index 000000000..e9556af9c --- /dev/null +++ b/src/vt/rocm/rocm_quant_dot.hip @@ -0,0 +1,1066 @@ +// ROCm keep-quant GGUF k-quant GEMM (KERNEL-QUANT-CIQ-GEMM-ROCM W1) — the +// kROCM provider for `OpId::kMatmulBTQuant` and `OpId::kMatmulBTQuantGrouped`. +// +// Port of src/vt/cuda/cuda_quant_dot.cu (the kCUDA provider), which is itself +// a port of the CPU oracle: +// src/vt/cpu/cpu_quant_gemm.cpp MatmulBTQuantKernel (the GEMM wiring) +// src/vt/cpu/cpu_quant_dot.cpp VecDot{Q2_K,Q3_K,Q4_K,Q5_K,Q6_K, +// IQ2_XXS,IQ3_XXS,IQ2_S,IQ1_S,IQ1_XXXS}Q8_K +// src/vt/cpu/cpu_quant_act.cpp QuantizeRowQ8_K (the activation quant) +// The device numeric helpers below are the SAME bit-exact ports of +// src/vt/dtype.cpp + cpu_quant_act.cpp, so the Q8_K activation bytes — and +// therefore the whole INTEGER dot — are IDENTICAL to the CPU reference. Only +// the per-super-block float scale sum is reassociated (warp reduction vs the +// CPU's sequential add), so the gate is the CUDA sibling's gate: INTEGER core +// bit-exact, final scale within the NMSE band test_ops_quant_dot uses. +// +// * Scratch uses hipMallocAsync (present in ROCm 7.14) with the same retire +// -never-free discipline as the CUDA side (graph_safe_scratch.h). +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "vt/cpu/cpu_quant_blocks.h" // vt::cpu::Block* struct mirror + // (single source; plain C++) +#include "vt/cuda/cuda_quant_iq_tables.cuh" // d_iq2xxs_grid / d_iq3xxs_grid / + // d_iq2s_grid / d_iq1s_grid / + // d_iq1xxxs_grid / d_ksigns / + // d_kmask (single source; pure + // __device__ syntax, no CUDA) +#include "vt/cuda/graph_safe_scratch.h" // RetireGraphScratch (portable) +#include "vt/ops.h" +#include "vt/quant.h" + +namespace vt::rocm { +// Forward declaration: the K-quant (Q4_K/Q5_K/Q6_K) path lives in +// rocm_grouped_gemm.hip (external linkage) with the MMVQ arm and +// norm-quant bridge. This file delegates to it for those types. +void MatmulBTQuantKernelRocmGdn(Queue& q, Tensor& out, + const Tensor& a, const Tensor& b); +void MatmulBTQuantGroupedKernelRocmGdn( + Queue& q, Tensor& out, const Tensor& act, + const Tensor& weight, const Tensor& expert_ids); + +namespace { + +// dtype.cpp F32ToF16 — round-to-nearest-even, subnormals, inf/nan. Used only +// for the Q8_0 activation scale (the CPU Q8_0 vec_dot's f16 round-trip). +__device__ inline uint16_t DF32ToF16(float f) { + uint32_t u = __float_as_uint(f); + uint16_t sign = static_cast((u >> 16) & 0x8000); + int32_t exp = static_cast((u >> 23) & 0xFF) - 127 + 15; + uint32_t mant = u & 0x7FFFFF; + if (((u >> 23) & 0xFF) == 0xFF) + return static_cast(sign | 0x7C00 | (mant ? 0x200 | (mant >> 13) : 0)); + if (exp >= 0x1F) return static_cast(sign | 0x7C00); + if (exp <= 0) { + if (exp < -10) return sign; + mant |= 0x800000; + uint32_t shift = static_cast(14 - exp); + uint32_t half = mant >> shift; + uint32_t rem = mant & ((1u << shift) - 1); + uint32_t mid = 1u << (shift - 1); + if (rem > mid || (rem == mid && (half & 1))) ++half; + return static_cast(sign | half); + } + uint32_t half = static_cast(exp << 10) | (mant >> 13); + uint32_t rem = mant & 0x1FFF; + if (rem > 0x1000 || (rem == 0x1000 && (half & 1))) ++half; + return static_cast(sign | half); +} + +// Load one activation element (dtype-decoded, exactly like cpu LoadActF32). + +using vt::cpu::BlockIQ1_S; +using vt::cpu::BlockIQ1_XXXS; +using vt::cpu::BlockIQ2_S; +using vt::cpu::BlockIQ2_XXS; +using vt::cpu::BlockIQ3_XXS; +using vt::cpu::BlockQ2_K; +using vt::cpu::BlockQ3_K; +using vt::cpu::BlockQ4_K; +using vt::cpu::BlockQ5_K; +using vt::cpu::BlockQ6_K; +using vt::cpu::BlockQ8_K; +using vt::cpu::BlockQ8_0; +using vt::cpu::kQK_K; +using vt::cpu::kQK8_0; + +void CheckHip(hipError_t err, const char* what) { + if (err != hipSuccess) { + throw std::runtime_error(std::string("vt rocm: matmul_bt_quant: ") + what + + ": " + hipGetErrorString(err)); + } +} + +// --- device numeric helpers — bit-exact ports of src/vt/dtype.cpp ------------- +__device__ inline float DF16ToF32(uint16_t h) { + uint32_t sign = static_cast(h & 0x8000) << 16; + uint32_t exp = (h >> 10) & 0x1F; + uint32_t mant = h & 0x3FF; + if (exp == 0x1F) return __int_as_float(sign | 0x7F800000 | (mant << 13)); + if (exp == 0) { + if (mant == 0) return __int_as_float(sign); + int shift = 0; + while ((mant & 0x400) == 0) { + mant <<= 1; + ++shift; + } + mant &= 0x3FF; + return __int_as_float(sign | ((113 - shift) << 23) | (mant << 13)); + } + return __int_as_float(sign | ((exp + 112) << 23) | (mant << 13)); +} + +__device__ inline float DBF16ToF32(uint16_t b) { + return __int_as_float(static_cast(b) << 16); +} + +__device__ inline uint16_t DF32ToBF16(float f) { + uint32_t u = __float_as_int(f); + if ((u & 0x7F800000) == 0x7F800000 && (u & 0x7FFFFF)) { + return static_cast((u >> 16) | 0x0040); + } + uint32_t rounding = 0x7FFF + ((u >> 16) & 1); + return static_cast((u + rounding) >> 16); +} + +// cpu_quant_act.cpp NearestInt (ggml-quants.c:563) — magic-constant round-to-even. +__device__ inline int DNearestInt(float fval) { + float val = fval + 12582912.0f; + int i = __float_as_int(val); + return (i & 0x007fffff) - 0x00400000; +} + +enum class ActDT : int { kF32 = 0, kF16 = 1, kBF16 = 2 }; + +__device__ inline float DLoadAct(const void* base, ActDT dt, int64_t idx) { + switch (dt) { + case ActDT::kF32: return static_cast(base)[idx]; + case ActDT::kF16: return DF16ToF32(static_cast(base)[idx]); + default: return DBF16ToF32(static_cast(base)[idx]); + } +} + +// --------------------------------------------------------------------------- +// GPU activation quantizer — ds4-parity grid: ONE BLOCK per (super-block, row), +// 256 threads (one thread per element). Bit-exact port of QuantizeRowQ8_K with +// the argmax tie broken by LOWEST original index (== the sequential first- +// occurrence scan). Byte-identical to both CUDA quantizers by construction. +// --------------------------------------------------------------------------- +__global__ void QuantizeQ8KKernel(BlockQ8_K* __restrict__ scratch, + const void* __restrict__ a, ActDT adt, + int64_t a_rs, int64_t m, int64_t nsb) { + const int64_t b = static_cast(blockIdx.x); // super-block within row + const int64_t i = static_cast(blockIdx.y); // activation row + if (b >= nsb || i >= m) return; + const int tid = static_cast(threadIdx.x); + const int64_t elem0 = i * a_rs + b * kQK_K; + const float v = DLoadAct(a, adt, elem0 + tid); + + __shared__ float sabs[kQK_K]; + __shared__ float sval[kQK_K]; + __shared__ int sidx[kQK_K]; + sabs[tid] = fabsf(v); + sval[tid] = v; + sidx[tid] = tid; + __syncthreads(); +#pragma unroll + for (int stride = kQK_K >> 1; stride > 0; stride >>= 1) { + if (tid < stride) { + const float oa = sabs[tid + stride]; + if (oa > sabs[tid] || (oa == sabs[tid] && sidx[tid + stride] < sidx[tid])) { + sabs[tid] = oa; + sval[tid] = sval[tid + stride]; + sidx[tid] = sidx[tid + stride]; + } + } + __syncthreads(); + } + const float mx = sval[0]; + const float amax = sabs[0]; + + BlockQ8_K& y = scratch[i * nsb + b]; + if (amax == 0.0f) { + if (tid == 0) y.d = 0.0f; + y.qs[tid] = 0; + if (tid < kQK_K / 16) y.bsums[tid] = 0; + return; + } + const float iscale = -127.0f / mx; + const int qv = DNearestInt(iscale * v); + y.qs[tid] = static_cast(qv < 127 ? qv : 127); + __syncthreads(); + if (tid < kQK_K / 16) { + int sum = 0; + for (int ii = 0; ii < 16; ++ii) sum += y.qs[tid * 16 + ii]; + y.bsums[tid] = static_cast(sum); + } + if (tid == 0) y.d = 1.0f / iscale; +} + +// --------------------------------------------------------------------------- +// Per-super-block integer dots. PORTABLE SCALAR forms of the CPU reference +// bodies (cpu_quant_dot.cpp), NOT the CUDA __dp4a forms: gfx1100 has no +// signed byte dot (see the HIP DELTAS header note). Exactness is preserved by +// keeping each dot's accumulation ORDER identical to the CPU body it mirrors: +// Q2_K sums sub-blocks in k/j/l order; Q3_K/Q6_K use the 8-wide aux32 split; +// Q4_K/Q5_K accumulate per-16 bsums then per-32 scale groups in order. +// --------------------------------------------------------------------------- +__device__ inline float DotQ2K(const BlockQ2_K* xb, const BlockQ8_K* yb) { + const uint8_t* q2 = xb->qs; + const int8_t* q8 = yb->qs; + const uint8_t* sc = xb->scales; + int summs = 0; + for (int j = 0; j < 16; ++j) summs += yb->bsums[j] * (sc[j] >> 4); + const float dall = yb->d * DF16ToF32(xb->d); + const float dmin = yb->d * DF16ToF32(xb->dmin); + int isum = 0; + int is = 0; + for (int k = 0; k < kQK_K / 128; ++k) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + int d = sc[is++] & 0xF; + int isuml = 0; + for (int l = 0; l < 16; ++l) isuml += q8[l] * ((q2[l] >> shift) & 3); + isum += d * isuml; + d = sc[is++] & 0xF; + isuml = 0; + for (int l = 16; l < 32; ++l) isuml += q8[l] * ((q2[l] >> shift) & 3); + isum += d * isuml; + shift += 2; + q8 += 32; + } + q2 += 32; + } + return dall * isum - dmin * summs; +} + +__device__ inline float DotQ3K(const BlockQ3_K* xb, const BlockQ8_K* yb) { + const uint32_t kmask1 = 0x03030303; + const uint32_t kmask2 = 0x0f0f0f0f; + const uint8_t* hm = xb->hmask; + const int8_t* q8 = yb->qs; + int8_t aux8[kQK_K]; + int8_t* a = aux8; + const uint8_t* q3 = xb->qs; + uint8_t m = 1; + for (int jj = 0; jj < kQK_K; jj += 128) { + for (int l = 0; l < 32; ++l) a[l] = q3[l] & 3; + for (int l = 0; l < 32; ++l) a[l] = static_cast(a[l] - ((hm[l] & m) ? 0 : 4)); + a += 32; m = static_cast(m << 1); + for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 2) & 3; + for (int l = 0; l < 32; ++l) a[l] = static_cast(a[l] - ((hm[l] & m) ? 0 : 4)); + a += 32; m = static_cast(m << 1); + for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 4) & 3; + for (int l = 0; l < 32; ++l) a[l] = static_cast(a[l] - ((hm[l] & m) ? 0 : 4)); + a += 32; m = static_cast(m << 1); + for (int l = 0; l < 32; ++l) a[l] = (q3[l] >> 6) & 3; + for (int l = 0; l < 32; ++l) a[l] = static_cast(a[l] - ((hm[l] & m) ? 0 : 4)); + a += 32; m = static_cast(m << 1); + q3 += 32; + } + uint32_t auxs[4]; + memcpy(auxs, xb->scales, 12); + const int8_t* scales = reinterpret_cast(auxs); + uint32_t tmp = auxs[2]; + auxs[2] = ((auxs[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + auxs[3] = ((auxs[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + auxs[0] = (auxs[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + auxs[1] = (auxs[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + a = aux8; + const int8_t* q8p = q8; + int32_t aux32[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int j = 0; j < kQK_K / 16; ++j) { + for (int l = 0; l < 8; ++l) aux32[l] += (scales[j] - 32) * (q8p[l] * a[l]); + q8p += 8; a += 8; + for (int l = 0; l < 8; ++l) aux32[l] += (scales[j] - 32) * (q8p[l] * a[l]); + q8p += 8; a += 8; + } + const float d = DF16ToF32(xb->d) * yb->d; + int isum = 0; + for (int l = 0; l < 8; ++l) isum += aux32[l]; + return d * isum; +} + +__device__ inline float DotQ4K(const BlockQ4_K* xb, const BlockQ8_K* yb) { + const uint32_t kmask1 = 0x3f3f3f3f; + const uint32_t kmask2 = 0x0f0f0f0f; + const uint32_t kmask3 = 0x03030303; + const uint8_t* q4 = xb->qs; + const int8_t* q8 = yb->qs; + uint32_t utmp[4]; + memcpy(utmp, xb->scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + const uint8_t* scales = reinterpret_cast(&utmp[0]); + const uint8_t* mins = reinterpret_cast(&utmp[2]); + int sumi = 0; + for (int j = 0; j < kQK_K / 16; ++j) sumi += yb->bsums[j] * mins[j / 2]; + // Portable nibble walk in the CPU body's order: 64-element groups decode + // low nibble then high nibble into aux8, then the per-32 scale groups. + int8_t aux8[kQK_K]; + int8_t* a = aux8; + for (int j = 0; j < kQK_K / 64; ++j) { + for (int l = 0; l < 32; ++l) a[l] = static_cast(q4[l] & 0xF); + a += 32; + for (int l = 0; l < 32; ++l) a[l] = static_cast(q4[l] >> 4); + a += 32; + q4 += 32; + } + a = aux8; + int32_t aux32[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int j = 0; j < kQK_K / 32; ++j) { + const int32_t scale = scales[j]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8[l] * a[l]); + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8[l] * a[l]); + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8[l] * a[l]); + q8 += 8; a += 8; + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8[l] * a[l]); + q8 += 8; a += 8; + } + const float d = DF16ToF32(xb->d) * yb->d; + const float dmin = DF16ToF32(xb->dmin) * yb->d; + int isum = 0; + for (int l = 0; l < 8; ++l) isum += aux32[l]; + return d * isum - dmin * sumi; +} + +__device__ inline float DotQ5K(const BlockQ5_K* xb, const BlockQ8_K* yb) { + const uint32_t kmask1 = 0x3f3f3f3f; + const uint32_t kmask2 = 0x0f0f0f0f; + const uint32_t kmask3 = 0x03030303; + const uint8_t* q4 = xb->qs; + const uint8_t* hm = xb->qh; + const int8_t* q8 = yb->qs; + uint32_t utmp[4]; + memcpy(utmp, xb->scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + const uint8_t* scales = reinterpret_cast(&utmp[0]); + const uint8_t* mins = reinterpret_cast(&utmp[2]); + int sumi = 0; + for (int j = 0; j < kQK_K / 16; ++j) sumi += yb->bsums[j] * mins[j / 2]; + int8_t aux8[kQK_K]; + int8_t* a = aux8; + uint8_t m = 1; + for (int j = 0; j < kQK_K / 64; ++j) { + for (int l = 0; l < 32; ++l) a[l] = static_cast(q4[l] & 0xF); + for (int l = 0; l < 32; ++l) + a[l] = static_cast(a[l] + ((hm[l] & m) ? 16 : 0)); + a += 32; + m = static_cast(m << 1); + for (int l = 0; l < 32; ++l) a[l] = static_cast(q4[l] >> 4); + for (int l = 0; l < 32; ++l) + a[l] = static_cast(a[l] + ((hm[l] & m) ? 16 : 0)); + a += 32; + m = static_cast(m << 1); + q4 += 32; + } + a = aux8; + int32_t aux32[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int j = 0; j < kQK_K / 32; ++j) { + const int32_t scale = scales[j]; + for (int r = 0; r < 4; ++r) { + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8[l] * a[l]); + q8 += 8; a += 8; + } + } + const float d = DF16ToF32(xb->d) * yb->d; + const float dmin = DF16ToF32(xb->dmin) * yb->d; + int isum = 0; + for (int l = 0; l < 8; ++l) isum += aux32[l]; + return d * isum - dmin * sumi; +} + +__device__ inline float DotQ6K(const BlockQ6_K* xb, const BlockQ8_K* yb) { + const uint8_t* q4 = xb->ql; + const uint8_t* qh = xb->qh; + const int8_t* q8 = yb->qs; + int8_t aux8[kQK_K]; + int8_t* a = aux8; + for (int j = 0; j < kQK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + a[l + 0] = static_cast( + static_cast((q4[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32); + a[l + 32] = static_cast( + static_cast((q4[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32); + a[l + 64] = static_cast( + static_cast((q4[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32); + a[l + 96] = static_cast( + static_cast((q4[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32); + } + a += 128; q4 += 64; qh += 32; + } + a = aux8; + const int8_t* q8p = q8; + int32_t aux32[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int j = 0; j < kQK_K / 16; ++j) { + const int scale = xb->scales[j]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8p[l] * a[l]); + q8p += 8; a += 8; + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8p[l] * a[l]); + q8p += 8; a += 8; + } + const float d = DF16ToF32(xb->d) * yb->d; + int isum = 0; + for (int l = 0; l < 8; ++l) isum += aux32[l]; + return d * isum; +} + +__device__ inline float DotIQ2XXS(const BlockIQ2_XXS* xb, const BlockQ8_K* yb) { + const float d = DF16ToF32(xb->d) * yb->d; + const uint16_t* qs = xb->qs; + const int8_t* q8 = yb->qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < kQK_K / 32; ++ib32) { + uint32_t aux32[2]; + memcpy(aux32, qs + 4 * ib32, 2 * sizeof(uint32_t)); + const uint32_t ls = 2 * (aux32[1] >> 28) + 1; + int32_t sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint8_t* grid = reinterpret_cast( + &vt::cuda::d_iq2xxs_grid[(aux32[0] >> (8 * l)) & 0xff]); + const uint8_t signs = + vt::cuda::d_ksigns_iq2xs[(aux32[1] >> (7 * l)) & 127]; + for (int j = 0; j < 8; ++j) + sumi += grid[j] * q8[j] * ((signs & vt::cuda::d_kmask_iq2xs[j]) ? -1 : 1); + q8 += 8; + } + bsum += sumi * static_cast(ls); + } + return d * bsum; // final *0.125 applied after the warp reduction +} + +__device__ inline float DotIQ3XXS(const BlockIQ3_XXS* xb, const BlockQ8_K* yb) { + const float d = DF16ToF32(xb->d) * yb->d; + const uint8_t* q3 = xb->qs; + const uint8_t* gas = xb->qs + kQK_K / 4; + const int8_t* q8 = yb->qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < kQK_K / 32; ++ib32) { + uint32_t a32; + memcpy(&a32, gas, sizeof(uint32_t)); + gas += sizeof(uint32_t); + const uint32_t ls = 2 * (a32 >> 28) + 1; + int32_t sumi = 0; + for (int l = 0; l < 4; ++l) { + const uint32_t g1 = vt::cuda::d_iq3xxs_grid[q3[2 * l + 0]]; + const uint32_t g2 = vt::cuda::d_iq3xxs_grid[q3[2 * l + 1]]; + const uint8_t signs = vt::cuda::d_ksigns_iq2xs[(a32 >> (7 * l)) & 127]; + for (int j = 0; j < 4; ++j) { + const int b1 = static_cast((g1 >> (8 * j)) & 0xff); + const int b2 = static_cast((g2 >> (8 * j)) & 0xff); + sumi += b1 * q8[j + 0] * ((signs & vt::cuda::d_kmask_iq2xs[j + 0]) ? -1 : 1); + sumi += b2 * q8[j + 4] * ((signs & vt::cuda::d_kmask_iq2xs[j + 4]) ? -1 : 1); + } + q8 += 8; + } + q3 += 8; + bsum += sumi * static_cast(ls); + } + return d * bsum; // final *0.25 applied after the warp reduction +} + +__device__ inline float DotIQ2S(const BlockIQ2_S* xb, const BlockQ8_K* yb) { + const float d = DF16ToF32(xb->d) * yb->d; + const int8_t* q8 = yb->qs; + const uint8_t* qs = xb->qs; + const uint8_t* qh = xb->qh; + const uint8_t* signs = qs + kQK_K / 8; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < kQK_K / 32; ++ib32) { + const int ls1 = 1 + 2 * (xb->scales[ib32] & 0xf); + const int ls2 = 1 + 2 * (xb->scales[ib32] >> 4); + int sumi1 = 0; + int sumi2 = 0; + for (int l = 0; l < 2; ++l) { + const uint8_t* grid = reinterpret_cast( + &vt::cuda::d_iq2s_grid[qs[l] | ((qh[ib32] << (8 - 2 * l)) & 0x300)]); + for (int j = 0; j < 8; ++j) + sumi1 += q8[j] * grid[j] * ((signs[l] & vt::cuda::d_kmask_iq2xs[j]) ? -1 : 1); + q8 += 8; + } + for (int l = 2; l < 4; ++l) { + const uint8_t* grid = reinterpret_cast( + &vt::cuda::d_iq2s_grid[qs[l] | ((qh[ib32] << (8 - 2 * l)) & 0x300)]); + for (int j = 0; j < 8; ++j) + sumi2 += q8[j] * grid[j] * ((signs[l] & vt::cuda::d_kmask_iq2xs[j]) ? -1 : 1); + q8 += 8; + } + bsum += ls1 * sumi1 + ls2 * sumi2; + qs += 4; + signs += 4; + } + return d * bsum; // final *0.125 applied after the warp reduction +} + +__device__ inline float DotIQ1S(const BlockIQ1_S* xb, const BlockQ8_K* yb) { + const int8_t* q8 = yb->qs; + const uint8_t* qs = xb->qs; + const uint16_t* qh = xb->qh; + int32_t sumi = 0; + int32_t sumi1 = 0; + for (int ib = 0; ib < kQK_K / 32; ++ib) { + const int ls = 2 * ((qh[ib] >> 12) & 7) + 1; + const int delta = (qh[ib] & 0x8000) ? -1 : 1; + int lsum = 0; + for (int l = 0; l < 4; ++l) { + const int8_t* grid = reinterpret_cast( + &vt::cuda::d_iq1s_grid[qs[l] | (((qh[ib] >> (3 * l)) & 7) << 8)]); + for (int j = 0; j < 8; ++j) lsum += q8[j] * grid[j]; + q8 += 8; + } + sumi += ls * lsum; + sumi1 += ls * delta * (yb->bsums[2 * ib + 0] + yb->bsums[2 * ib + 1]); + qs += 4; + } + return DF16ToF32(xb->d) * yb->d * + (static_cast(sumi) + 0.125f * static_cast(sumi1)); +} + +__device__ inline float DotIQ1XXXS(const BlockIQ1_XXXS* xb, const BlockQ8_K* yb) { + const int8_t* q8 = yb->qs; + const uint8_t* qs = xb->qs; + const uint8_t* sc = xb->sc; + int32_t sumi = 0; + int32_t sumi1 = 0; + for (int ib = 0; ib < kQK_K / 32; ++ib) { + const int nib = (sc[ib / 2] >> (4 * (ib & 1))) & 0xf; + const int ls = 2 * (nib & 7) + 1; + const int delta = (nib & 8) ? -1 : 1; + int lsum = 0; + for (int l = 0; l < 4; ++l) { + const int8_t* grid = + reinterpret_cast(&vt::cuda::d_iq1xxxs_grid[qs[l]]); + for (int j = 0; j < 8; ++j) lsum += q8[j] * grid[j]; + q8 += 8; + } + sumi += ls * lsum; + sumi1 += ls * delta * (yb->bsums[2 * ib + 0] + yb->bsums[2 * ib + 1]); + qs += 4; + } + return DF16ToF32(xb->d) * yb->d * + (static_cast(sumi) + 0.125f * static_cast(sumi1)); +} + +// --------------------------------------------------------------------------- +// WType tags + DotSuperblock dispatch — mirrors the CUDA file's table. +// --------------------------------------------------------------------------- +enum class WType : int { + kIQ2_XXS = 0, + kIQ3_XXS = 1, + kQ2_K = 2, + kQ3_K = 3, + kQ4_K = 4, + kQ5_K = 5, + kQ6_K = 6, + kIQ2_S = 7, + kIQ1_S = 8, + kIQ1_XXXS = 9, +}; + +template +__device__ inline float DotSuperblock(const void* w_sb, const BlockQ8_K* a_sb); + +template <> +__device__ inline float DotSuperblock(const void* w, const BlockQ8_K* a) { + return DotIQ2XXS(static_cast(w), a); +} +template <> +__device__ inline float DotSuperblock(const void* w, const BlockQ8_K* a) { + return DotIQ3XXS(static_cast(w), a); +} +template <> +__device__ inline float DotSuperblock(const void* w, const BlockQ8_K* a) { + return DotQ2K(static_cast(w), a); +} +template <> +__device__ inline float DotSuperblock(const void* w, const BlockQ8_K* a) { + return DotQ3K(static_cast(w), a); +} +template <> +__device__ inline float DotSuperblock(const void* w, const BlockQ8_K* a) { + return DotQ4K(static_cast(w), a); +} +template <> +__device__ inline float DotSuperblock(const void* w, const BlockQ8_K* a) { + return DotQ5K(static_cast(w), a); +} +template <> +__device__ inline float DotSuperblock(const void* w, const BlockQ8_K* a) { + return DotQ6K(static_cast(w), a); +} +template <> +__device__ inline float DotSuperblock(const void* w, const BlockQ8_K* a) { + return DotIQ2S(static_cast(w), a); +} +template <> +__device__ inline float DotSuperblock(const void* w, const BlockQ8_K* a) { + return DotIQ1S(static_cast(w), a); +} +template <> +__device__ inline float DotSuperblock(const void* w, const BlockQ8_K* a) { + return DotIQ1XXXS(static_cast(w), a); +} + +template +__device__ constexpr float FinalFactor() { + return (W == WType::kIQ2_XXS || W == WType::kIQ2_S) + ? 0.125f + : (W == WType::kIQ3_XXS ? 0.25f : 1.0f); +} + +// --------------------------------------------------------------------------- +// The MMVQ-style GEMM: one WARP per output element (i,j). Lanes split the K +// super-blocks; the warp reduction sums the partials. HIP delta: the shuffle +// mask is 64-bit on this target. Determinism note unchanged from CUDA: the +// integer core is exact; only the scale sum reassociates (within NMSE). +// --------------------------------------------------------------------------- +template +__global__ void QuantDotGemmKernel(OutT* __restrict__ out, + const uint8_t* __restrict__ weight, + const BlockQ8_K* __restrict__ act, int64_t m, + int64_t n, int64_t nsb, size_t w_row_bytes, + size_t w_block_bytes) { + const int64_t warp = static_cast(blockIdx.x) * (blockDim.x >> 5) + + (threadIdx.x >> 5); + if (warp >= m * n) return; + const int64_t i = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x & 31; + + const uint8_t* w_row = weight + static_cast(j) * w_row_bytes; + const BlockQ8_K* a_row = act + i * nsb; + + float partial = 0.0f; + for (int64_t sb = lane; sb < nsb; sb += 32) { + const void* w_sb = w_row + static_cast(sb) * w_block_bytes; + partial += DotSuperblock(w_sb, a_row + sb); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + partial += __shfl_down_sync(0xffffffffffffffffull, partial, off); + + if (lane == 0) { + const float v = FinalFactor() * partial; + if constexpr (sizeof(OutT) == 4) { + out[i * n + j] = v; + } else { + out[i * n + j] = DF32ToBF16(v); + } + } +} + +// GROUPED variant: warp per (p, n); weight row selected by expert_ids[p]. +template +__global__ void QuantDotGemmGroupedKernel(OutT* __restrict__ out, + const uint8_t* __restrict__ weight, + const BlockQ8_K* __restrict__ act, + const int32_t* __restrict__ expert_ids, + int64_t P, int64_t n, int64_t nsb, + size_t w_row_bytes, + size_t w_block_bytes, bool bcast) { + const int64_t warp = static_cast(blockIdx.x) * (blockDim.x >> 5) + + (threadIdx.x >> 5); + if (warp >= P * n) return; + const int64_t p = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x & 31; + + const int64_t e = expert_ids[p]; + const uint8_t* w_row = weight + static_cast(e * n + j) * w_row_bytes; + // Broadcast activation: the routed gate/up share ONE quantized hidden. + const BlockQ8_K* a_row = act + (bcast ? 0 : p) * nsb; + + float partial = 0.0f; + for (int64_t sb = lane; sb < nsb; sb += 32) { + const void* w_sb = w_row + static_cast(sb) * w_block_bytes; + partial += DotSuperblock(w_sb, a_row + sb); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + partial += __shfl_down_sync(0xffffffffffffffffull, partial, off); + + if (lane == 0) { + const float v = FinalFactor() * partial; + if constexpr (sizeof(OutT) == 4) { + out[p * n + j] = v; + } else { + out[p * n + j] = DF32ToBF16(v); + } + } +} + +// --------------------------------------------------------------------------- +// Host wiring: scratch, launches, providers, registrar. +// --------------------------------------------------------------------------- +struct StreamScratch { + void* buf = nullptr; + size_t bytes = 0; +}; + +std::mutex& ScratchMutex() { + static std::mutex mu; + return mu; +} + +StreamScratch& ScratchFor(hipStream_t s) { + static std::unordered_map map; + return map[s]; +} + +void CheckHipLaunch(const char* what) { CheckHip(hipGetLastError(), what); } + +void* EnsureScratch(size_t need, hipStream_t s) { + std::lock_guard lock(ScratchMutex()); + StreamScratch& sc = ScratchFor(s); + if (need > sc.bytes) { + // Retire (never free): a captured hipGraph may have baked this pointer. + vt::cuda::RetireGraphScratch(sc.buf); + CheckHip(hipMallocAsync(&sc.buf, need, s), "hipMallocAsync q8_K act scratch"); + sc.bytes = need; + } + return sc.buf; +} + +inline ActDT ActDtOf(DType dt) { + return dt == DType::kF32 ? ActDT::kF32 : dt == DType::kF16 ? ActDT::kF16 : ActDT::kBF16; +} + +void LaunchQuantizeQ8K(BlockQ8_K* qact, const void* data, ActDT adt, int64_t a_rs, + int64_t rows, int64_t nsb, hipStream_t s) { + dim3 qgrid(static_cast(nsb), static_cast(rows), 1); + QuantizeQ8KKernel<<>>(qact, data, adt, a_rs, rows, nsb); + CheckHipLaunch("quantize_q8_K launch"); +} + +bool IsRocmKeepQuantSupported(DType dt, WType* out) { + switch (dt) { + case DType::kIQ2_XXS: *out = WType::kIQ2_XXS; return true; + case DType::kIQ3_XXS: *out = WType::kIQ3_XXS; return true; + case DType::kQ2_K: *out = WType::kQ2_K; return true; + case DType::kQ3_K: *out = WType::kQ3_K; return true; + case DType::kQ4_K: *out = WType::kQ4_K; return true; + case DType::kQ5_K: *out = WType::kQ5_K; return true; + case DType::kQ6_K: *out = WType::kQ6_K; return true; + case DType::kIQ2_S: *out = WType::kIQ2_S; return true; + case DType::kIQ1_S: *out = WType::kIQ1_S; return true; + case DType::kIQ1_XXXS: *out = WType::kIQ1_XXXS; return true; + // Q4_0 / Q8_0 / MXFP4 dot a Q8_0 activation and have no native arm here. + default: return false; + } +} + +template +void LaunchGemm(Tensor& out, const uint8_t* weight, const BlockQ8_K* act, + int64_t m, int64_t n, int64_t nsb, size_t w_row_bytes, + size_t w_block_bytes, hipStream_t s) { + // Wave32 geometry: one warp per output; 8 warps (256 threads) per block. + constexpr int kWarpsPerBlock = 8; + dim3 block(32 * kWarpsPerBlock, 1, 1); + const int64_t warps = m * n; + const unsigned grid = + static_cast((warps + kWarpsPerBlock - 1) / kWarpsPerBlock); + if (out.dtype == DType::kF32) { + QuantDotGemmKernel<<>>( + static_cast(out.data), weight, act, m, n, nsb, w_row_bytes, + w_block_bytes); + } else { + QuantDotGemmKernel<<>>( + static_cast(out.data), weight, act, m, n, nsb, w_row_bytes, + w_block_bytes); + } + CheckHipLaunch("matmul_bt_quant launch"); +} + +template +void LaunchGroupedGemm(Tensor& out, const uint8_t* weight, const BlockQ8_K* act, + const int32_t* expert_ids, int64_t P, int64_t n, + int64_t nsb, size_t w_row_bytes, size_t w_block_bytes, + bool bcast, hipStream_t s) { + constexpr int kWarpsPerBlock = 8; + dim3 block(32 * kWarpsPerBlock, 1, 1); + const int64_t warps = P * n; + const unsigned grid = + static_cast((warps + kWarpsPerBlock - 1) / kWarpsPerBlock); + if (out.dtype == DType::kF32) { + QuantDotGemmGroupedKernel<<>>( + static_cast(out.data), weight, act, expert_ids, P, n, nsb, + w_row_bytes, w_block_bytes, bcast); + } else { + QuantDotGemmGroupedKernel<<>>( + static_cast(out.data), weight, act, expert_ids, P, n, nsb, + w_row_bytes, w_block_bytes, bcast); + } + CheckHipLaunch("matmul_bt_quant_grouped launch"); +} + + + +// Q8_0 (legacy 32-block, Q8_0-activation) arm. Self-contained: quantize the +// activation to Q8_0 on the device, then the Q8_0xQ8_0 integer dot. The dot is +// the PORTABLE SCALAR form of cpu_quant_dot.cpp VecDotQ8_0Q8_0 (quants.c:400): +// sumi0 over the first 16 elements, sumi1 over the second, ONE int32 add, then +// the f16-scale fold — accumulation order preserved for exactness. The CUDA +// file's __dp4a form is byte-identical to this on a signed-dot target; gfx1100 +// has no signed byte dot, so W1 ships the scalar body. +// --------------------------------------------------------------------------- +__global__ void QuantizeQ8_0Kernel(BlockQ8_0* __restrict__ scratch, + const void* __restrict__ a, ActDT adt, + int64_t a_rs, int64_t m, int64_t nb) { + const int64_t t = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (t >= m * nb) return; + const int64_t i = t / nb; + const int64_t b = t % nb; + const int64_t elem0 = i * a_rs + b * kQK8_0; + float amax = 0.0f; + for (int j = 0; j < kQK8_0; ++j) { + const float av = fabsf(DLoadAct(a, adt, elem0 + j)); + amax = amax > av ? amax : av; + } + BlockQ8_0& y = scratch[t]; + const float d = amax / 127.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + y.d = DF32ToF16(d); + for (int j = 0; j < kQK8_0; ++j) { + y.qs[j] = static_cast(roundf(DLoadAct(a, adt, elem0 + j) * id)); + } +} + +template +__global__ void QuantDotGemmQ8_0Kernel(OutT* __restrict__ out, + const uint8_t* __restrict__ weight, + const BlockQ8_0* __restrict__ act, + int64_t m, int64_t n, int64_t nb, + size_t w_row_bytes) { + const int64_t warp = static_cast(blockIdx.x) * (blockDim.x >> 5) + + (threadIdx.x >> 5); + if (warp >= m * n) return; + const int64_t i = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x & 31; + const uint8_t* w_row = weight + static_cast(j) * w_row_bytes; + const BlockQ8_0* a_row = act + i * nb; + float partial = 0.0f; + for (int64_t b = lane; b < nb; b += 32) { + const BlockQ8_0* wb = + reinterpret_cast(w_row + static_cast(b) * + sizeof(BlockQ8_0)); + const BlockQ8_0* ab = a_row + b; + // CPU order: first half then second half, one add, then scale. + int sumi0 = 0; + for (int j2 = 0; j2 < kQK8_0 / 2; ++j2) + sumi0 += ab->qs[j2] * wb->qs[j2]; + int sumi1 = 0; + for (int j2 = kQK8_0 / 2; j2 < kQK8_0; ++j2) + sumi1 += ab->qs[j2] * wb->qs[j2]; + partial += (sumi0 + sumi1) * (DF16ToF32(wb->d) * DF16ToF32(ab->d)); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + partial += __shfl_down_sync(0xffffffffffffffffull, partial, off); + if (lane == 0) { + if constexpr (sizeof(OutT) == 4) + out[i * n + j] = partial; + else + out[i * n + j] = DF32ToBF16(partial); + } +} + +void MatmulQ8_0Rocm(Tensor& out, const Tensor& a, const Tensor& b, + hipStream_t s) { + const int64_t m = a.shape[0], k = a.shape[1], n = b.shape[0]; + if (m == 0 || n == 0) return; + if (k % kQK8_0 != 0) + throw std::runtime_error( + "vt rocm: matmul_bt_quant Q8_0: K must be a multiple of 32"); + const int64_t nb = k / kQK8_0; + const size_t w_row_bytes = static_cast(nb) * sizeof(BlockQ8_0); + const size_t act_bytes = + static_cast(m) * static_cast(nb) * sizeof(BlockQ8_0); + BlockQ8_0* act = static_cast(EnsureScratch(act_bytes, s)); + { + constexpr int kQBlock = 128; + const int64_t grid = (m * nb + kQBlock - 1) / kQBlock; + QuantizeQ8_0Kernel<<(grid), kQBlock, 0, s>>>( + act, a.data, ActDtOf(a.dtype), a.stride[0], m, nb); + CheckHipLaunch("quantize_q8_0 launch"); + } + constexpr int kWarpsPerBlock = 8; + dim3 block(32 * kWarpsPerBlock, 1, 1); + const unsigned grid = + static_cast((m * n + kWarpsPerBlock - 1) / kWarpsPerBlock); + if (out.dtype == DType::kF32) + QuantDotGemmQ8_0Kernel<<>>( + static_cast(out.data), static_cast(b.data), + act, m, n, nb, w_row_bytes); + else + QuantDotGemmQ8_0Kernel<<>>( + static_cast(out.data), static_cast(b.data), + act, m, n, nb, w_row_bytes); + CheckHipLaunch("matmul_bt_quant Q8_0 launch"); +} +} // namespace + +// The kROCM provider for OpId::kMatmulBTQuant. Validation already done by +// vt::MatmulBTQuant (ops.cpp). Contract: b is [N,K] block-quant, a [M,K] +// f32/bf16 row-packed, out [M,N]. On this DISCRETE backend an unsupported +// dtype cannot fall back to the CPU kernel (it would follow device pointers), +// so it throws naming the dtype — VT_GGUF_KEEP_QUANT=0 restores the load-time +// bf16 expansion for such files. +// +// External linkage on purpose: rocm_ops.hip declares and registers these two +// entry points for kROCM, and rocm_moe_gate_up_swiglu.hip calls the grouped +// one. Internal-linkage definitions here left those references unresolved at +// link. The helpers above stay internal; only the two entry points are +// exported, and registration stays with rocm_ops.hip's single registrar. +void MatmulBTQuantKernelRocm(Queue& q, Tensor& out, const Tensor& a, + const Tensor& b) { + hipStream_t s = static_cast(q.handle); + const int64_t m = a.shape[0]; + const int64_t k = a.shape[1]; + const int64_t n = b.shape[0]; + if (b.dtype == DType::kQ8_0) { + MatmulQ8_0Rocm(out, a, b, s); + return; + } + + // Delegate K-quant (Q4_K/Q5_K/Q6_K) to rocm_grouped_gemm.hip which + // has the MMVQ arm (VT_GEMV_MMVQ) and norm-quant bridge (lever-C). + if (b.dtype == DType::kQ4_K || b.dtype == DType::kQ5_K || + b.dtype == DType::kQ6_K) { + MatmulBTQuantKernelRocmGdn(q, out, a, b); + return; + } + + WType w{}; + if (!IsRocmKeepQuantSupported(b.dtype, &w)) { + throw std::runtime_error( + std::string("vt rocm: matmul_bt_quant: no keep-quant kernel for dtype ") + + Name(b.dtype) + + " (this discrete backend has no host fallback; set VT_GGUF_KEEP_QUANT=0 " + "to expand the file to bf16 at load)"); + } + if (k % kQK_K != 0) { + throw std::runtime_error( + "vt rocm: matmul_bt_quant: K must be a whole number of 256-element " + "Q8_K super-blocks"); + } + const int64_t nsb = k / kQK_K; + const size_t w_block_bytes = static_cast(vt::BlockBytes(b.dtype)); + const size_t w_row_bytes = static_cast(nsb) * w_block_bytes; + + const size_t act_bytes = + static_cast(m) * static_cast(nsb) * sizeof(BlockQ8_K); + BlockQ8_K* act = static_cast(EnsureScratch(act_bytes, s)); + LaunchQuantizeQ8K(act, a.data, ActDtOf(a.dtype), a.stride[0], m, nsb, s); + + const uint8_t* weight = static_cast(b.data); + switch (w) { + case WType::kIQ2_XXS: LaunchGemm(out, weight, act, m, n, nsb, w_row_bytes, w_block_bytes, s); break; + case WType::kIQ3_XXS: LaunchGemm(out, weight, act, m, n, nsb, w_row_bytes, w_block_bytes, s); break; + case WType::kQ2_K: LaunchGemm(out, weight, act, m, n, nsb, w_row_bytes, w_block_bytes, s); break; + case WType::kQ3_K: LaunchGemm(out, weight, act, m, n, nsb, w_row_bytes, w_block_bytes, s); break; + case WType::kQ4_K: LaunchGemm(out, weight, act, m, n, nsb, w_row_bytes, w_block_bytes, s); break; + case WType::kQ5_K: LaunchGemm(out, weight, act, m, n, nsb, w_row_bytes, w_block_bytes, s); break; + case WType::kQ6_K: LaunchGemm(out, weight, act, m, n, nsb, w_row_bytes, w_block_bytes, s); break; + case WType::kIQ2_S: LaunchGemm(out, weight, act, m, n, nsb, w_row_bytes, w_block_bytes, s); break; + case WType::kIQ1_S: LaunchGemm(out, weight, act, m, n, nsb, w_row_bytes, w_block_bytes, s); break; + case WType::kIQ1_XXXS: LaunchGemm(out, weight, act, m, n, nsb, w_row_bytes, w_block_bytes, s); break; + // IsRocmKeepQuantSupported answered yes, so a missing case must be LOUD: + // launching nothing leaves `out` stale while callers see success (#967). + default: + throw std::runtime_error( + std::string("vt rocm: matmul_bt_quant: no keep-quant kernel for dtype ") + + Name(b.dtype)); + } +} + +void MatmulBTQuantGroupedKernelRocm(Queue& q, Tensor& out, const Tensor& act, + const Tensor& weight, + const Tensor& expert_ids) { + hipStream_t s = static_cast(q.handle); + const int64_t P = out.shape[0]; + const int64_t n = out.shape[1]; + const int64_t k = act.shape[1]; + if (P == 0 || n == 0) return; + + // Delegate Q8_0/Q4_K/Q5_K/Q6_K to the optimized grouped kernels. Q8_0 has no + // Q8_K-superblock arm in this file (it dots a Q8_0 activation), so leaving it + // out of this list turns a format the *Gdn kernel serves today into a throw. + if (weight.dtype == DType::kQ8_0 || weight.dtype == DType::kQ4_K || + weight.dtype == DType::kQ5_K || weight.dtype == DType::kQ6_K) { + MatmulBTQuantGroupedKernelRocmGdn(q, out, act, weight, expert_ids); + return; + } + + WType w{}; + if (!IsRocmKeepQuantSupported(weight.dtype, &w)) { + throw std::runtime_error( + std::string("vt rocm: matmul_bt_quant_grouped: no keep-quant kernel for " + "dtype ") + + Name(weight.dtype) + + " (set VT_GGUF_KEEP_QUANT=0 to expand at load)"); + } + if (k % kQK_K != 0) { + throw std::runtime_error( + "vt rocm: matmul_bt_quant_grouped: K must be a whole number of " + "256-element Q8_K super-blocks"); + } + const int64_t nsb = k / kQK_K; + const size_t w_block_bytes = static_cast(vt::BlockBytes(weight.dtype)); + const size_t w_row_bytes = static_cast(nsb) * w_block_bytes; + + // Broadcast activation (preq-reuse): ONE quantized hidden feeds every routed + // expert slot; bit-identical because identical input yields identical Q8_K. + const int64_t Pa = act.shape[0]; + const bool bcast = (Pa == 1 && P > 1); + + const size_t act_bytes = + static_cast(Pa) * static_cast(nsb) * sizeof(BlockQ8_K); + BlockQ8_K* qact = static_cast(EnsureScratch(act_bytes, s)); + LaunchQuantizeQ8K(qact, act.data, ActDtOf(act.dtype), act.stride[0], Pa, nsb, + s); + + const uint8_t* wt = static_cast(weight.data); + const int32_t* eids = static_cast(expert_ids.data); + switch (w) { + case WType::kIQ2_XXS: LaunchGroupedGemm(out, wt, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast, s); break; + case WType::kIQ3_XXS: LaunchGroupedGemm(out, wt, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast, s); break; + case WType::kQ2_K: LaunchGroupedGemm(out, wt, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast, s); break; + case WType::kQ3_K: LaunchGroupedGemm(out, wt, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast, s); break; + case WType::kQ4_K: LaunchGroupedGemm(out, wt, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast, s); break; + case WType::kQ5_K: LaunchGroupedGemm(out, wt, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast, s); break; + case WType::kQ6_K: LaunchGroupedGemm(out, wt, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast, s); break; + case WType::kIQ2_S: LaunchGroupedGemm(out, wt, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast, s); break; + case WType::kIQ1_S: LaunchGroupedGemm(out, wt, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast, s); break; + case WType::kIQ1_XXXS: LaunchGroupedGemm(out, wt, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast, s); break; + default: + throw std::runtime_error( + std::string("vt rocm: matmul_bt_quant_grouped: no grouped kernel for " + "keep-quant dtype ") + + Name(weight.dtype)); + } +} + +// The two providers above are registered exactly once, by rocm_ops.hip's +// static registrar — this file deliberately carries no Registrar of its own. +// A second registration of the same OpId+DeviceType would make which registrar +// wins depend on static-init order across translation units. + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_rmsnorm.hip b/src/vt/rocm/rocm_rmsnorm.hip index f0aebeabd..aa2d183f6 100644 --- a/src/vt/rocm/rocm_rmsnorm.hip +++ b/src/vt/rocm/rocm_rmsnorm.hip @@ -1,28 +1,7 @@ // ROCm RmsNorm — the ONE op the W0 skeleton registers (BACKEND-ROCM, W0). // -// BUILD STATE, kept honest per issue #41, in the shape rocm_backend.hip uses. -// **THIS TU COMPILES.** `hipcc` (ROCm 7.2.4) built it clean for `gfx1151` with -// ZERO warnings on `strix:gpu0`, under `rc` job -// `d99a30fa-3336-4b90-bc9e-877717c9b54e`, from tree -// `2d1fdc19bd60cffec64fd2e66eb8abb5b5c1e1c5` -- the revision that gave this file -// its independent gamma type `Tw` (#2492). -// -// This line used to read "**UNBUILT.** Never compiled: no AMD GPU on the -// authoring machine", and that half is STILL TRUE of the authoring machine: -// there is no `hipcc` and no HIP header there, and no lane in -// `.github/workflows/` sets `VLLM_CPP_HIP`, which defaults to `AUTO` = OFF -// (CMakeLists.txt:124). The compile came from a leased device, not from CI, so -// this is a DATED REPORT ABOUT ONE REVISION and not a standing gate: if this -// file changes again, the report does not follow it, and the next editor is in -// the same position the previous one was. -// -// It is a COMPILE report and nothing more. No kernel in this file has been RUN -// on an AMD device. What makes even the compile mean something is the negative -// control the job ran: it broke the `case DType::kF16` arm, asserted the break -// applied, and required `hipcc` to REJECT the rebuild -- because a green `ninja` -// can otherwise mean "nothing to do". That, the five source guards, and the -// absent-object precondition are recorded in -// `.agents/specs/rmsnorm-gamma-dtype-twins.md`. +// **UNBUILT.** Never compiled: no AMD GPU on the authoring machine. See the +// header of rocm_backend.hip. // // PORTED FROM: src/vt/cuda/cuda_ops.cu:96-126 `RmsNormRowKernel` (one block per // row, shared-memory f32 tree reduction), whose own upstream counterpart is @@ -51,12 +30,20 @@ #include #include +#include #include +#include + +#include "vt/cpu/cpu_quant_blocks.h" +#include "vt/rocm/rocm_act_quant.h" +#include "vt/rocm/rocm_norm_quant_bridge.h" #include "vt/ops.h" #include "vt/rocm/rocm_device_bind.h" namespace vt::rocm { + +using vt::cpu::kQK_K; namespace { // Block width, from cuda_ops.cu:24. Kept at 256 rather than raised to a multiple @@ -69,13 +56,6 @@ __device__ inline float Load(const float* p, int64_t i) { return p[i]; } __device__ inline float Load(const __hip_bfloat16* p, int64_t i) { return __bfloat162float(p[i]); } -// GAMMA-ONLY element type. `vt::RmsNorm` admits `IsFloat(weight.dtype)` -// (ops.cpp:24, :1030), which includes kF16, and the CPU sibling widens a kF16 -// gamma like any other (`WidenRowToF32`, cpu_ops.cpp:554-557) -- so a device arm -// that refused it would be a divergence to record (cuda_qwen4_exp.cu:60-62). It -// is NOT admitted as an ACTIVATION (`Tin`) or as an output: `RmsNormKernelRocm` -// still refuses a kF16 `x`, which is a SEPARATE divergence of the same class, -// recorded rather than fixed here (#2542). __device__ inline float Load(const __half* p, int64_t i) { return __half2float(p[i]); } __device__ inline void Store(float* p, int64_t i, float v) { p[i] = v; } __device__ inline void Store(__hip_bfloat16* p, int64_t i, float v) { @@ -97,21 +77,17 @@ __device__ inline float ResRound<__hip_bfloat16>(float v) { return __bfloat162float(__float2bfloat16(v)); } -// `Tw` is the GAMMA's own dtype and is INDEPENDENT of `Tin` (#2492, the ROCm twin -// of #2477). It used to be `Tin`, which welded the two and forced -// `RmsNormKernelRocm` to refuse `w.dtype != x.dtype` -- a refusal that is a -// property of THIS kernel's signature and of nothing upstream. vLLM never couples -// them (`GemmaRMSNorm` reads `normalized * (1.0 + self.weight.float())`, -// vllm/models/qwen4_exp/nvidia/ple_layer.py:80 at vLLM origin/main cdefd9d499, -// which is AHEAD of this project's pin 5559679229 -- a forward reference, see -// #2502), this tree's CPU sibling widens `w` and `x` separately -// (cpu_ops.cpp:554-557, :576), and the CUDA sibling was decoupled by #2493 -// (cuda_ops.cu `RmsNormRowKernel`). `Load()` is overloaded for every gamma element -// type below, so the body is unchanged and its f32 arithmetic is bit-identical for -// every pairing that already worked. +// Lever C: `q8_out` (nullptr on every pre-existing path) turns the kernel into +// a PRODUCER of Q8_K activation scratch: after the output row is stored, one +// thread per 256-wide superblock requantizes the STORED bf16/f32 row through +// the SAME shared QuantQ8KSBlock body the standalone QuantizeQ8KK uses, so +// the scratch is byte-identical to the standalone path by construction +// (asserted in tests/vt/test_rocm_quant_dot.cpp). Requires nsb <= kBlock: one +// thread per superblock within this single-row block. template __global__ void RmsNormRowKernel(Tout* out, const Tin* x, const Tw* w, Tres* residual, int64_t h, - float eps, bool gemma) { + float eps, bool gemma, vt::cpu::BlockQ8_K* q8_out, int nsb, + ActDT q8_adt) { const int64_t row = blockIdx.x; const Tin* xrow = x + row * h; Tout* orow = out + row * h; @@ -140,38 +116,342 @@ __global__ void RmsNormRowKernel(Tout* out, const Tin* x, const Tw* w, Tres* res if (gemma) wj += 1.0f; Store(orow, j, v * inv * wj); } + if (q8_out != nullptr) { + // Barrier first: the epilogue reads the WHOLE stored row (written by all + // 256 threads) back from global memory, exactly as the standalone + // QuantizeQ8KK would -- same bytes, same serial per-superblock walk, so + // the amax first-occurrence tie-break and every scale bit match. + __syncthreads(); + const int sb = static_cast(threadIdx.x); + if (sb < nsb) + QuantQ8KSBlock(q8_out[static_cast(row) * nsb + sb], orow, q8_adt, + static_cast(sb) * kQK_K); + } +} + +// --- T8 (GFX1100-TG200): cooperative single-row remap ----------------------- +// Decode launches ONE 256-thread block per row; the ported body chains three +// strided scalar passes, a nine-step __syncthreads shared-memory tree, and a +// fused-q8 epilogue whose per-superblock QuantQ8KSBlock walks 256 elements +// serially on ONE thread (nsb<=10 of 256 threads active). rocpd prices the +// fused instantiation at 65 launches/tok x 18.1us = 1.178 ms/tok against a +// microsecond-class floor. This arm keeps the grid shape and the byte +// contracts but rebuilds the internals, behind VT_RMSNORM_ROW_COOP=1 +// (default OFF leaves every path byte-unchanged): +// +// 1. Reduction: two wavefront shfl_down trees + one cross-wavefront +// combine through shared memory -- TWO __syncthreads instead of nine. +// The float association CHANGES (recorded adjudication required; the +// opt-in flag rides the campaign config like GDN_SCAN_COOP). +// 2. Vector passes: 16-byte loads/stores where the row base allows, +// scalar fallback otherwise (uniform per launch). +// 3. Cooperative q8 epilogue: the whole block quantizes ONE superblock +// at a time, thread i owning element i. BYTE CONTRACT PRESERVED BY +// CONSTRUCTION: the (mx, amax) pair comes from a LEFT-BIASED max over +// ascending element positions (shfl trees keep the earlier element on +// magnitude ties, which is exactly the scalar scan's first-occurrence +// rule), iscale/DNearestInt/clamp arithmetic is verbatim, and bsums +// stay exact integer sums. Asserted against the standalone quantizer +// by the focused test under BOTH flag states. +template +__global__ void RmsNormRowCoopKernel(Tout* out, const Tin* x, const Tw* w, Tres* residual, + int64_t h, float eps, bool gemma, + vt::cpu::BlockQ8_K* q8_out, int nsb, ActDT q8_adt, + bool lds_quant) { + const int64_t row = blockIdx.x; + const int tid = static_cast(threadIdx.x); + const Tin* xrow = x + row * h; + Tout* orow = out + row * h; + Tres* rrow = residual == nullptr ? nullptr : residual + row * h; + const int waveSz = warpSize; // runtime width (32 on RDNA default); never hardcode + const int kWaves = kBlock / waveSz; + __shared__ float partial[kBlock]; + // Statically sized for the narrowest wavefront (32); kWaves is runtime. + __shared__ float wave_mx[kBlock / 32]; + __shared__ float wave_amax[kBlock / 32]; + extern __shared__ char dyn_smem[]; + Tout* lds_out = reinterpret_cast(dyn_smem); + // The vector body decodes the gamma as 8 bf16 halves per uint4, so a + // non-bf16 gamma (kF32/kF16 are admitted by DispatchRmsNormWeight, + // independent of Tin) MUST take the scalar path, whose Load() overloads + // decode per dtype. sizeof(Tw)==2 alone would still misread an f16 gamma. + const bool vec = + std::is_same_v && sizeof(Tin) == 2 && + (rrow == nullptr || sizeof(Tres) == 2) && + (reinterpret_cast(xrow) & 15u) == 0 && + (reinterpret_cast(orow) & 15u) == 0 && + (rrow == nullptr || (reinterpret_cast(rrow) & 15u) == 0); + + // Pass 1: sum of squares (+ residual fold-in). Per-thread element set is + // the plain kernel's; the vector path regroups WITHIN a thread's step. + float acc = 0.0f; + if (vec && sizeof(Tin) == 2) { + const uint4* xv = reinterpret_cast(xrow); + uint4* rv = reinterpret_cast(rrow); + const int step8 = kBlock * 8; // elements per block-step (bf16) + for (int j8 = tid * 8; j8 < static_cast(h); j8 += step8) { + const int elems = j8 + 8 <= static_cast(h) ? 8 : static_cast(h) - j8; + if (elems == 8) { + const uint4 rawx = xv[j8 / 8]; + uint16_t hx[8] = {static_cast(rawx.x & 0xFFFF), static_cast(rawx.x >> 16), + static_cast(rawx.y & 0xFFFF), static_cast(rawx.y >> 16), + static_cast(rawx.z & 0xFFFF), static_cast(rawx.z >> 16), + static_cast(rawx.w & 0xFFFF), static_cast(rawx.w >> 16)}; + float v[8]; +#pragma unroll + for (int u = 0; u < 8; ++u) { + float fv = DBF16ToF32(hx[u]); + if (rrow != nullptr) { + fv = ResRound(fv + Load(rrow, j8 + u)); + } + v[u] = fv; + } + if (rrow != nullptr) { + uint4 rout; + uint16_t hout[8]; + for (int u = 0; u < 8; ++u) hout[u] = DF32ToBF16(v[u]); + rout.x = static_cast(hout[0]) | (static_cast(hout[1]) << 16); + rout.y = static_cast(hout[2]) | (static_cast(hout[3]) << 16); + rout.z = static_cast(hout[4]) | (static_cast(hout[5]) << 16); + rout.w = static_cast(hout[6]) | (static_cast(hout[7]) << 16); + rv[j8 / 8] = rout; + } +#pragma unroll + for (int u = 0; u < 8; ++u) acc += v[u] * v[u]; + } else { + for (int j = j8; j < j8 + elems; ++j) { + float v = Load(xrow, j); + if (rrow != nullptr) { + v = ResRound(v + Load(rrow, j)); + Store(rrow, j, v); + } + acc += v * v; + } + } + } + } else { + for (int64_t j = tid; j < h; j += kBlock) { + float v = Load(xrow, j); + if (rrow != nullptr) { + v = ResRound(v + Load(rrow, j)); + Store(rrow, j, v); + } + acc += v * v; + } + } + + // Two-level reduction: wavefront shfl trees, one cross-wavefront combine. + float wacc = acc; +#pragma unroll + for (int off = waveSz / 2; off > 0; off >>= 1) + wacc += __shfl_down_sync(0xffffffffULL, wacc, off); + if (tid % waveSz == 0) partial[tid / waveSz] = wacc; + __syncthreads(); + if (tid == 0) { + float t = partial[0]; +#pragma unroll + for (int wv = 1; wv < kWaves; ++wv) t += partial[wv]; + partial[0] = t; + } + __syncthreads(); + const float inv = 1.0f / sqrtf(partial[0] / static_cast(h) + eps); + + // Pass 3: scaled output. + if (vec && sizeof(Tin) == 2 && sizeof(Tout) == 2) { + const uint4* xv = reinterpret_cast( + rrow != nullptr ? static_cast(rrow) : static_cast(xrow)); + uint4* ov = reinterpret_cast(orow); + const uint4* wv = reinterpret_cast(w); + const bool waligned = (reinterpret_cast(w) & 15u) == 0; + for (int j8 = tid * 8; j8 < static_cast(h); j8 += kBlock * 8) { + const int elems = j8 + 8 <= static_cast(h) ? 8 : static_cast(h) - j8; + if (elems == 8 && waligned) { + const uint4 rawx = xv[j8 / 8]; + const uint4 raww = wv[j8 / 8]; + uint16_t hx[8] = {static_cast(rawx.x & 0xFFFF), static_cast(rawx.x >> 16), + static_cast(rawx.y & 0xFFFF), static_cast(rawx.y >> 16), + static_cast(rawx.z & 0xFFFF), static_cast(rawx.z >> 16), + static_cast(rawx.w & 0xFFFF), static_cast(rawx.w >> 16)}; + uint16_t hw[8] = {static_cast(raww.x & 0xFFFF), static_cast(raww.x >> 16), + static_cast(raww.y & 0xFFFF), static_cast(raww.y >> 16), + static_cast(raww.z & 0xFFFF), static_cast(raww.z >> 16), + static_cast(raww.w & 0xFFFF), static_cast(raww.w >> 16)}; + uint16_t ho[8]; +#pragma unroll + for (int u = 0; u < 8; ++u) { + float wj = DBF16ToF32(hw[u]); + if (gemma) wj += 1.0f; + ho[u] = DF32ToBF16(DBF16ToF32(hx[u]) * inv * wj); + } + uint4 o; + o.x = static_cast(ho[0]) | (static_cast(ho[1]) << 16); + o.y = static_cast(ho[2]) | (static_cast(ho[3]) << 16); + o.z = static_cast(ho[4]) | (static_cast(ho[5]) << 16); + o.w = static_cast(ho[6]) | (static_cast(ho[7]) << 16); + ov[j8 / 8] = o; + if (q8_out != nullptr && lds_quant) { +#pragma unroll + for (int u = 0; u < 8; ++u) + *reinterpret_cast(lds_out + j8 + u) = ho[u]; + } + } else { + for (int j = j8; j < j8 + elems; ++j) { + const float v = rrow != nullptr ? Load(rrow, j) : Load(xrow, j); + float wj = Load(w, j); + if (gemma) wj += 1.0f; + Store(orow, j, v * inv * wj); + if (q8_out != nullptr && lds_quant) lds_out[j] = static_cast(v * inv * wj); + } + } + } + } else { + for (int64_t j = tid; j < h; j += kBlock) { + const float v = rrow != nullptr ? Load(rrow, j) : Load(xrow, j); + float wj = Load(w, j); + if (gemma) wj += 1.0f; + Store(orow, j, v * inv * wj); + if (q8_out != nullptr && lds_quant) lds_out[j] = static_cast(v * inv * wj); + } + } + + if (q8_out != nullptr) { + // Barrier: the epilogue reads the WHOLE stored row back, exactly as the + // standalone QuantizeQ8KK would. + __syncthreads(); + __shared__ int8_t sq[kQK_K]; + __shared__ float lead_iscale; + const void* quant_src = lds_quant ? static_cast(lds_out) + : static_cast(orow); + for (int sb = 0; sb < nsb; ++sb) { + vt::cpu::BlockQ8_K& y = q8_out[static_cast(row) * nsb + sb]; + const float x = DLoadAct(quant_src, q8_adt, static_cast(sb) * kQK_K + tid); + // Seed 0 like the standalone QuantizeQ8KK (strict > keeps the first + // occurrence and rejects NaN): seeding from x let a NaN element poison + // d and diverge from the scalar contract's zero block. + float mx = 0.0f, amax = 0.0f; + if (x == x) { mx = x; amax = fabsf(x); } + for (int off = waveSz / 2; off > 0; off >>= 1) { + const float ox = __shfl_down_sync(0xffffffffULL, mx, off); + const float oa = __shfl_down_sync(0xffffffffULL, amax, off); + if (oa > amax) { amax = oa; mx = ox; } + } + if (tid % waveSz == 0) { + wave_mx[tid / waveSz] = mx; + wave_amax[tid / waveSz] = amax; + } + __syncthreads(); + if (tid == 0) { + float bmx = wave_mx[0], bamax = wave_amax[0]; +#pragma unroll + for (int wv = 1; wv < kWaves; ++wv) { + if (wave_amax[wv] > bamax) { bamax = wave_amax[wv]; bmx = wave_mx[wv]; } + } + lead_iscale = bamax == 0.0f ? 0.0f : -127.0f / bmx; + } + __syncthreads(); + if (lead_iscale == 0.0f) { + if (tid < kQK_K) y.qs[tid] = 0; + if (tid < kQK_K / 16) y.bsums[tid] = 0; + if (tid == 0) y.d = 0.0f; + } else { + const int qv = DNearestInt(lead_iscale * x); + const int8_t q = static_cast(qv < 127 ? qv : 127); + y.qs[tid] = q; + sq[tid] = q; + __syncthreads(); + if (tid < kQK_K / 16) { + int sum = 0; +#pragma unroll + for (int ii = 0; ii < 16; ++ii) sum += sq[tid * 16 + ii]; + y.bsums[tid] = static_cast(sum); + } + if (tid == 0) y.d = 1.0f / lead_iscale; + } + __syncthreads(); // sq reuse guard across superblocks + } + } } template void LaunchRmsNormRes(hipStream_t s, Tensor& out, const Tensor& x, const Tensor& w, - const RmsNormArgs& args, Tensor* residual, unsigned rows, int64_t h) { + const RmsNormArgs& args, Tensor* residual, unsigned rows, int64_t h, + vt::cpu::BlockQ8_K* q8_out) { + const int nsb = q8_out != nullptr ? static_cast(h / kQK_K) : 0; + // T8 opt-in arm (read PER CALL like the sibling flags): the cooperative + // remap requires one 256-wide superblock per thread and h a multiple of + // kQK_K; anything else falls back to the ported body. + const char* coop_e = std::getenv("VT_RMSNORM_ROW_COOP"); + const bool coop = coop_e != nullptr && coop_e[0] == '1' && coop_e[1] == '\0' && + h % kQK_K == 0 && h / kQK_K <= kBlock; + if (coop) { + // T24: LDS-buffered quant epilogue (VT_RMSNORM_LDS_QUANT, default ON). + // Stores the normalized row to LDS during Pass 3 so the Q8_K quant + // epilogue reads from LDS instead of re-reading global memory. + const char* lds_e = std::getenv("VT_RMSNORM_LDS_QUANT"); + const bool lds_quant = + lds_e == nullptr || lds_e[0] == '1'; // default ON when q8_out active + const size_t smem = + (q8_out != nullptr && lds_quant) ? static_cast(h) * sizeof(Tout) : 0; + if (residual != nullptr && residual->dtype == DType::kBF16) { + RmsNormRowCoopKernel<<>>( + out.Ptr(), x.Ptr(), w.Ptr(), residual->Ptr<__hip_bfloat16>(), h, + args.eps, args.gemma, q8_out, nsb, ActDtOf(out.dtype), lds_quant); + } else { + float* res = residual == nullptr ? nullptr : residual->Ptr(); + RmsNormRowCoopKernel<<>>( + out.Ptr(), x.Ptr(), w.Ptr(), res, h, args.eps, args.gemma, q8_out, + nsb, ActDtOf(out.dtype), lds_quant); + } + return; + } if (residual != nullptr && residual->dtype == DType::kBF16) { RmsNormRowKernel<<>>( out.Ptr(), x.Ptr(), w.Ptr(), residual->Ptr<__hip_bfloat16>(), h, args.eps, - args.gemma); + args.gemma, q8_out, nsb, ActDtOf(out.dtype)); } else { float* res = residual == nullptr ? nullptr : residual->Ptr(); RmsNormRowKernel<<>>( - out.Ptr(), x.Ptr(), w.Ptr(), res, h, args.eps, args.gemma); + out.Ptr(), x.Ptr(), w.Ptr(), res, h, args.eps, args.gemma, q8_out, + nsb, ActDtOf(out.dtype)); } } +// Lever C opt-in gate: VT_NORM_QUANT_FUSED=1 (read PER CALL like the sibling +// arms' flags so in-process tests and graph capture pick it up at dispatch +// time), dense rows, whole Q8_K superblocks, and at most one thread's worth of +// superblocks per row. Default OFF leaves every path byte-unchanged. +vt::cpu::BlockQ8_K* NormQuantEpilogueFor(int64_t rows, int64_t h, hipStream_t s, + size_t* scratch_bytes) { + const char* e = std::getenv("VT_NORM_QUANT_FUSED"); + if (e == nullptr || e[0] != '1' || e[1] != '\0') return nullptr; + if (rows <= 0 || h % kQK_K != 0 || h / kQK_K > kBlock) return nullptr; + *scratch_bytes = static_cast(rows) * static_cast(h / kQK_K) * + sizeof(vt::cpu::BlockQ8_K); + return static_cast(NormQuantProducerScratch(*scratch_bytes, s)); +} + template void LaunchRmsNorm(hipStream_t s, Tensor& out, const Tensor& x, const Tensor& w, const RmsNormArgs& args, Tensor* residual) { const int64_t t = x.shape[0], h = x.shape[1]; if (t == 0 || h == 0) return; const unsigned rows = static_cast(t); + size_t q8_bytes = 0; + vt::cpu::BlockQ8_K* q8_out = + NormQuantEpilogueFor(t, h, s, &q8_bytes); // nullptr unless opt-in flag switch (out.dtype) { case DType::kF32: - LaunchRmsNormRes(s, out, x, w, args, residual, rows, h); + LaunchRmsNormRes(s, out, x, w, args, residual, rows, h, q8_out); break; case DType::kBF16: - LaunchRmsNormRes(s, out, x, w, args, residual, rows, h); + LaunchRmsNormRes(s, out, x, w, args, residual, rows, h, q8_out); break; default: VT_CHECK(false, "rocm rmsnorm: unsupported out dtype (f32/bf16 only)"); } + if (q8_out != nullptr) { + NormQuantRecordProducer(out.data, t, h, /*row_stride=*/h, out.dtype, q8_out, s); + } } // The gamma's dtype is dispatched SEPARATELY from the activation's, mirroring diff --git a/src/vt/rocm/rocm_skinny_gemm.hip b/src/vt/rocm/rocm_skinny_gemm.hip index b4d7dd755..b55a2c4aa 100644 --- a/src/vt/rocm/rocm_skinny_gemm.hip +++ b/src/vt/rocm/rocm_skinny_gemm.hip @@ -52,11 +52,28 @@ inline int mindiv(int N, int div1, int div2) { } // A (activation) fits LDS. N = decode batch (our M). bf16, f32 accum. -template +// Output-store polymorphism: the adopted bf16 decode arm rounds the f32 +// accumulator to bf16; lever B2 (evidence 15.1) serves the Qwen3.5 GDN BA +// projections, which must EMIT f32 — same geometry/reduction tree, only the +// store type differs. No rounding on the f32 arm. +template +__device__ __forceinline__ OutT WvStoreCast(float v); +template <> +__device__ __forceinline__ __hip_bfloat16 WvStoreCast<__hip_bfloat16>(float v) { + return __float2bfloat16(v); +} +template <> +__device__ __forceinline__ float WvStoreCast(float v) { + return v; +} + +// A (activation) fits LDS. N = decode batch (our M). bf16 in, OutT out, +// f32 accum. +template __global__ void __launch_bounds__(kWvPrGrp * kThrds) wvSplitKSml(const int K, const int Kbp, const int Kap, const int M, const __hip_bfloat16* __restrict__ B, const __hip_bfloat16* __restrict__ A, - __hip_bfloat16* C, const int _WvPrGrp, const int CuCount) { + OutT* C, const int _WvPrGrp, const int CuCount) { constexpr int max_lds_len = kLdsSize / 2; // bf16 elements __shared__ __hip_bfloat16 s[max_lds_len]; @@ -119,7 +136,8 @@ __global__ void __launch_bounds__(kWvPrGrp * kThrds) } if (threadIdx.x == (kThrds - 1)) { for (int n = 0; n < N; n++) - for (int y = 0; y < kYtile; y++) C[m + y + n * M] = __float2bfloat16(sum[n][y]); + for (int y = 0; y < kYtile; y++) + C[m + y + n * M] = WvStoreCast(sum[n][y]); } m += CuCount * _WvPrGrp * kYtile; } @@ -139,31 +157,30 @@ int DeviceCuCount(int device) { return cache; } -} // namespace -// out[M,N] = a[M,K] @ b[N,K]^T, bf16 in/out. Only called for the decode-skinny -// gate below (M in 1..4, K%8==0); the caller keeps every other shape on the -// BLAS path. -void WvSplitKBT(hipStream_t s, void* out, const void* a, const void* b, int M, int N, - int K, int device) { +// Shared launch math; the store dtype is the only difference between the +// adopted bf16 arm (WvSplitKBT) and lever B2's f32-out arm (WvSplitKBTToF32). +template +void WvSplitKBTLaunch(hipStream_t s, void* out, const void* a, const void* b, + int M, int N, int K, int device) { const int cu = DeviceCuCount(device); dim3 grid(cu), block(kThrds, kWvPrGrp); const int wvPrGrp = mindiv(N, cu * kYtile, kWvPrGrp); - auto* C = static_cast<__hip_bfloat16*>(out); + auto* C = static_cast(out); auto* A = static_cast(a); auto* B = static_cast(b); switch (M) { case 1: - wvSplitKSml<1><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + wvSplitKSml<1, OutT><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); break; case 2: - wvSplitKSml<2><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + wvSplitKSml<2, OutT><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); break; case 3: - wvSplitKSml<3><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + wvSplitKSml<3, OutT><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); break; case 4: - wvSplitKSml<4><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); + wvSplitKSml<4, OutT><<>>(K, K, K, N, B, A, C, wvPrGrp, cu); break; default: throw std::runtime_error("vt rocm: wvSplitK unsupported M=" + std::to_string(M)); @@ -172,4 +189,22 @@ void WvSplitKBT(hipStream_t s, void* out, const void* a, const void* b, int M, i throw std::runtime_error("vt rocm: wvSplitK launch failed"); } +} // namespace + +// out[M,N] = a[M,K] @ b[N,K]^T, bf16 in/out. Only called for the decode-skinny +// gate below (M in 1..4, K%8==0); the caller keeps every other shape on the +// BLAS path. +void WvSplitKBT(hipStream_t s, void* out, const void* a, const void* b, int M, int N, + int K, int device) { + WvSplitKBTLaunch<__hip_bfloat16>(s, out, a, b, M, N, K, device); +} + +// Lever B2 (evidence 15.1): same geometry, f32 OUTPUT — the Qwen3.5 GDN BA +// projections consume f32 (g/beta chain) and today starve on rocBLAS's +// large-M tile. Dispatched by MatmulBTKernelRocm under VT_SKINNY_BF16=1. +void WvSplitKBTToF32(hipStream_t s, void* out, const void* a, const void* b, + int M, int N, int K, int device) { + WvSplitKBTLaunch(s, out, a, b, M, N, K, device); +} + } // namespace vt::rocm diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index a3117cc61..21ea59959 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -2666,6 +2666,14 @@ vllm_cpp_add_test(test_rocm_fp8_kv_cache vt/test_rocm_fp8_kv_cache.cpp) # the multiply, matching the CPU oracle RoundThrough and upstream vLLM silu_kernel. # Self-skipping without a ROCm device, mirroring test_rocm_backend.cpp guard. vllm_cpp_add_test(test_ops_rocm_silu_rounding vt/test_ops_rocm_silu_rounding.cpp) +# KERNEL-QUANT-CIQ-GEMM-ROCM: the ROCm keep-quant GEMM (kROCM provider for +# kMatmulBTQuant/kMatmulBTQuantGrouped). Gates the device dequant-in-kernel dot +# against the CPU keep-quant oracle + an f64 dequant reference on the ten +# Q8_K-family encodings. Skips coherently with no AMD GPU. +vllm_cpp_add_test(test_rocm_quant_dot vt/test_rocm_quant_dot.cpp) +target_include_directories(test_rocm_quant_dot PRIVATE ${CMAKE_SOURCE_DIR}/src) +vllm_cpp_add_test(test_rocm_skinny_f32 vt/test_rocm_skinny_f32.cpp) +target_include_directories(test_rocm_skinny_f32 PRIVATE ${CMAKE_SOURCE_DIR}/src) # #785 P1 GPU product-seam witness. Executable only — NOT add_test. # Ordinary CTest must not see this target. Runner fail-closes on 77/nonzero. add_executable(test_ops_paged_attn_sharedk_wmma_p1_gpu diff --git a/tests/vt/test_rocm_quant_dot.cpp b/tests/vt/test_rocm_quant_dot.cpp new file mode 100644 index 000000000..d22efaa29 --- /dev/null +++ b/tests/vt/test_rocm_quant_dot.cpp @@ -0,0 +1,1330 @@ +// ROCm keep-quant GEMM gate (GFX1100-TG200). The campaign spec names +// `tests/vt/test_rocm_quant_dot.cpp` as the quant-path lever gate; until T4a +// that file DID NOT EXIST — the GPU-parity cases lived in +// tests/vt/test_cuda_quant_dot.cpp behind HasCuda() and so SKIPPED on this +// ROCm-only box (the exact T3a blind spot: op-level green while the engine +// produced garbage). This file is the fix: a focused gate for the ROCm +// kMatmulBTQuant provider (src/vt/rocm/rocm_grouped_gemm.hip) guarded on ROCM +// availability, never on CUDA. +// +// RED-first contract: before the dispatch arm exists VT_GEMV_MMVQ=1 is inert, +// so ON==OFF trivially; the dispatch-gate cases below fail if the flag never +// engages the arm. +// +// T4a REPAIR ROUND numerics contract: the arm must be BYTE-IDENTICAL TO THE +// DEFAULT (warp-reduction) KERNEL — the engine-safety property the FIRST +// round lacked. Round 1 was bit-exact vs the CPU ORACLE while the ENGINE +// degraded: oracle association != baseline tree association, and greedy +// near-ties flipped (extended ON-vs-OFF sweep red at N=2304..248320, +// isolated first-diverging rows). This gate therefore asserts ON==OFF raw +// byte identity on every case below, PLUS the standard 1e-6 NMSE band vs +// the CPU oracle for the ON arm (the same band the default arm is held to). +// +// Skips cleanly (returns) when the build has HIP but the box has no AMD GPU, +// so the CPU CI leg stays green. +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "vt/backend.h" +#include "vt/device.h" +#include "vt/dtype.h" +#include "vt/ops.h" +#include "vt/quant.h" +#include "vt/rocm/rocm_runtime.h" +#include "vt/tensor.h" + +using vt::Backend; +using vt::Device; +using vt::DeviceType; +using vt::DType; +using vt::Queue; +using vt::Tensor; + +namespace vt::rocm { +void MmvqQuantScratchForTesting(Queue& q, void* dst, const Tensor& a, + bool fused_semantics); + +// T4a REPAIR-ROUND-2 routing witness (review findings F1/F2): the HOST-side +// dispatch counters exposed by rocm_grouped_gemm.hip. ON and OFF arms are +// BIT-EQUAL on outputs by design, so no output comparison can witness which +// dispatch branch a call took -- these integer counters can. +struct MmvqRouteCounts { + long long baseline; // KQuantGemmK warp-reduction dispatches + long long gemv_mmvq; // non-fused MMVQ GEMV dispatches (standalone quant) + long long gemv_fused; // fused-fold sub-branch dispatches +}; +// Lever C (GFX1100-TG200-NORMQ): producer-fused Q8_K norm-epilogue witnesses. +// The RmsNormRowKernel producer emits the row's Q8_K blocks alongside its +// normal output under VT_NORM_QUANT_FUSED=1 and records a producer token; +// MatmulBTQuant's K-quant branch SKIPS the standalone QuantizeQ8KK when the +// consuming activation matches that token. These counters make the ROUTE +// observable (outputs are bit-equal either way by contract). +struct NormQuantCounts { + long long producers; // epilogue-enabled RmsNorm dispatches + long long consumers_fused; // K-quant matvec dispatches that skipped the standalone quant + long long consumers_standalone; // K-quant matvec dispatches that launched QuantizeQ8KK +}; +NormQuantCounts NormQuantCountsForTesting(); +void NormQuantResetForTesting(); +// Device pointer of the Q8_K scratch written by the LAST producer-fused +// RmsNorm dispatch (rows * (h/256) BlockQ8_K blocks) -- lets tests assert the +// epilogue bytes are IDENTICAL to the standalone quantizer's. +const void* NormQuantLastScratchForTesting(); +MmvqRouteCounts MmvqRouteCountsForTesting(); +void MmvqResetRouteCountsForTesting(); +} // namespace vt::rocm + +namespace { +Device Cpu() { return Device{DeviceType::kCPU, 0}; } +Device GpuDev() { return Device{DeviceType::kROCM, 0}; } + +constexpr double kMaxNmseErr = 5e-4; // test-backend-ops.cpp:4277 band +constexpr double kMaxNmseVsCpu = 1e-6; // integer core exact; scale sum only + +struct WeightCase { + DType dtype; + int64_t block_elems; + int64_t block_bytes; + int d_off; + int dmin_off; + const char* name; + // f64-dequant ceiling override (0 = kMaxNmseErr); see the CUDA table for why + // the IQ1 family needs a wider ACTIVATION-error band while the ROCm-vs-CPU + // bound below stays shared and unrelaxed. + double nmse_ref_max = 0.0; +}; + +// The ten Q8_K-family encodings the CUDA sibling serves (test_cuda_quant_dot +// .cpp's WeightCase table), plus Q8_0 which the *Gdn kernels serve. +const WeightCase kCases[] = { + {DType::kIQ2_XXS, 256, 66, 0, -1, "iq2_xxs"}, + {DType::kIQ3_XXS, 256, 98, 0, -1, "iq3_xxs"}, + {DType::kIQ2_S, 256, 82, 0, -1, "iq2_s"}, + {DType::kIQ1_S, 256, 50, 0, -1, "iq1_s", 2e-3}, + {DType::kIQ1_XXXS, 256, 38, 0, -1, "iq1_xxxs", 2e-3}, + {DType::kQ2_K, 256, 84, 80, 82, "q2_K"}, + {DType::kQ3_K, 256, 110, 108, -1, "q3_K"}, + {DType::kQ4_K, 256, 144, 0, 2, "q4_K"}, + {DType::kQ5_K, 256, 176, 0, 2, "q5_K"}, + {DType::kQ6_K, 256, 210, 208, -1, "q6_K"}, + // Q8_0 is NOT a Q8_K-superblock encoding: it dots a Q8_0 activation and is + // served by the *Gdn kernels this file delegates to. + {DType::kQ8_0, 32, 34, 0, -1, "q8_0"}, +}; + +// Same table discipline as test_cuda_quant_dot.cpp:113 (offsets restated from +// ggml-common.h): the three K-quants the ROCm provider serves natively. +const WeightCase kKQuantCases[] = { + {DType::kQ4_K, 256, 144, 0, 2, "q4_K"}, + {DType::kQ5_K, 256, 176, 0, 2, "q5_K"}, + {DType::kQ6_K, 256, 210, 208, -1, "q6_K"}, +}; + +void GenerateData(float offset, size_t n, float* dst) { + for (size_t i = 0; i < n; i++) + dst[i] = 0.1F + 2 * std::cos(static_cast(i) + offset); +} + +std::vector RandomBlocks(const WeightCase& c, int64_t nblocks, + uint32_t seed) { + std::mt19937 rng(seed); + std::vector bytes(static_cast(nblocks * c.block_bytes)); + for (uint8_t& b : bytes) b = static_cast(rng() & 0xFF); + for (int64_t i = 0; i < nblocks; ++i) { + uint8_t* blk = bytes.data() + i * c.block_bytes; + auto put_f16 = [&](int off, float v) { + const uint16_t h = vt::F32ToF16(v); + std::memcpy(blk + off, &h, sizeof(h)); + }; + const float jitter = 1.0F + 0.05F * static_cast(i % 7); + if (c.d_off >= 0) put_f16(c.d_off, 0.0125F * jitter); + if (c.dmin_off >= 0) put_f16(c.dmin_off, 0.0075F * jitter); + // IQ1 sub-block scales live INSIDE the weight (qh bits 12-14 / sc nibbles): + // narrow them to encoder-plausible values exactly as the CUDA table does. + if (c.dtype == DType::kIQ1_S) { + for (int ib = 0; ib < 8; ++ib) { + uint16_t qh = 0; + std::memcpy(&qh, blk + 34 + 2 * ib, sizeof(qh)); + const uint16_t ls = static_cast(2 + ((i + ib) % 3)); + qh = static_cast((qh & 0x8FFFU) | (ls << 12)); + std::memcpy(blk + 34 + 2 * ib, &qh, sizeof(qh)); + } + } + if (c.dtype == DType::kIQ1_XXXS) { + for (int ib = 0; ib < 8; ++ib) { + uint8_t& byte = blk[34 + ib / 2]; + const int shift = 4 * (ib & 1); + const uint8_t ls = static_cast(2 + ((i + ib) % 3)); + const uint8_t keep_sign = static_cast((byte >> shift) & 0x8); + byte = static_cast((byte & ~(0xFU << shift)) | + ((keep_sign | ls) << shift)); + } + } + } + return bytes; +} + +double Nmse(const std::vector& got, const std::vector& ref) { + double num = 0, den = 0; + for (size_t i = 0; i < ref.size(); ++i) { + const double d = static_cast(got[i]) - static_cast(ref[i]); + num += d * d; + den += static_cast(ref[i]) * static_cast(ref[i]); + } + return num / den; +} + +Tensor DevTensor(void* p, DType dt, const std::vector& shape) { + Tensor t; + t.data = p; + t.dtype = dt; + t.device = GpuDev(); + t.rank = static_cast(shape.size()); + int64_t stride = 1; + for (int i = t.rank - 1; i >= 0; --i) { + t.shape[i] = shape[static_cast(i)]; + t.stride[i] = stride; + stride *= shape[static_cast(i)]; + } + return t; +} + +struct EnvGuard { + explicit EnvGuard(bool on) { ::setenv("VT_GEMV_MMVQ", on ? "1" : "0", 1); } + ~EnvGuard() { ::unsetenv("VT_GEMV_MMVQ"); } +}; + +} // namespace + +TEST_CASE("ROCm keep-quant GEMM == CPU reference and f64 dequant (Q8_K family)") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no ROCm backend on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + Queue cq{Cpu(), nullptr}; + + for (const WeightCase& c : kCases) { + const int64_t k = 8 * c.block_elems; + for (int64_t m : {int64_t{1}, int64_t{4}, int64_t{32}, int64_t{512}}) { + for (int64_t n : {int64_t{1}, int64_t{7}, int64_t{16}}) { + CAPTURE(std::string(c.name)); + CAPTURE(m); + CAPTURE(k); + CAPTURE(n); + + std::vector wq = + RandomBlocks(c, n * (k / c.block_elems), 0x5EEDU); + std::vector a(static_cast(m * k)); + GenerateData(1.0F, a.size(), a.data()); + + // --- CPU oracle (the landed keep-quant kernel over host tensors) ------ + std::vector cpu_out(static_cast(m * n), 0.0F); + { + Tensor at = Tensor::Contiguous(a.data(), DType::kF32, Cpu(), {m, k}); + Tensor bt = + Tensor::Contiguous(wq.data(), DType::kF32, Cpu(), {n, k}); + bt.dtype = c.dtype; + Tensor ot = + Tensor::Contiguous(cpu_out.data(), DType::kF32, Cpu(), {m, n}); + vt::MatmulBTQuant(cq, ot, at, bt); + } + + // --- ROCm path (device tensors; discrete card, so real staging) ------ + void* d_a = gpu.Alloc(a.size() * sizeof(float)); + void* d_w = gpu.Alloc(wq.size()); + void* d_o = gpu.Alloc(static_cast(m * n) * sizeof(float)); + gpu.Copy(gq, d_a, a.data(), a.size() * sizeof(float)); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + Tensor at = DevTensor(d_a, DType::kF32, {m, k}); + Tensor bt = DevTensor(d_w, c.dtype, {n, k}); + Tensor ot = DevTensor(d_o, DType::kF32, {m, n}); + vt::MatmulBTQuant(gq, ot, at, bt); + std::vector rocm_out(static_cast(m * n), 0.0F); + gpu.Copy(gq, rocm_out.data(), d_o, rocm_out.size() * sizeof(float)); + gpu.Synchronize(gq); + gpu.Free(d_a); + gpu.Free(d_w); + gpu.Free(d_o); + + // --- f64 independent reference -------------------------------------- + std::vector w(static_cast(n * k)); + vt::cpu::BlockToFloat(c.dtype)(wq.data(), w.data(), n * k); + + double num_ref = 0, den_ref = 0, num_cpu = 0, den_cpu = 0; + for (int64_t i = 0; i < m; ++i) { + for (int64_t jj = 0; jj < n; ++jj) { + double ref = 0; + for (int64_t p = 0; p < k; ++p) + ref += static_cast(a[static_cast(i * k + p)]) * + static_cast(w[static_cast(jj * k + p)]); + const double got = + rocm_out[static_cast(i * n + jj)]; + const double cpu = cpu_out[static_cast(i * n + jj)]; + num_ref += (got - ref) * (got - ref); + den_ref += ref * ref; + num_cpu += (got - cpu) * (got - cpu); + den_cpu += cpu * cpu; + REQUIRE(std::isfinite(got)); + } + } + const double nmse_ref = den_ref > 0 ? num_ref / den_ref : num_ref; + const double nmse_cpu = den_cpu > 0 ? num_cpu / den_cpu : num_cpu; + CAPTURE(nmse_ref); + CAPTURE(nmse_cpu); + const double ref_ceiling = + c.nmse_ref_max > 0 ? c.nmse_ref_max : kMaxNmseErr; + CHECK(nmse_ref <= ref_ceiling); // quantization error vs f64 dequant + CHECK(nmse_cpu <= kMaxNmseVsCpu); // matches the CPU oracle (int core exact) + } + } + } + gpu.DestroyQueue(gq); +} + +TEST_CASE("ROCm keep-quant registers the native kROCM providers") { + // The registration flips the GGUF loader's keep-quant default ON on a ROCm + // device (GgufQuantComputeAvailable -> OpRegistered(kMatmulBTQuant,kROCM)). + // Present only in a HIP build. + if (!vt::rocm::DeviceAvailable()) return; + CHECK(vt::OpRegistered(vt::OpId::kMatmulBTQuant, DeviceType::kROCM)); + CHECK(vt::OpRegistered(vt::OpId::kMatmulBTQuantGrouped, DeviceType::kROCM)); +} + +TEST_CASE( + "ROCm grouped keep-quant GEMM == CPU grouped golden and it WRITES the " + "output") { + if (!vt::rocm::DeviceAvailable()) return; + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + Queue cq{Cpu(), nullptr}; + + // All ten encodings, decode + prefill shapes, broadcast and per-row arms — + // the same matrix the CUDA grouped gate runs, over a POISONED output buffer. + struct GroupedShape { + int64_t P; + int64_t n; + int64_t E; + bool bcast; + }; + const GroupedShape kGroupedShapes[] = { + {6, 3, 4, false}, {32, 7, 8, false}, {16, 5, 2, true}}; + int64_t combos = 0; + for (const WeightCase& c : kCases) { + const int64_t k = 8 * c.block_elems; + for (const GroupedShape& g : kGroupedShapes) { + CAPTURE(std::string(c.name)); + CAPTURE(g.P); + CAPTURE(g.n); + CAPTURE(g.E); + CAPTURE(g.bcast); + const int64_t arows = g.bcast ? 1 : g.P; + std::vector wq = + RandomBlocks(c, g.E * g.n * (k / c.block_elems), 0x5EEDU); + std::vector af(static_cast(arows * k)); + GenerateData(1.0F, af.size(), af.data()); + std::vector ids(g.P); + for (int64_t p = 0; p < g.P; ++p) ids[static_cast(p)] = p % g.E; + const size_t outn = static_cast(g.P * g.n); + + // --- CPU golden (the landed grouped keep-quant kernel over host tensors) + std::vector cpu_out(outn, 1337.0F); + { + Tensor at = + Tensor::Contiguous(af.data(), DType::kF32, Cpu(), {arows, k}); + Tensor wt = + Tensor::Contiguous(wq.data(), DType::kF32, Cpu(), {g.E * g.n, k}); + wt.dtype = c.dtype; + Tensor et = + Tensor::Contiguous(ids.data(), DType::kI32, Cpu(), {g.P}); + Tensor ot = + Tensor::Contiguous(cpu_out.data(), DType::kF32, Cpu(), {g.P, g.n}); + vt::MatmulBTQuantGrouped(cq, ot, at, wt, et); + } + + // --- ROCm path over a POISONED output buffer ------------------------- + void* d_a = gpu.Alloc(af.size() * sizeof(float)); + void* d_w = gpu.Alloc(wq.size()); + void* d_e = gpu.Alloc(ids.size() * sizeof(int32_t)); + void* d_o = gpu.Alloc(outn * sizeof(float)); + std::vector poison(outn, 1337.0F); + gpu.Copy(gq, d_a, af.data(), af.size() * sizeof(float)); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + gpu.Copy(gq, d_e, ids.data(), ids.size() * sizeof(int32_t)); + gpu.Copy(gq, d_o, poison.data(), poison.size() * sizeof(float)); + gpu.Synchronize(gq); + Tensor at = DevTensor(d_a, DType::kF32, {arows, k}); + Tensor wt = DevTensor(d_w, c.dtype, {g.E * g.n, k}); + Tensor et = DevTensor(d_e, DType::kI32, {g.P}); + Tensor ot = DevTensor(d_o, DType::kF32, {g.P, g.n}); + vt::MatmulBTQuantGrouped(gq, ot, at, wt, et); + std::vector got(outn, 0.0F); + gpu.Copy(gq, got.data(), d_o, got.size() * sizeof(float)); + gpu.Synchronize(gq); + gpu.Free(d_a); + gpu.Free(d_w); + gpu.Free(d_e); + gpu.Free(d_o); + + int poisoned = 0; + int nonfinite = 0; + double num = 0, den = 0; + for (size_t i = 0; i < got.size(); ++i) { + if (got[i] == 1337.0F) ++poisoned; + if (!std::isfinite(got[i])) ++nonfinite; + num += (got[i] - cpu_out[i]) * (got[i] - cpu_out[i]); + den += cpu_out[i] * cpu_out[i]; + } + const double nmse = den > 0 ? num / den : num; + CAPTURE(nmse); + CHECK(poisoned == 0); // a dispatch that launches nothing lands HERE + CHECK(nonfinite == 0); + CHECK(nmse <= kMaxNmseVsCpu); + ++combos; + } + } + // doctest prints "SUCCESS!" for a loop that never ran. Say how many it ran. + CAPTURE(combos); + CHECK(combos == + static_cast(std::size(kCases) * std::size(kGroupedShapes))); + CHECK(combos > 0); + gpu.DestroyQueue(gq); +} + +TEST_CASE("ROCm K-quant decode arm (VT_GEMV_MMVQ=1) is BYTE-EXACT vs the default arm and within the oracle NMSE band") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + Queue cq{Cpu(), nullptr}; + + // m=1 (the decode shape the arm serves), Q4_K/Q5_K/Q6_K, nsb edges + // (nsb=1 -> one partial pass; nsb=3 -> ragged tail pass) and odd-but-valid + // N (warp-guard edge). + for (const WeightCase& c : kKQuantCases) { + for (int64_t nsb : {int64_t{1}, int64_t{3}, int64_t{10}}) { + const int64_t k = nsb * c.block_elems; + for (int64_t n : {int64_t{1}, int64_t{7}, int64_t{129}}) { + for (uint32_t seed : {0x5EEDU, 0xA11CEU}) { + CAPTURE(c.name); + CAPTURE(k); + CAPTURE(n); + CAPTURE(seed); + + std::vector wq = RandomBlocks(c, n * nsb, seed); + // Engine-realistic dtypes too: the model runs these projections with + // bf16 activations and bf16 outputs; f32-only tests were the blind + // spot that let the first fused build pass ops while the engine + // degraded. Activation storage is generated in `adt`. + for (DType adt : {DType::kF32, DType::kBF16, DType::kF16}) { + for (DType odt : {DType::kF32, DType::kBF16}) { + CAPTURE(adt); + CAPTURE(odt); + std::vector af(static_cast(k)); + GenerateData(static_cast(seed) + 0.5F * static_cast(int(adt)), + af.size(), af.data()); + std::vector abuf(af.size() * + (adt == DType::kF32 ? 4 : 2)); + for (size_t i2 = 0; i2 < af.size(); ++i2) { + if (adt == DType::kF32) + std::memcpy(abuf.data() + 4 * i2, &af[i2], 4); + else if (adt == DType::kBF16) { + const uint16_t h = vt::F32ToBF16(af[i2]); + std::memcpy(abuf.data() + 2 * i2, &h, 2); + } else { + const uint16_t h = vt::F32ToF16(af[i2]); + std::memcpy(abuf.data() + 2 * i2, &h, 2); + } + } + + // --- CPU oracle (host tensors, generic nrc==1 tier at m==1) ------- + std::vector cpu_out(static_cast(n), 0.0F); + { + Tensor at = Tensor::Contiguous(abuf.data(), adt, Cpu(), {1, k}); + Tensor bt = + Tensor::Contiguous(wq.data(), DType::kF32, Cpu(), {n, k}); + bt.dtype = c.dtype; + Tensor ot = + Tensor::Contiguous(cpu_out.data(), DType::kF32, Cpu(), {1, n}); + vt::MatmulBTQuant(cq, ot, at, bt); + } + + // --- ROCm path: BOTH arms at this shape; ON must equal OFF + // byte-for-byte (raw output buffer), and ON stays within the + // 1e-6 NMSE band vs the CPU oracle (the default arm's band) ---- + const size_t oesz = odt == DType::kF32 ? 4 : 2; + void* d_a = gpu.Alloc(abuf.size()); + void* d_w = gpu.Alloc(wq.size()); + gpu.Copy(gq, d_a, abuf.data(), abuf.size()); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + std::vector> arm_out(2); + std::vector> arm_raw(2); + for (int arm = 0; arm < 2; ++arm) { + void* d_o = gpu.Alloc(oesz * static_cast(n)); + { + EnvGuard on(arm == 1); + Tensor at = DevTensor(d_a, adt, {1, k}); + Tensor bt = DevTensor(d_w, c.dtype, {n, k}); + Tensor ot = DevTensor(d_o, odt, {1, n}); + vt::MatmulBTQuant(gq, ot, at, bt); + arm_raw[arm].resize(oesz * static_cast(n)); + gpu.Copy(gq, arm_raw[arm].data(), d_o, arm_raw[arm].size()); + arm_out[arm].resize(static_cast(n), 0.0F); + for (size_t i2 = 0; i2 < arm_out[arm].size(); ++i2) + arm_out[arm][i2] = + odt == DType::kF32 + ? reinterpret_cast(arm_raw[arm].data())[i2] + : vt::BF16ToF32( + reinterpret_cast(arm_raw[arm].data())[i2]); + gpu.Synchronize(gq); + } + gpu.Free(d_o); + } + gpu.Free(d_a); + gpu.Free(d_w); + + // ON arm must be BYTE-IDENTICAL to the default kernel + CHECK(std::memcmp(arm_raw[0].data(), arm_raw[1].data(), + arm_raw[0].size()) == 0); + // CPU side mirrors the output dtype conversion exactly + std::vector cpu_ref(cpu_out.size()); + for (size_t i2 = 0; i2 < cpu_out.size(); ++i2) + cpu_ref[i2] = odt == DType::kF32 + ? cpu_out[i2] + : vt::BF16ToF32(vt::F32ToBF16(cpu_out[i2])); + const double nmse_on = Nmse(arm_out[1], cpu_ref); + CAPTURE(nmse_on); + CHECK(nmse_on <= kMaxNmseVsCpu); + } // odt + } // adt + } + } + } + } + gpu.DestroyQueue(gq); +} + +TEST_CASE("T4a repair: MULTI-M calls stay byte-exact ON-vs-OFF (the m-gate red)") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + // RED-first for the TRUE defect-1: the arm's non-fused branch originally + // gated ONLY the fused fold on m==1, so ENGINE PREFILL calls (observed + // m=39) took the GEMV kernel, which writes row 0 only -- rows 1..m-1 of + // the output were left UNWRITTEN while every op-level test (m==1) stayed + // green. This case runs m>1 batches and asserts the FULL m x n output is + // byte-identical between the arms. + struct MCase { DType wt; int64_t m, n, k; }; + const std::vector cases = { + {DType::kQ4_K, 3, 7, 2560}, + {DType::kQ4_K, 39, 18432, 2560}, // the engine's observed prefill shape + {DType::kQ6_K, 5, 129, 9216}, + {DType::kQ6_K, 2, 248320, 2560}, // lm_head-class with m=2 + }; + for (const MCase& mc : cases) { + const WeightCase* c = nullptr; + for (const WeightCase& wc : kKQuantCases) + if (wc.dtype == mc.wt) c = &wc; + const int64_t nsb = mc.k / c->block_elems; + CAPTURE(mc.m); + CAPTURE(mc.n); + CAPTURE(mc.k); + std::vector wq = RandomBlocks(*c, mc.n * nsb, 0x5EEDU); + const size_t aesz = 2; // bf16 activations, engine-realistic + std::vector abuf(static_cast(mc.m * mc.k)); + for (size_t i = 0; i < abuf.size(); ++i) + abuf[i] = static_cast((i * 2654435761u) >> 11); + void* d_w = gpu.Alloc(wq.size()); + void* d_a = gpu.Alloc(abuf.size() * aesz); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + gpu.Copy(gq, d_a, abuf.data(), abuf.size() * aesz); + constexpr size_t kOesz = 2; + std::vector> outs(2); + for (int arm = 0; arm < 2; ++arm) { + void* d_o = gpu.Alloc(kOesz * static_cast(mc.m * mc.n)); + { + // Canary-fill so any UNWRITTEN row is detected rather than + // coincidentally matching stale allocation contents. + std::vector canary(kOesz * static_cast(mc.m * mc.n), + arm == 1 ? 0xAB : 0xCD); + gpu.Copy(gq, d_o, canary.data(), canary.size()); + EnvGuard guard(arm == 1); + Tensor at = DevTensor(d_a, DType::kBF16, {mc.m, mc.k}); + Tensor bt = DevTensor(d_w, c->dtype, {mc.n, mc.k}); + Tensor ot = DevTensor(d_o, DType::kBF16, {mc.m, mc.n}); + vt::MatmulBTQuant(gq, ot, at, bt); + outs[arm].resize(kOesz * static_cast(mc.m * mc.n)); + gpu.Copy(gq, outs[arm].data(), d_o, outs[arm].size()); + gpu.Synchronize(gq); + } + gpu.Free(d_o); + } + gpu.Free(d_w); + gpu.Free(d_a); + size_t first_bad = outs[0].size(); + for (size_t i = 0; i < outs[0].size(); ++i) + if (outs[0][i] != outs[1][i]) { first_bad = i; break; } + CAPTURE(first_bad); + CHECK(outs[0] == outs[1]); + } + gpu.DestroyQueue(gq); +} + +TEST_CASE("T4a repair: ON-vs-OFF BYTE identity over the ENGINE shape set (incl. lm_head-sized N)") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + + // The REAL (dtype, N, K) set the engine serves, from the GGUF tensor + // manifest of the acceptance checkpoint (Qwen3.5-4B-Q4_K_M: ne0=K, + // ne1=N) plus the operator's ON-capture grids (grid = ceil(N/4) at + // kGemvWarps=4 -> grid 80/576/256/7760) and the contract-named + // lm_head-class probes. Defect-1 hypothesis under test: a 32-bit + // offset/index overflow once N*w_row_bytes grows past 2^31 (lm_head + // N x nsb x 210B ~= 0.5 GB at these shapes -- near the int32 edge). + struct ShapeCase { + DType wt; + const char* name; + int64_t n, k; + bool all_act_dtypes; // giants run bf16-only (engine-realistic) to bound suite time + }; + const std::vector shapes = { + {DType::kQ6_K, "q6_K grid=80", 320, 2560, true}, + {DType::kQ4_K, "q4_K grid=80", 320, 2560, true}, + {DType::kQ4_K, "q4_K grid=576", 2304, 2560, true}, + {DType::kQ6_K, "q6_K blk.out", 1024, 2560, true}, + {DType::kQ4_K, "q4_K ffn-out", 2560, 4096, true}, + {DType::kQ5_K, "q5_K ffn-out", 2560, 4096, true}, + {DType::kQ5_K, "q5_K gate_up", 8192, 2560, true}, + {DType::kQ4_K, "q4_K gate_up", 8192, 2560, true}, + {DType::kQ4_K, "q4_K down", 2560, 9216, true}, + {DType::kQ6_K, "q6_K down", 2560, 9216, true}, + {DType::kQ6_K, "q6_K grid=7760 (operator lm_head-class)", 31040, 4096, false}, + {DType::kQ6_K, "q6_K lm_head-class N=151936 (contract-named)", 151936, 4096, false}, + {DType::kQ6_K, "q6_K lm_head REAL N=248320", 248320, 2560, false}, + // Exact tuples observed from the ENGINE dispatch trace (bf16 x bf16): + {DType::kQ4_K, "ENGINE q4_K n=18432 k=2560", 18432, 2560, false}, + {DType::kQ4_K, "ENGINE q4_K n=1024 k=2560", 1024, 2560, true}, + {DType::kQ4_K, "ENGINE q4_K n=2560 k=4096", 2560, 4096, true}, + {DType::kQ4_K, "ENGINE q4_K n=8192 k=2560", 8192, 2560, true}, + }; + + for (const ShapeCase& sc : shapes) { + const WeightCase* c = nullptr; + for (const WeightCase& wc : kKQuantCases) + if (wc.dtype == sc.wt) c = &wc; + const int64_t nsb = sc.k / c->block_elems; + CHECK(sc.k % c->block_elems == 0); + CAPTURE(std::string(sc.name)); + CAPTURE(sc.n); + CAPTURE(sc.k); + + std::vector wq = RandomBlocks(*c, sc.n * nsb, 0x5EEDU); + const size_t wbytes = wq.size(); + void* d_w = gpu.Alloc(wbytes); + gpu.Copy(gq, d_w, wq.data(), wbytes); + + std::vector adts{DType::kBF16}; + if (sc.all_act_dtypes) adts = {DType::kF32, DType::kBF16, DType::kF16}; + for (DType adt : adts) { + CAPTURE(adt); + const size_t aesz = adt == DType::kF32 ? 4 : 2; + // One fixed activation row, magnitudes the engine actually sees. + std::vector af(static_cast(sc.k)); + GenerateData(3.0F, af.size(), af.data()); + std::vector abuf(af.size() * aesz); + for (size_t i = 0; i < af.size(); ++i) { + if (adt == DType::kF32) + std::memcpy(abuf.data() + 4 * i, &af[i], 4); + else if (adt == DType::kBF16) { + const uint16_t h = vt::F32ToBF16(af[i]); + std::memcpy(abuf.data() + 2 * i, &h, 2); + } else { + const uint16_t h = vt::F32ToF16(af[i]); + std::memcpy(abuf.data() + 2 * i, &h, 2); + } + } + void* d_a = gpu.Alloc(abuf.size()); + gpu.Copy(gq, d_a, abuf.data(), abuf.size()); + + // Run BOTH arms at the SAME output dtype (bf16, engine-realistic) + // and compare RAW output bytes. + constexpr size_t kOesz = 2; // bf16 + std::vector> outs(2); + for (int arm = 0; arm < 2; ++arm) { + void* d_o = gpu.Alloc(kOesz * static_cast(sc.n)); + { + EnvGuard guard(arm == 1); + Tensor at = DevTensor(d_a, adt, {1, sc.k}); + Tensor bt = DevTensor(d_w, c->dtype, {sc.n, sc.k}); + Tensor ot = DevTensor(d_o, DType::kBF16, {1, sc.n}); + vt::MatmulBTQuant(gq, ot, at, bt); + outs[arm].resize(kOesz * static_cast(sc.n)); + gpu.Copy(gq, outs[arm].data(), d_o, outs[arm].size()); + gpu.Synchronize(gq); + } + gpu.Free(d_o); + } + // Byte identity: locate and report the FIRST divergence for triage. + size_t first_bad = outs[0].size(); + for (size_t i = 0; i < outs[0].size(); ++i) + if (outs[0][i] != outs[1][i]) { first_bad = i; break; } + CAPTURE(first_bad); + CHECK(outs[0] == outs[1]); + gpu.Free(d_a); + } + gpu.Free(d_w); + } + gpu.DestroyQueue(gq); +} + +TEST_CASE("ROCm K-quant DEFAULT arm (env unset) stays within 1e-6 NMSE vs CPU") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + Queue cq{Cpu(), nullptr}; + + // Default-OFF inertness probe: with no VT_GEMV_MMVQ in the environment the + // baseline warp-reduction kernel must be untouched by the T4a change. The + // baseline's shfl tree reassociates the float sum, so this holds it to the + // SAME 1e-6 NMSE-vs-CPU band as the CUDA sibling gate — not bit-exactness. + const WeightCase& c = kKQuantCases[0]; // q4_K + const int64_t nsb = 10, k = nsb * c.block_elems, n = 7; + std::vector wq = RandomBlocks(c, n * nsb, 0x5EEDU); + std::vector a(static_cast(k)); + GenerateData(1.0F, a.size(), a.data()); + + std::vector cpu_out(static_cast(n), 0.0F); + { + Tensor at = Tensor::Contiguous(a.data(), DType::kF32, Cpu(), {1, k}); + Tensor bt = Tensor::Contiguous(wq.data(), DType::kF32, Cpu(), {n, k}); + bt.dtype = c.dtype; + Tensor ot = Tensor::Contiguous(cpu_out.data(), DType::kF32, Cpu(), {1, n}); + vt::MatmulBTQuant(cq, ot, at, bt); + } + + void* d_a = gpu.Alloc(a.size() * sizeof(float)); + void* d_w = gpu.Alloc(wq.size()); + void* d_o = gpu.Alloc(sizeof(float) * static_cast(n)); + gpu.Copy(gq, d_a, a.data(), a.size() * sizeof(float)); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + std::vector rocm_out(static_cast(n), 0.0F); + { + EnvGuard off(false); // explicitly "0": the arm must NOT engage + Tensor at = DevTensor(d_a, DType::kF32, {1, k}); + Tensor bt = DevTensor(d_w, c.dtype, {n, k}); + Tensor ot = DevTensor(d_o, DType::kF32, {1, n}); + vt::MatmulBTQuant(gq, ot, at, bt); + gpu.Copy(gq, rocm_out.data(), d_o, rocm_out.size() * sizeof(float)); + gpu.Synchronize(gq); + } + gpu.Free(d_a); + gpu.Free(d_w); + gpu.Free(d_o); + + const double nmse = Nmse(rocm_out, cpu_out); + CAPTURE(nmse); + CHECK(nmse <= kMaxNmseVsCpu); + gpu.DestroyQueue(gq); +} + +TEST_CASE("Fused-prologue Q8_K quantization is BYTE-IDENTICAL to the standalone quantizer") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + // nsb=10 covers this model's decode K; inputs: pseudo-random rows plus an + // ADVERSARIAL tied-amax row (+max first, equal-magnitude negative later, so + // the amax FIRST-occurrence tie-break is what decides mx's sign) and an + // all-zero row. + const int64_t k = 10 * 256; + std::mt19937 rng(0xB00B5U); + std::vector> rows; + for (int r = 0; r < 4; ++r) { + std::vector a(static_cast(k)); + for (float& v : a) v = static_cast(static_cast(rng() % 2001) - 1000) / 500.0F; + rows.push_back(std::move(a)); + } + { + std::vector a(static_cast(k), 0.0F); + a[0] = 3.5F; + a[17] = -3.5F; // exact fabs tie; FIRST occurrence (index 0) must win + a[291] = -3.5F; // another tie, still after index 0 + rows.push_back(std::move(a)); + } + rows.push_back(std::vector(static_cast(k), 0.0F)); + + for (size_t r = 0; r < rows.size(); ++r) { + CAPTURE(r); + const std::vector& a = rows[r]; + void* d_a = gpu.Alloc(a.size() * sizeof(float)); + void* d_sa = gpu.Alloc(10 * 292); // sizeof(BlockQ8_K), pinned by static_assert + void* d_sb = gpu.Alloc(10 * 292); + gpu.Copy(gq, d_a, a.data(), a.size() * sizeof(float)); + Tensor at = DevTensor(d_a, DType::kF32, {1, k}); + vt::rocm::MmvqQuantScratchForTesting(gq, d_sa, at, false); + vt::rocm::MmvqQuantScratchForTesting(gq, d_sb, at, true); + std::vector sa(10 * 292), sb(10 * 292); + gpu.Copy(gq, sa.data(), d_sa, sa.size()); + gpu.Copy(gq, sb.data(), d_sb, sb.size()); + gpu.Synchronize(gq); + gpu.Free(d_a); gpu.Free(d_sa); gpu.Free(d_sb); + CHECK(std::memcmp(sa.data(), sb.data(), sa.size()) == 0); + } + gpu.DestroyQueue(gq); +} + +TEST_CASE("T4a repair: per-grid OFF-vs-ON timing at the operator's captured grids") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + // Evidence-only case (no correctness assert): median us/call per arm at + // the grids the operator's rocprofv3 parse captured (grid = ceil(N/4)), + // plus the real lm_head shape. bf16 act/out, engine-realistic. + struct BenchShape { DType wt; const char* name; int64_t n, k; int reps; }; + const std::vector shapes = { + {DType::kQ6_K, "grid=80 Li2 (320x2560)", 320, 2560, 30}, + {DType::kQ4_K, "grid=80 Li0 (320x2560)", 320, 2560, 30}, + {DType::kQ4_K, "grid=576 Li0 (2304x2560)", 2304, 2560, 30}, + {DType::kQ6_K, "grid=7760 Li2 (31040x4096)", 31040, 4096, 12}, + {DType::kQ6_K, "lm_head real (248320x2560)", 248320, 2560, 8}, + }; + for (const BenchShape& sc : shapes) { + const WeightCase* c = nullptr; + for (const WeightCase& wc : kKQuantCases) + if (wc.dtype == sc.wt) c = &wc; + const int64_t nsb = sc.k / c->block_elems; + std::vector wq = RandomBlocks(*c, sc.n * nsb, 0x5EEDU); + std::vector af(static_cast(sc.k)); + GenerateData(3.0F, af.size(), af.data()); + std::vector abuf(af.size()); + for (size_t i = 0; i < af.size(); ++i) + abuf[i] = vt::F32ToBF16(af[i]); + void* d_w = gpu.Alloc(wq.size()); + void* d_a = gpu.Alloc(abuf.size() * 2); + void* d_o = gpu.Alloc(2 * static_cast(sc.n)); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + gpu.Copy(gq, d_a, abuf.data(), abuf.size() * 2); + MESSAGE(sc.name); + double med[2] = {0, 0}; + for (int arm = 0; arm < 2; ++arm) { + EnvGuard guard(arm == 1); + Tensor at = DevTensor(d_a, DType::kBF16, {1, sc.k}); + Tensor bt = DevTensor(d_w, c->dtype, {sc.n, sc.k}); + Tensor ot = DevTensor(d_o, DType::kBF16, {1, sc.n}); + for (int w = 0; w < 3; ++w) { // warmup + vt::MatmulBTQuant(gq, ot, at, bt); + gpu.Synchronize(gq); + } + std::vector t; + for (int r = 0; r < sc.reps; ++r) { + const auto t0 = std::chrono::steady_clock::now(); + vt::MatmulBTQuant(gq, ot, at, bt); + gpu.Synchronize(gq); + const auto t1 = std::chrono::steady_clock::now(); + t.push_back(std::chrono::duration(t1 - t0).count()); + } + std::sort(t.begin(), t.end()); + med[arm] = t[t.size() / 2]; + } + char buf[128]; + std::snprintf(buf, sizeof(buf), + " OFF %9.1f us/call | ON %9.1f us/call | ratio ON/OFF %.2fx", + med[0], med[1], med[1] / med[0]); + MESSAGE(buf); + gpu.Free(d_w); + gpu.Free(d_a); + gpu.Free(d_o); + } + gpu.DestroyQueue(gq); +} + +// --------------------------------------------------------------------------- +// T4a REPAIR ROUND 2 (reviewer findings F1/F2). The round-1 gate could not +// witness ROUTING: EnvGuard(false) writes "0" (never a true unset), and since +// ON==OFF are bit-equal by design, every output comparison is blind to which +// dispatch branch ran. These two cases pin routing itself via the host-side +// dispatch counters. + +// F1: with VT_GEMV_MMVQ TRULY ABSENT (unsetenv, not "0") the call must take +// the BASELINE branch; with VT_GEMV_MMVQ=1 it must NOT. Catches an inverted +// getenv default (mutation M3) that outputs cannot see. +TEST_CASE("T4a repair-2 F1: ROUTING WITNESS -- env truly unset routes to BASELINE; ON routes to the GEMV arm") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + const WeightCase& c = kKQuantCases[0]; // q4_K + const int64_t nsb = 10, k = nsb * c.block_elems, n = 7; + std::vector wq = RandomBlocks(c, n * nsb, 0x5EEDU); + std::vector a(static_cast(k)); + GenerateData(1.5F, a.size(), a.data()); + + void* d_a = gpu.Alloc(a.size() * sizeof(float)); + void* d_w = gpu.Alloc(wq.size()); + void* d_o = gpu.Alloc(sizeof(float) * static_cast(n)); + gpu.Copy(gq, d_a, a.data(), a.size() * sizeof(float)); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + + auto run_once = [&] { + Tensor at = DevTensor(d_a, DType::kF32, {1, k}); + Tensor bt = DevTensor(d_w, c.dtype, {n, k}); + Tensor ot = DevTensor(d_o, DType::kF32, {1, n}); + vt::MatmulBTQuant(gq, ot, at, bt); + gpu.Synchronize(gq); + }; + + // TRUE unset: the flag string must be absent from the environment -- NOT + // EnvGuard(false), which sets "0". Default-OFF inertness means the + // BASELINE counter advances and no GEMV counter moves. + ::unsetenv("VT_GEMV_MMVQ"); + vt::rocm::MmvqResetRouteCountsForTesting(); + run_once(); + const auto off_counts = vt::rocm::MmvqRouteCountsForTesting(); + CHECK(off_counts.baseline == 1); + CHECK(off_counts.gemv_mmvq == 0); + CHECK(off_counts.gemv_fused == 0); + + // Paired ON case: exactly the reverse. n=7 <= kMmvqFoldMaxRows, so the + // arm engages via its FUSED sub-branch; either way the baseline counter + // must not move. + { + EnvGuard on(true); + vt::rocm::MmvqResetRouteCountsForTesting(); + run_once(); + const auto on_counts = vt::rocm::MmvqRouteCountsForTesting(); + CHECK(on_counts.baseline == 0); + CHECK(on_counts.gemv_fused == 1); + CHECK(on_counts.gemv_mmvq == 0); + } + ::unsetenv("VT_GEMV_MMVQ"); + gpu.Free(d_a); + gpu.Free(d_w); + gpu.Free(d_o); + gpu.DestroyQueue(gq); +} + +// F2: fold-crossover WITNESS. With the arm ON, n=256 (<= kMmvqFoldMaxRows) +// must dispatch through the FUSED sub-branch and n=2304 (> 512, within the +// reviewer's mutated range (512,4096]) must dispatch through the NON-FUSED +// GEMV branch. Catches a kMmvqFoldMaxRows drift (mutation M4: 512 -> 4096) +// that flips measured per-call ratios while staying output-green. +TEST_CASE("T4a repair-2 F2: FOLD-CROSSOVER WITNESS -- fused sub-branch only at n <= kMmvqFoldMaxRows") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + const WeightCase& c = kKQuantCases[0]; // q4_K + const int64_t nsb = 10, k = nsb * c.block_elems; + + struct FoldShape { const char* name; int64_t n; long long want_fused, want_gemv, want_baseline; }; + const FoldShape shapes[] = { + {"n=256 (fold expected)", 256, 1, 0, 0}, + {"n=2304 (fold NOT expected)", 2304, 0, 1, 0}, + }; + for (const FoldShape& sc : shapes) { + CAPTURE(sc.name); + std::vector wq = RandomBlocks(c, sc.n * nsb, 0x5EEDU); + std::vector a(static_cast(k)); + GenerateData(2.5F, a.size(), a.data()); + void* d_a = gpu.Alloc(a.size() * sizeof(float)); + void* d_w = gpu.Alloc(wq.size()); + void* d_o = gpu.Alloc(sizeof(float) * static_cast(sc.n)); + gpu.Copy(gq, d_a, a.data(), a.size() * sizeof(float)); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + { + EnvGuard on(true); + vt::rocm::MmvqResetRouteCountsForTesting(); + Tensor at = DevTensor(d_a, DType::kF32, {1, k}); + Tensor bt = DevTensor(d_w, c.dtype, {sc.n, k}); + Tensor ot = DevTensor(d_o, DType::kF32, {1, sc.n}); + vt::MatmulBTQuant(gq, ot, at, bt); + gpu.Synchronize(gq); + const auto counts = vt::rocm::MmvqRouteCountsForTesting(); + CHECK(counts.gemv_fused == sc.want_fused); + CHECK(counts.gemv_mmvq == sc.want_gemv); + CHECK(counts.baseline == sc.want_baseline); + } + ::unsetenv("VT_GEMV_MMVQ"); + gpu.Free(d_a); + gpu.Free(d_w); + gpu.Free(d_o); + } + gpu.DestroyQueue(gq); +} + +// F3 (lever B1, GFX1100-TG200): the fold crossover becomes RUNTIME-TUNABLE +// via VT_GEMV_MMVQ_FOLD_MAX (integer rows; default = kMmvqFoldMaxRowsDefault +// = 512; invalid/empty = default). The suite constants above keep pinning +// DEFAULT behavior; THIS case asserts the env actually moves ROUTING via the +// same host-side dispatch counters: +// - unset : n=256 folds, n=2304 does NOT (default pinned) +// - "4096" : n=2304 FOLDS (knob widens the gate) [RED pre-knob: env inert] +// - "128" : n=256 does NOT fold (knob narrows the gate) [RED pre-knob: env inert] +// - "256" : n=256 still folds (boundary is INCLUSIVE <=) +// - garbage: behaves exactly like unset (invalid falls back to default) +// RED-first contract: before the knob exists VT_GEMV_MMVQ_FOLD_MAX is +// inert, so the "4096" and "128" legs fail while routing stays at defaults. +namespace { +struct FoldMaxGuard { + explicit FoldMaxGuard(const char* v) { + if (v != nullptr) ::setenv("VT_GEMV_MMVQ_FOLD_MAX", v, 1); + else ::unsetenv("VT_GEMV_MMVQ_FOLD_MAX"); + } + ~FoldMaxGuard() { ::unsetenv("VT_GEMV_MMVQ_FOLD_MAX"); } +}; +} // namespace + +TEST_CASE("T4a lever-B1 F3: FOLD-MAX KNOB WITNESS -- VT_GEMV_MMVQ_FOLD_MAX moves routing at runtime; invalid values fall back to the default") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + const WeightCase& c = kKQuantCases[0]; // q4_K + const int64_t nsb = 10, k = nsb * c.block_elems; + + struct Leg { const char* name; const char* fold_max; int64_t n; + long long want_fused, want_gemv, want_baseline; }; + const Leg legs[] = { + {"unset n=256 (default pins fold)", nullptr, 256, 1, 0, 0}, + {"unset n=2304 (default pins non-fused)", nullptr, 2304, 0, 1, 0}, + {"4096 n=2304 (knob WIDENS -> fold)", "4096", 2304, 1, 0, 0}, + {"128 n=256 (knob NARROWS -> gemv)", "128", 256, 0, 1, 0}, + {"256 n=256 (boundary is inclusive)", "256", 256, 1, 0, 0}, + {"garbage n=256 (invalid -> default fold)", "not-a-number", 256, 1, 0, 0}, + {"garbage n=2304 (invalid -> default gemv)", "not-a-number", 2304, 0, 1, 0}, + }; + for (const Leg& sc : legs) { + CAPTURE(sc.name); + std::vector wq = RandomBlocks(c, sc.n * nsb, 0x5EEDU); + std::vector a(static_cast(k)); + GenerateData(2.5F, a.size(), a.data()); + void* d_a = gpu.Alloc(a.size() * sizeof(float)); + void* d_w = gpu.Alloc(wq.size()); + void* d_o = gpu.Alloc(sizeof(float) * static_cast(sc.n)); + gpu.Copy(gq, d_a, a.data(), a.size() * sizeof(float)); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + { + EnvGuard on(true); + FoldMaxGuard fm(sc.fold_max); + vt::rocm::MmvqResetRouteCountsForTesting(); + Tensor at = DevTensor(d_a, DType::kF32, {1, k}); + Tensor bt = DevTensor(d_w, c.dtype, {sc.n, k}); + Tensor ot = DevTensor(d_o, DType::kF32, {1, sc.n}); + vt::MatmulBTQuant(gq, ot, at, bt); + gpu.Synchronize(gq); + const auto counts = vt::rocm::MmvqRouteCountsForTesting(); + CHECK(counts.gemv_fused == sc.want_fused); + CHECK(counts.gemv_mmvq == sc.want_gemv); + CHECK(counts.baseline == sc.want_baseline); + } + ::unsetenv("VT_GEMV_MMVQ_FOLD_MAX"); + ::unsetenv("VT_GEMV_MMVQ"); + gpu.Free(d_a); + gpu.Free(d_w); + gpu.Free(d_o); + } + gpu.DestroyQueue(gq); +} + +// --- Lever C (GFX1100-TG200-NORMQ): producer-fused Q8_K norm epilogue ------- +// +// RED-FIRST contract: before the epilogue exists VT_NORM_QUANT_FUSED=1 is +// inert, so the ON-leg witness expectations (producers>=1, standalone skipped) +// FAIL while the OFF leg trivially holds; the scratch byte-equality case also +// fails because NormQuantLastScratchForTesting() has no producer to observe. +namespace { + +struct EnvNormQuantGuard { + explicit EnvNormQuantGuard(bool on) { + ::setenv("VT_NORM_QUANT_FUSED", on ? "1" : "0", 1); + } + ~EnvNormQuantGuard() { ::unsetenv("VT_NORM_QUANT_FUSED"); } +}; + +std::vector RunNormQuantChain(Backend& gpu, Queue& gq, + void* d_x, void* d_nw, void* d_w, + void* d_o, int64_t k, int64_t n) { + std::vector out_raw(sizeof(uint16_t) * static_cast(n)); + Tensor xt = DevTensor(d_x, DType::kBF16, {1, k}); + Tensor wt = DevTensor(d_nw, DType::kBF16, {k}); + void* d_norm = gpu.Alloc(sizeof(uint16_t) * static_cast(k)); + Tensor nout = DevTensor(d_norm, DType::kBF16, {1, k}); + vt::RmsNorm(gq, nout, xt, wt, vt::RmsNormArgs{1e-6f, false}); + Tensor bt = DevTensor(d_w, DType::kQ4_K, {n, k}); + Tensor oo = DevTensor(d_o, DType::kBF16, {1, n}); + vt::MatmulBTQuant(gq, oo, nout, bt); + gpu.Copy(gq, out_raw.data(), d_o, out_raw.size()); + gpu.Synchronize(gq); + gpu.Free(d_norm); + return out_raw; +} + +} // namespace + +TEST_CASE("Lever C red: VT_NORM_QUANT_FUSED=1 routes norm-produced activations through the fused epilogue (counter witnesses + byte identity)") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + const int64_t k = 10 * 256, n = 64; + // weight blocks for a Q4_K [n,k] matvec + std::vector wq = RandomBlocks(kKQuantCases[0], n * 10, 0xC0FFEEU); + // bf16 activation row (the engine's dtype on this path) + std::vector af(static_cast(k)); + GenerateData(0.75F, af.size(), af.data()); + std::vector abf(af.size()); + for (size_t i = 0; i < af.size(); ++i) abf[i] = vt::F32ToBF16(af[i]); + // bf16 norm weight + std::vector nw(static_cast(k)); + std::mt19937 rng(7U); + for (uint16_t& v : nw) v = vt::F32ToBF16(0.5F + static_cast(rng() % 100) / 200.0F); + + void* d_a = gpu.Alloc(abf.size() * 2); + void* d_nw = gpu.Alloc(nw.size() * 2); + void* d_w = gpu.Alloc(wq.size()); + void* d_o = gpu.Alloc(2 * static_cast(n)); + gpu.Copy(gq, d_a, abf.data(), abf.size() * 2); + gpu.Copy(gq, d_nw, nw.data(), nw.size() * 2); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + + // OFF leg: flag absent -> no producer epilogue, standalone quant runs. + std::vector off_raw; + { + vt::rocm::NormQuantResetForTesting(); + off_raw = RunNormQuantChain(gpu, gq, d_a, d_nw, d_w, d_o, k, n); + const auto c = vt::rocm::NormQuantCountsForTesting(); + CHECK(c.producers == 0); + CHECK(c.consumers_fused == 0); + CHECK(c.consumers_standalone == 1); + } + // ON leg: epilogue fires, the consumer SKIPS the standalone quant, and a + // second consumer of the SAME activation (the attn q/k/v pattern: three + // matvecs re-quantizing one normalized row) skips too. Outputs must stay + // byte-identical to the OFF arm. + { + EnvNormQuantGuard on(true); + vt::rocm::NormQuantResetForTesting(); + // run the chain twice manually to keep the same normalized buffer alive + // across two consumers + Tensor xt = DevTensor(d_a, DType::kBF16, {1, k}); + Tensor wt = DevTensor(d_nw, DType::kBF16, {k}); + void* d_norm = gpu.Alloc(sizeof(uint16_t) * static_cast(k)); + Tensor nout = DevTensor(d_norm, DType::kBF16, {1, k}); + vt::RmsNorm(gq, nout, xt, wt, vt::RmsNormArgs{1e-6f, false}); + Tensor bt = DevTensor(d_w, DType::kQ4_K, {n, k}); + std::vector on_raw(sizeof(uint16_t) * static_cast(n)); + for (int consumer = 0; consumer < 2; ++consumer) { + Tensor oo = DevTensor(d_o, DType::kBF16, {1, n}); + vt::MatmulBTQuant(gq, oo, nout, bt); + gpu.Copy(gq, on_raw.data(), d_o, on_raw.size()); + gpu.Synchronize(gq); + } + gpu.Free(d_norm); + const auto c = vt::rocm::NormQuantCountsForTesting(); + CHECK(c.producers == 1); + CHECK(c.consumers_fused == 2); + CHECK(c.consumers_standalone == 0); + CHECK(std::memcmp(on_raw.data(), off_raw.data(), on_raw.size()) == 0); + } + gpu.Free(d_a); gpu.Free(d_nw); gpu.Free(d_w); gpu.Free(d_o); + gpu.DestroyQueue(gq); +} + +TEST_CASE("Lever C: fused norm-epilogue Q8_K scratch is BYTE-IDENTICAL to the standalone QuantizeQ8KK (random, tied-amax, zero rows; m=1 and m=3)") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + constexpr size_t kQ8KBytes = 292; // sizeof(BlockQ8_K), pinned by static_assert + for (int64_t nsb : {int64_t{1}, int64_t{3}, int64_t{10}}) { + const int64_t k = nsb * 256; + CAPTURE(k); + for (int64_t rows : {int64_t{1}, int64_t{3}}) { + CAPTURE(rows); + // row set: pseudo-random x(rows), an adversarial tied-amax row (fabs + // tie decided by FIRST occurrence -> index 0 wins; inverting the + // tie-break flips mx's sign and the whole block), an all-zero row. + std::mt19937 rng(0xB00B5U + static_cast(rows)); + std::vector> rowset; + // rows-1 pseudo-random rows, then the adversarial tied-amax row (fabs + // tie decided by FIRST occurrence -> index 0 wins; inverting the + // tie-break flips mx's sign and the whole block). For rows>=3 a final + // all-zero row rides along. + for (int r = 0; r < rows - 1; ++r) { + std::vector a(static_cast(k)); + for (float& v : a) v = static_cast(static_cast(rng() % 2001) - 1000) / 500.0F; + rowset.push_back(std::move(a)); + } + { + std::vector a(static_cast(k), 0.0F); + a[0] = 3.5F; + a[17] = -3.5F; + if (k > 300) a[291] = -3.5F; + rowset.push_back(std::move(a)); + } + if (rows >= 3) rowset.push_back(std::vector(static_cast(k), 0.0F)); + + const size_t abuf_bytes = rowset.size() * static_cast(k) * 2; + std::vector abf(rowset.size() * static_cast(k)); + std::vector nw(static_cast(k)); + for (size_t i = 0; i < nw.size(); ++i) nw[i] = vt::F32ToBF16(0.5F); + for (size_t r = 0; r < rowset.size(); ++r) + for (int64_t j = 0; j < k; ++j) abf[r * static_cast(k) + static_cast(j)] = vt::F32ToBF16(rowset[r][static_cast(j)]); + + void* d_a = gpu.Alloc(abuf_bytes); + void* d_nw = gpu.Alloc(nw.size() * 2); + gpu.Copy(gq, d_a, abf.data(), abuf_bytes); + gpu.Copy(gq, d_nw, nw.data(), nw.size() * 2); + + // The fused epilogue quantizes the NORM'S OUTPUT rows, so the reference + // is the standalone quantizer over those SAME output rows: run the + // producer-fused RmsNorm first, then hook the standalone QuantizeQ8KK + // on the produced out tensor (device dst, copied back after). + void* d_out = gpu.Alloc(abuf_bytes); + EnvNormQuantGuard on(true); + vt::rocm::NormQuantResetForTesting(); + Tensor xt = DevTensor(d_a, DType::kBF16, {static_cast(rowset.size()), k}); + Tensor wt = DevTensor(d_nw, DType::kBF16, {k}); + Tensor ot = DevTensor(d_out, DType::kBF16, {static_cast(rowset.size()), k}); + vt::RmsNorm(gq, ot, xt, wt, vt::RmsNormArgs{1e-6f, false}); + const void* scratch = vt::rocm::NormQuantLastScratchForTesting(); + REQUIRE(scratch != nullptr); + + void* d_ref = gpu.Alloc(rowset.size() * static_cast(nsb) * kQ8KBytes); + for (size_t r = 0; r < rowset.size(); ++r) { + Tensor rt = DevTensor(static_cast(d_out) + r * static_cast(k) * 2, DType::kBF16, {1, k}); + vt::rocm::MmvqQuantScratchForTesting(gq, static_cast(d_ref) + r * static_cast(nsb) * kQ8KBytes, rt, false); + } + + std::vector ref(rowset.size() * nsb * kQ8KBytes); + gpu.Copy(gq, ref.data(), d_ref, ref.size()); + std::vector got(rowset.size() * nsb * kQ8KBytes); + gpu.Copy(gq, got.data(), scratch, got.size()); + gpu.Synchronize(gq); + gpu.Free(d_ref); + CHECK(std::memcmp(got.data(), ref.data(), got.size()) == 0); + // HOST-ORACLE leg: vt::cpu::QuantizeRowQ8_K over the bf16-rounded norm + // outputs. The two GPU paths above share one device body, so a drift in + // that body moves BOTH identically -- this independent oracle is what + // actually pins the tie-break (lowest-index first occurrence) and the + // d-scale arithmetic down. + const auto from_float = vt::cpu::BlockFromFloat(DType::kQ8_K); + REQUIRE(from_float != nullptr); + std::vector out_host(rowset.size() * static_cast(k)); + gpu.Copy(gq, out_host.data(), d_out, out_host.size() * 2); + gpu.Synchronize(gq); + for (size_t r = 0; r < rowset.size(); ++r) { + std::vector xf(static_cast(k)); + for (int64_t j = 0; j < k; ++j) + xf[static_cast(j)] = + vt::BF16ToF32(out_host[r * static_cast(k) + static_cast(j)]); + std::vector want(nsb * kQ8KBytes); + from_float(xf.data(), want.data(), k); + CAPTURE(r); + CHECK(std::memcmp(got.data() + r * nsb * kQ8KBytes, want.data(), + nsb * kQ8KBytes) == 0); + } + gpu.Free(d_out); + gpu.Free(d_a); + gpu.Free(d_nw); + } + } + gpu.DestroyQueue(gq); +} + +TEST_CASE("Lever C: a non-matching K-quant consumer invalidates the producer token (stale-scratch guard)") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + const int64_t k = 10 * 256, n = 32, k2 = 3 * 256; + std::vector wq = RandomBlocks(kKQuantCases[0], n * 10, 0xD00DU); + std::vector wq2 = RandomBlocks(kKQuantCases[0], n * 3, 0xD01DU); + std::vector abf(static_cast(k)), a2bf(static_cast(k2)); + for (size_t i = 0; i < abf.size(); ++i) abf[i] = vt::F32ToBF16(0.1F * static_cast(i % 31)); + for (size_t i = 0; i < a2bf.size(); ++i) a2bf[i] = vt::F32ToBF16(0.2F * static_cast(i % 17)); + std::vector nw(static_cast(k)); + for (size_t i = 0; i < nw.size(); ++i) nw[i] = vt::F32ToBF16(0.5F); + void* d_a = gpu.Alloc(abf.size() * 2); + void* d_a2 = gpu.Alloc(a2bf.size() * 2); + void* d_nw = gpu.Alloc(nw.size() * 2); + void* d_w = gpu.Alloc(wq.size()); + void* d_w2 = gpu.Alloc(wq2.size()); + void* d_o = gpu.Alloc(2 * static_cast(n)); + gpu.Copy(gq, d_a, abf.data(), abf.size() * 2); + gpu.Copy(gq, d_a2, a2bf.data(), a2bf.size() * 2); + gpu.Copy(gq, d_nw, nw.data(), nw.size() * 2); + gpu.Copy(gq, d_w, wq.data(), wq.size()); + gpu.Copy(gq, d_w2, wq2.data(), wq2.size()); + + EnvNormQuantGuard on(true); + vt::rocm::NormQuantResetForTesting(); + // produce a token for d_a + Tensor xt = DevTensor(d_a, DType::kBF16, {1, k}); + Tensor wt = DevTensor(d_nw, DType::kBF16, {k}); + void* d_norm = gpu.Alloc(sizeof(uint16_t) * static_cast(k)); + Tensor nout = DevTensor(d_norm, DType::kBF16, {1, k}); + vt::RmsNorm(gq, nout, xt, wt, vt::RmsNormArgs{1e-6f, false}); + // non-matching consumer (different ptr/shape): must take the standalone + // quant AND invalidate the token... + Tensor at2 = DevTensor(d_a2, DType::kBF16, {1, k2}); + Tensor bt2 = DevTensor(d_w2, DType::kQ4_K, {n, k2}); + Tensor oo = DevTensor(d_o, DType::kBF16, {1, n}); + vt::MatmulBTQuant(gq, oo, at2, bt2); + gpu.Synchronize(gq); + auto c = vt::rocm::NormQuantCountsForTesting(); + CHECK(c.producers == 1); + CHECK(c.consumers_fused == 0); + CHECK(c.consumers_standalone == 1); + // ...so even a shape-matching call on the OLD buffer now goes standalone + Tensor bt = DevTensor(d_w, DType::kQ4_K, {n, k}); + Tensor nout2 = DevTensor(d_norm, DType::kBF16, {1, k}); + vt::MatmulBTQuant(gq, oo, nout2, bt); + gpu.Synchronize(gq); + c = vt::rocm::NormQuantCountsForTesting(); + CHECK(c.consumers_fused == 0); + CHECK(c.consumers_standalone == 2); + gpu.Free(d_norm); + gpu.Free(d_a); gpu.Free(d_a2); gpu.Free(d_nw); gpu.Free(d_w); gpu.Free(d_w2); gpu.Free(d_o); + gpu.DestroyQueue(gq); +} diff --git a/tests/vt/test_rocm_skinny_f32.cpp b/tests/vt/test_rocm_skinny_f32.cpp new file mode 100644 index 000000000..fb9dbc677 --- /dev/null +++ b/tests/vt/test_rocm_skinny_f32.cpp @@ -0,0 +1,300 @@ +// vllm.cpp original (vt runtime); no upstream mirror. +// +// GFX1100-TG200 lever B2 focused gate: the f32-OUTPUT decode-skinny arm +// (VT_SKINNY_BF16=1) for bf16-in/f32-out MatmulBT at M<=4. The engine +// population that motivates it is the Qwen3.5 GDN BA pair +// (ProjectGdnBA, qwen3_5.cpp:3663-3664): N=32, K=2560, m=1, which today +// falls through every decode-skinny gate in MatmulBTKernelRocm (all require +// a bf16 output) onto hipblasGemmEx -> rocBLAS's large-M Tensile tile +// MT128x32x16 (~73.7us to stream a 164 KiB weight; evidence file section +// 15.1). +// +// Numerics contract: the arm is NOT bit-exact vs the default route by +// construction (different reduction order), so unlike test_rocm_quant_dot +// this gate asserts the sibling 1e-6 NMSE band vs the CPU oracle on BOTH +// arms, a tight ON-vs-OFF agreement band, and ROUTING witnesses through +// host-side dispatch counters (outputs cannot witness routing here because +// both arms are numerically correct). +// +// RED-first contract: before the seam exists this file fails to LINK +// (SkinnyF32RouteCountsForTesting undefined) and the routing cases fail +// behaviorally once counters exist but the env arm does not engage. +// +// Skips cleanly when the build has HIP but the box has no AMD GPU. +#include + +#include +#include +#include +#include +#include +#include + +#include "vt/backend.h" +#include "vt/device.h" +#include "vt/dtype.h" +#include "vt/ops.h" +#include "vt/rocm/rocm_runtime.h" + +using vt::Backend; +using vt::Device; +using vt::DeviceType; +using vt::DType; +using vt::Queue; +using vt::Tensor; + +namespace vt::rocm { +// Host-side routing witness (the test_rocm_quant_dot.cpp F1/F2 convention): +// process-global counters bumped on exactly the branch taken per +// bf16-in/f32-out MatmulBT dispatch. Both arms are numerically correct, so +// no output comparison can witness routing -- these integers can. +struct SkinnyF32RouteCounts { + long long blas; // fell through to hipblasGemmEx (default route) + long long skinny; // took the VT_SKINNY_BF16 wvSplitK-class arm +}; +SkinnyF32RouteCounts SkinnyF32RouteCountsForTesting(); +void SkinnyF32ResetRouteCountsForTesting(); +} // namespace vt::rocm + +namespace { + +Device GpuDev() { return Device{DeviceType::kROCM, 0}; } + +// test_rocm_quant_dot.cpp:79 — the band the sibling gates hold their arms to. +constexpr double kMaxNmseVsCpu = 1e-6; + +double Nmse(const std::vector& got, const std::vector& ref) { + double num = 0, den = 0; + for (size_t i = 0; i < ref.size(); ++i) { + const double d = static_cast(got[i]) - static_cast(ref[i]); + num += d * d; + den += static_cast(ref[i]) * static_cast(ref[i]); + } + return den > 0 ? num / den : num; +} + +Tensor DevTensor(void* p, DType dt, const std::vector& shape) { + Tensor t; + t.data = p; + t.dtype = dt; + t.device = GpuDev(); + t.rank = static_cast(shape.size()); + int64_t stride = 1; + for (int i = t.rank - 1; i >= 0; --i) { + t.shape[i] = shape[static_cast(i)]; + t.stride[i] = stride; + stride *= shape[static_cast(i)]; + } + return t; +} + +struct EnvGuard { + explicit EnvGuard(bool on) { ::setenv("VT_SKINNY_BF16", on ? "1" : "0", 1); } + ~EnvGuard() { ::unsetenv("VT_SKINNY_BF16"); } +}; + +std::vector RandomBf16(size_t n, uint32_t seed) { + std::vector v(n); + uint32_t s = seed; + for (size_t i = 0; i < n; ++i) { + s = s * 1664525u + 1013904223u; + // Small-magnitude values: keeps both arms' f32 accumulation well- + // conditioned so the NMSE bands measure reduction order, not conditioning. + const float f = (static_cast(s >> 8) / 8388608.0f - 1.0f) * 0.125f; + v[i] = vt::F32ToBF16(f); + } + return v; +} + +std::vector CpuOracleBt(const std::vector& a_bf16, + const std::vector& b_bf16, int64_t m, + int64_t n, int64_t k) { + std::vector out(static_cast(m * n), 0.0f); + for (int64_t i = 0; i < m; ++i) + for (int64_t j = 0; j < n; ++j) { + double acc = 0.0; + for (int64_t l = 0; l < k; ++l) + acc += static_cast( + vt::BF16ToF32(a_bf16[static_cast(i * k + l)])) * + static_cast( + vt::BF16ToF32(b_bf16[static_cast(j * k + l)])); + out[static_cast(i * n + j)] = static_cast(acc); + } + return out; +} + +// One shape, both arms: returns per-arm outputs and asserts the shared +// contract (oracle band on both, ON-vs-OFF agreement). Route deltas are +// returned so callers can assert routing too. +struct ArmRun { + std::vector> out; + long long blas_delta; + long long skinny_delta; +}; + +ArmRun RunBothArms(Backend& gpu, Queue gq, const std::vector& a_bf16, + const std::vector& b_bf16, int64_t m, int64_t n, + int64_t k) { + ArmRun run; + void* d_a = gpu.Alloc(a_bf16.size() * 2); + void* d_b = gpu.Alloc(b_bf16.size() * 2); + gpu.Copy(gq, d_a, a_bf16.data(), a_bf16.size() * 2); + gpu.Copy(gq, d_b, b_bf16.data(), b_bf16.size() * 2); + run.out.resize(2); + vt::rocm::SkinnyF32ResetRouteCountsForTesting(); + const auto before = vt::rocm::SkinnyF32RouteCountsForTesting(); + for (int arm = 0; arm < 2; ++arm) { + void* d_o = gpu.Alloc(4 * static_cast(m * n)); + { + EnvGuard guard(arm == 1); + Tensor at = DevTensor(d_a, DType::kBF16, {m, k}); + Tensor bt = DevTensor(d_b, DType::kBF16, {n, k}); + Tensor ot = DevTensor(d_o, DType::kF32, {m, n}); + vt::MatmulBT(gq, ot, at, bt); + run.out[arm].resize(static_cast(m * n), 0.0f); + gpu.Copy(gq, run.out[arm].data(), d_o, run.out[arm].size() * 4); + gpu.Synchronize(gq); + } + gpu.Free(d_o); + } + const auto after = vt::rocm::SkinnyF32RouteCountsForTesting(); + run.blas_delta = after.blas - before.blas; + run.skinny_delta = after.skinny - before.skinny; + gpu.Free(d_a); + gpu.Free(d_b); + return run; +} + +} // namespace + +TEST_CASE("ROCm f32-out decode-skinny arm (VT_SKINNY_BF16=1): NMSE vs CPU oracle and routing witnesses") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm f32-out skinny gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + + struct ShapeCase { + int64_t m, n, k; + bool arm_serves; // expected VT_SKINNY_BF16=1 routing decision + const char* name; + }; + // Engine-realistic anchor first: EXACTLY the ProjectGdnBA decode shape + // (evidence 15.1 rows 3+4). Then gate-boundary edges: even-N variants the + // kernel serves, and the exclusions (odd N, K%8!=0, m beyond the skinny + // range) that must stay on the default BLAS route under ON. + const std::vector shapes = { + {1, 32, 2560, true, "gdn-ba-engine-shape"}, + {1, 64, 4096, true, "even-n-larger-k"}, + {1, 10, 512, true, "minimal-even-n"}, + {4, 32, 2560, true, "m-at-upper-edge"}, + {1, 33, 2560, false, "odd-n-stays-blas"}, + {2, 33, 2560, false, "odd-n-and-m2-stays-blas"}, + {1, 32, 12, false, "k-not-multiple-of-8"}, + {5, 32, 2560, false, "m-past-skinny-range"}, + {1, 8, 2560, false, "n-at-feature-floor"}, + }; + for (const ShapeCase& sc : shapes) { + CAPTURE(sc.name); + CAPTURE(sc.m); + CAPTURE(sc.n); + CAPTURE(sc.k); + const std::vector a = RandomBf16( + static_cast(sc.m * sc.k), 0x5EEDu + static_cast(sc.n)); + const std::vector b = RandomBf16( + static_cast(sc.n * sc.k), 0xA11CEu + static_cast(sc.k)); + + const std::vector ref = + CpuOracleBt(a, b, sc.m, sc.n, sc.k); + + const ArmRun run = RunBothArms(gpu, gq, a, b, sc.m, sc.n, sc.k); + + // Routing witness over the TWO dispatches (OFF then ON). The counters + // only track the bf16-in/f32-out population with M in [1,4]; inside it, + // OFF always routes to BLAS and ON's branch is decided by the shape gate + // alone; outside it (e.g. m=5) neither dispatch is counted. + const bool in_pop = sc.m <= 4; + CHECK(run.blas_delta == (in_pop ? 1 : 0) + (in_pop && !sc.arm_serves ? 1 : 0)); + CHECK(run.skinny_delta == (sc.arm_serves ? 1 : 0)); + + for (int arm = 0; arm < 2; ++arm) { + CAPTURE(arm); + const double nmse = Nmse(run.out[static_cast(arm)], ref); + CAPTURE(nmse); + CHECK(nmse <= kMaxNmseVsCpu); + } + // Cross-arm agreement (same accumulator precision, different tree). + const double nmse_cross = Nmse(run.out[1], run.out[0]); + CAPTURE(nmse_cross); + CHECK(nmse_cross <= kMaxNmseVsCpu); + } + gpu.DestroyQueue(gq); +} + +TEST_CASE("ROCm f32-out skinny routing witness: TRUE-unset behaves like OFF (default-OFF inertness)") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm f32-out skinny gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + // EnvGuard(false) writes "0" -- it can NEVER witness a true unset. The + // first window below therefore constructs NO guard at all (the F-1 repair; + // test_rocm_quant_dot.cpp F1 convention): run_window only touches the + // environment for the explicit windows, so the true-unset dispatch sees + // getenv()==NULL and the engine default must route to BLAS exactly as an + // explicit "0" does. + const std::vector a = RandomBf16(2560, 0x5EEDu); + const std::vector b = RandomBf16(32 * 2560, 0xA11CEu); + void* d_a = gpu.Alloc(a.size() * 2); + void* d_b = gpu.Alloc(b.size() * 2); + gpu.Copy(gq, d_a, a.data(), a.size() * 2); + gpu.Copy(gq, d_b, b.data(), b.size() * 2); + + enum class WindowEnv { kTrueUnset, kExplicitOff, kExplicitOn }; + const auto run_window = [&](WindowEnv env) { + void* d_o = gpu.Alloc(4 * 32); + std::optional guard; + if (env != WindowEnv::kTrueUnset) { + guard.emplace(env == WindowEnv::kExplicitOn); + } + Tensor at = DevTensor(d_a, DType::kBF16, {1, 2560}); + Tensor bt = DevTensor(d_b, DType::kBF16, {32, 2560}); + Tensor ot = DevTensor(d_o, DType::kF32, {1, 32}); + vt::MatmulBT(gq, ot, at, bt); + gpu.Synchronize(gq); + gpu.Free(d_o); + }; + + vt::rocm::SkinnyF32ResetRouteCountsForTesting(); + { + ::unsetenv("VT_SKINNY_BF16"); // true-unset window + run_window(WindowEnv::kTrueUnset); + } + const auto unset_counts = vt::rocm::SkinnyF32RouteCountsForTesting(); + + vt::rocm::SkinnyF32ResetRouteCountsForTesting(); + { + EnvGuard guard(false); // explicit "0" + run_window(WindowEnv::kExplicitOff); + } + const auto off_counts = vt::rocm::SkinnyF32RouteCountsForTesting(); + + vt::rocm::SkinnyF32ResetRouteCountsForTesting(); + { + EnvGuard guard(true); // "1" + run_window(WindowEnv::kExplicitOn); + } + const auto on_counts = vt::rocm::SkinnyF32RouteCountsForTesting(); + + CHECK(unset_counts.blas == 1); + CHECK(unset_counts.skinny == 0); + CHECK(off_counts.blas == 1); + CHECK(off_counts.skinny == 0); + CHECK(on_counts.blas == 0); + CHECK(on_counts.skinny == 1); + gpu.Free(d_a); + gpu.Free(d_b); + gpu.DestroyQueue(gq); +}