diff --git a/arc-tools/cuda_compile_check.sh b/arc-tools/cuda_compile_check.sh index 9327ecede..ba32f7a77 100755 --- a/arc-tools/cuda_compile_check.sh +++ b/arc-tools/cuda_compile_check.sh @@ -75,8 +75,8 @@ fi step "2/5 compile Arc CUDA kernels (mistralrs-quant + arc-cuda-graph) [features: cuda]" # `--lib`: build.rs runs nvcc on every .cu, but no binary is linked. Keeps this # step GPU-less-safe (arc-cuda-graph's gemv_bench bin would otherwise link -lcuda). -cargo build -p mistralrs-quant --lib --features cuda 2>&1 | tail -5 || fail "mistralrs-quant cuda kernels did not compile" -cargo build -p arc-cuda-graph --lib --features cuda 2>&1 | tail -5 || fail "arc-cuda-graph cuda kernels did not compile" +cargo build -p mistralrs-quant --lib --features cuda 2>&1 | tail -40 || fail "mistralrs-quant --features cuda FAILED. This step does TWO things: nvcc compiles the .cu files, then rustc type-checks the cuda-gated Rust. A 'Compiling N of N kernels' line above means the KERNELS were fine and the RUST did not compile -- read the rustc errors, not the nvcc output." +cargo build -p arc-cuda-graph --lib --features cuda 2>&1 | tail -40 || fail "arc-cuda-graph --features cuda FAILED (nvcc kernels AND/OR cuda-gated Rust -- see above)" ok "Arc CUDA kernels compiled for sm_${CUDA_COMPUTE_CAP}" step "3/5 compile-check CUDA-gated tests (--no-run)" diff --git a/arc-tools/wave64_keystone_gate.sh b/arc-tools/wave64_keystone_gate.sh new file mode 100755 index 000000000..0a90fef96 --- /dev/null +++ b/arc-tools/wave64_keystone_gate.sh @@ -0,0 +1,168 @@ +#!/usr/bin/env bash +# +# Arc — wave64 keystone gate: trellis grouped-GEMM arch specialisation. +# +# What this answers, in order, so a cheap box can answer the first half and an +# H200 the second: +# +# STAGE 1 COMPILE (no GPU needed) — the new SM90+ kernel and the +# status-returning launchers compile for sm_80, sm_90a and sm_100a. +# nvcc cross-compiles, so ANY box with CUDA >= 12.8 can run this; +# an A30/A6000 at ~$0.40/hr is enough. This is the gate that has +# never been run for this change. +# STAGE 2 PARITY (needs sm_90+) — the grouped-GEMM CUDA parity tests. +# On an sm_90 box these exercise the NEW kernel; on sm_80 they +# exercise the old one. Both must pass. +# STAGE 3 A/B TIMING (needs sm_90+) — qtip_grouped_curve on this branch and +# on the baseline ref, same box, same fixture. +# +# D18: three distinct exit codes. 0 = pass. 1 = a genuine failure, the only +# result that should change a decision. 2 = the environment could not answer +# (no CUDA, no GPU, wrong arch) — NEVER conflate this with 1. Every terminal +# outcome, including a timeout, writes a line to the status file, so silence +# can never be read as success. +# +# Usage (main runs this; agents do not call runcrate — D15): +# GATE=/tmp/arc-gate-wave64-keystone +# mkdir -p $GATE && setsid nohup bash arc-tools/wave64_keystone_gate.sh \ +# < /dev/null > $GATE/run.log 2>&1 & +# +# Every artefact lives under GATE_DIR, which is namespaced per D19. Shared, +# un-namespaced log paths have already caused one agent's output to be read as +# another's, and two gates running concurrently on the same box would otherwise +# overwrite each other's status file and curve logs. +# +# Env: +# BASELINE_REF git ref for the A/B baseline (default: origin/master) +# GATE_DIR artefact dir (default: /tmp/arc-gate-wave64-keystone) +# STATUS status file (default: $GATE_DIR/status.txt) +# SKIP_AB set to 1 to run stages 1-2 only +# CURVE_ARGS extra args for qtip_grouped_curve + +set -uo pipefail + +GATE_DIR=${GATE_DIR:-/tmp/arc-gate-wave64-keystone} +mkdir -p "$GATE_DIR" +STATUS=${STATUS:-$GATE_DIR/status.txt} +BASELINE_REF=${BASELINE_REF:-origin/master} +REPO=$(git rev-parse --show-toplevel 2>/dev/null) || { + echo "STAGE0 env FAIL not a git repo" | tee -a "$STATUS"; exit 2; } +cd "$REPO" || exit 2 + +: > "$STATUS" +say() { echo "$*" | tee -a "$STATUS"; } +# Any unexpected exit still writes a terminal line. +trap 'rc=$?; [ "$rc" -ne 0 ] && say "TERMINAL unexpected-exit rc=$rc"; exit $rc' EXIT + +say "wave64 keystone gate branch=$(git rev-parse --abbrev-ref HEAD) head=$(git rev-parse --short HEAD)" + +# --------------------------------------------------------------------------- +# Environment probe. Distinguishes "cannot answer" from "failed". +# --------------------------------------------------------------------------- +if ! command -v nvcc >/dev/null 2>&1; then + export PATH=${CUDA_HOME:-/usr/local/cuda}/bin:$PATH +fi +if ! command -v nvcc >/dev/null 2>&1; then + say "STAGE1 compile ENV-CANNOT-ANSWER no nvcc on PATH" + say "TERMINAL exit=2" + exit 2 +fi +NVCC_VER=$(nvcc --version | sed -n 's/.*release \([0-9.]*\).*/\1/p' | head -1) +GPU_CC=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null \ + | head -1 | tr -d '.' || true) +say "env nvcc=$NVCC_VER gpu_cc=${GPU_CC:-none}" + +# --------------------------------------------------------------------------- +# STAGE 1 — compile for every arch this change claims to support. +# sm_100a needs CUDA >= 12.8; when the toolkit is older that is an +# environment answer, not a failure of the code. +# --------------------------------------------------------------------------- +STAGE1_FAIL=0 +for CC in 80 90 100; do + if [ "$CC" -ge 100 ]; then + case "$NVCC_VER" in + 12.[0-7]*|12.[0-7]|11.*) say "STAGE1 sm_${CC} ENV-CANNOT-ANSWER nvcc $NVCC_VER predates sm_100a"; continue ;; + esac + fi + if CUDA_COMPUTE_CAP=$CC FEATURES="cuda" RUN_GPU_TESTS=no \ + timeout 3600 bash arc-tools/cuda_compile_check.sh > "$GATE_DIR/compile_${CC}.log" 2>&1; then + say "STAGE1 sm_${CC} PASS" + else + rc=$? + [ "$rc" -eq 124 ] && say "STAGE1 sm_${CC} FAIL timeout" || say "STAGE1 sm_${CC} FAIL rc=$rc (see $GATE_DIR/compile_${CC}.log)" + STAGE1_FAIL=1 + fi +done +[ "$STAGE1_FAIL" -eq 1 ] && { say "TERMINAL exit=1 compile gate failed"; exit 1; } + +# --------------------------------------------------------------------------- +# STAGE 2 — grouped-GEMM parity on real silicon. +# --------------------------------------------------------------------------- +if [ -z "${GPU_CC:-}" ]; then + say "STAGE2 parity ENV-CANNOT-ANSWER no GPU visible" + say "STAGE3 ab ENV-CANNOT-ANSWER no GPU visible" + say "TERMINAL exit=0 compile-only (stages 2-3 unanswered)" + exit 0 +fi + +export CUDA_COMPUTE_CAP=$GPU_CC +PARITY_LOG=$GATE_DIR/parity.log +if timeout 5400 cargo test --release -p mistralrs-quant --features cuda \ + cuda_grouped_gemm -- --nocapture --test-threads=1 > "$PARITY_LOG" 2>&1; then + # D18 rule 1: a green result must PROVE work happened. `cargo test` reports + # success for zero tests run, which is exactly the shape being guarded here. + NRUN=$(grep -Eo '[0-9]+ passed' "$PARITY_LOG" | head -1 | grep -Eo '[0-9]+' || echo 0) + if [ "${NRUN:-0}" -lt 2 ]; then + say "STAGE2 parity FAIL only ${NRUN} grouped tests ran (expected >= 2); a filter that matches nothing reports success" + say "TERMINAL exit=1"; exit 1 + fi + say "STAGE2 parity PASS sm_${GPU_CC} tests_passed=${NRUN}" +else + rc=$? + [ "$rc" -eq 124 ] && say "STAGE2 parity FAIL timeout" || say "STAGE2 parity FAIL rc=$rc (see $PARITY_LOG)" + say "TERMINAL exit=1"; exit 1 +fi + +if [ "${SKIP_AB:-0}" = "1" ]; then + say "STAGE3 ab SKIPPED by request" + say "TERMINAL exit=0"; exit 0 +fi +if [ "$GPU_CC" -lt 90 ] 2>/dev/null; then + say "STAGE3 ab ENV-CANNOT-ANSWER sm_${GPU_CC} takes the Ampere path; the change under test is the SM90+ kernel" + say "TERMINAL exit=0 (parity green, delta unmeasured)"; exit 0 +fi + +# --------------------------------------------------------------------------- +# STAGE 3 — A/B against the baseline on the SAME box and fixture. +# --------------------------------------------------------------------------- +run_curve() { # $1 = label, $2 = dir + local out="$GATE_DIR/curve_$1.log" + if ! ( cd "$2" && timeout 5400 cargo run --release -p mistralrs-quant \ + --features cuda --example qtip_grouped_curve -- ${CURVE_ARGS:-} ) > "$out" 2>&1; then + local rc=$? + [ "$rc" -eq 124 ] && say "STAGE3 $1 FAIL timeout" || say "STAGE3 $1 FAIL rc=$rc (see $out)" + return 1 + fi + # The harness has previously printed a headline while every row was empty. + if ! grep -qE '^[[:space:]]*[0-9]+[[:space:]]*\|' "$out"; then + say "STAGE3 $1 FAIL ran clean but produced no data rows (see $out)" + return 1 + fi + say "STAGE3 $1 PASS rows=$(grep -cE '^[[:space:]]*[0-9]+[[:space:]]*\|' "$out") (see $out)" + return 0 +} + +run_curve head "$REPO" || { say "TERMINAL exit=1"; exit 1; } + +BASE_WT=$GATE_DIR/baseline-worktree +rm -rf "$BASE_WT" +if ! git worktree add --detach "$BASE_WT" "$BASELINE_REF" >/dev/null 2>&1; then + say "STAGE3 baseline ENV-CANNOT-ANSWER could not create a worktree at $BASELINE_REF" + say "TERMINAL exit=0 (head measured, baseline unmeasured)"; exit 0 +fi +run_curve baseline "$BASE_WT" || { say "TERMINAL exit=1"; exit 1; } +git worktree remove --force "$BASE_WT" >/dev/null 2>&1 || true + +say "STAGE3 ab PASS — compare $GATE_DIR/curve_head.log vs $GATE_DIR/curve_baseline.log" +say "TERMINAL exit=0" +exit 0 diff --git a/mistralrs-quant/examples/qtip_grouped_curve.rs b/mistralrs-quant/examples/qtip_grouped_curve.rs index 638a32791..733ddc7c7 100644 --- a/mistralrs-quant/examples/qtip_grouped_curve.rs +++ b/mistralrs-quant/examples/qtip_grouped_curve.rs @@ -76,7 +76,7 @@ use candle_core::{DType, Device, Result, Tensor}; use mistralrs_quant::{ qtip_expected_distinct_experts, Qtip2bLayer, QtipMode, QuantMethod, QTIP_GATHER_GEMV_MAX_PAIRS, - QTIP_GROUPED_TILE_K, QTIP_GROUPED_TILE_M, QTIP_ONDEVICE_MOE_MAX_TOKENS_ENV, + QTIP_GROUPED_TILE_K, QTIP_ONDEVICE_MOE_MAX_TOKENS_ENV, }; use std::sync::Arc; use std::time::Instant; @@ -313,11 +313,15 @@ fn assert_config_is_measurable(e: usize, batches: &[usize]) -> Result<()> { // expert read once per step" — the model behind the floor column — only // holds while the average pairs per woken expert stays within one tile. let per_expert = max_pairs as f64 / qtip_expected_distinct_experts(max_b, TOPK, e); - if per_expert > QTIP_GROUPED_TILE_M as f64 { + // D16: the m-tile is arch-dependent (16 on Ampere, 64 on SM90+). Ask the + // device -- assuming the Ampere tile on an H200 would bail four times too + // early and mis-state the traffic model below. + let tile_m = tile_m_for_device()?; + if per_expert > tile_m as f64 { candle_core::bail!( - "at B={max_b} each woken expert takes {per_expert:.1} pairs, past GROUPED_TILE_M=\ - {QTIP_GROUPED_TILE_M}: the grouped kernel re-stages weights per m-tile beyond that, \ - so the unique-expert floor column would understate its real traffic." + "at B={max_b} each woken expert takes {per_expert:.1} pairs, past this device's \ + grouped m-tile {tile_m}: the grouped kernel re-stages weights per m-tile beyond \ + that, so the unique-expert floor column would understate its real traffic." ); } Ok(()) @@ -437,6 +441,21 @@ struct Shape { calls_per_layer: f64, } +/// The grouped kernel's m-tile ON THIS DEVICE (D16: 16 on Ampere, 64 on +/// SM90+). Every traffic model in this harness is bounded by it, so it is +/// queried, never assumed -- a wrong tile silently mis-states the grouped +/// path's advantage. +fn tile_m_for_device() -> Result { + #[cfg(feature = "cuda")] + { + mistralrs_quant::qtip_grouped_tile_m() + } + #[cfg(not(feature = "cuda"))] + { + Ok(mistralrs_quant::QTIP_GROUPED_TILE_M) + } +} + /// Packed bytes per expert: 2-bit codes (n*k/4) + F32 row scales. fn bytes_per_expert(s: &Shape) -> f64 { s.n as f64 * s.k as f64 * 0.25 + s.n as f64 * 4.0 @@ -444,7 +463,12 @@ fn bytes_per_expert(s: &Shape) -> f64 { /// Build the routing draws once per batch size and share them across paths, so /// grouped and GEMV are compared on byte-identical traffic. -fn build_routing(dev: &Device, b: usize, e: usize) -> Result<(Vec, f64, f64, f64)> { +fn build_routing( + dev: &Device, + b: usize, + e: usize, + tile_m: usize, +) -> Result<(Vec, f64, f64, f64)> { let mut rng = Rng(ROUTING_SEED ^ (b as u64).wrapping_mul(0x9E37_79B9)); let mut sets = Vec::with_capacity(ROUTING_WINDOWS); let (mut sum_distinct, mut sum_tiles) = (0.0f64, 0.0f64); @@ -454,7 +478,7 @@ fn build_routing(dev: &Device, b: usize, e: usize) -> Result<(Vec, f64, sum_distinct += counts.iter().filter(|c| **c > 0).count() as f64; sum_tiles += counts .iter() - .map(|c| c.div_ceil(QTIP_GROUPED_TILE_M) as f64) + .map(|c| c.div_ceil(tile_m) as f64) .sum::(); sets.push(Tensor::from_vec(idx, (b, TOPK), dev)?); } @@ -479,7 +503,7 @@ fn time_curve( path.select(); for &b in batches { assert_cap_unchanged(cap_at_startup)?; - let (idx_sets, mean_distinct, mean_tiles, pairs) = build_routing(dev, b, e)?; + let (idx_sets, mean_distinct, mean_tiles, pairs) = build_routing(dev, b, e, tile_m_for_device()?)?; assert_routing_matches_model(b, e, mean_distinct)?; let issued_units = if path == Path::Gemv { pairs diff --git a/mistralrs-quant/kernels/qtip/qtip2b_common.cuh b/mistralrs-quant/kernels/qtip/qtip2b_common.cuh index f18d03d6f..c4503dab7 100644 --- a/mistralrs-quant/kernels/qtip/qtip2b_common.cuh +++ b/mistralrs-quant/kernels/qtip/qtip2b_common.cuh @@ -63,6 +63,30 @@ __device__ __forceinline__ uint32_t q2b_state_from_window(uint32_t win16) { return ((r & 0x5555u) << 1) | ((r >> 1) & 0x5555u); } +// --------------------------------------------------------------------------- +// SEQUENTIAL RUN ADVANCE — the cheap half of the same identity. +// +// `q2b_state_from_window` costs ~11 ALU ops (two smem words, two shifts, an +// OR, __brev, and the adjacent-bit swap). That is the price of decoding an +// ISOLATED position. Along a CONTIGUOUS run of symbols the original trellis +// recurrence is far cheaper — one SHF plus one LOP3: +// +// state(t+1) = ((state(t) << 2) | sym[t+1]) & 0xFFFF +// +// so a run of R symbols costs `window + (R-1)*advance` instead of +// `R * window`. At R = 32 that is ~11.3 ops/weight -> ~2 ops/weight of state +// reconstruction. This matters because the trellis GEMM is ALU-issue-bound on +// decode, not bandwidth-bound (98/98 gen-2 GEMV variants measured: latency +// hiding does not pay, per-symbol decode is the ceiling). +// +// The window identity is still what SEEDS each run, so runs stay independent: +// no cross-thread state chain, no warm-up replay. Mirrored and property-tested +// on the Rust side by `grouped::run_states_2b`. +// --------------------------------------------------------------------------- +__device__ __forceinline__ uint32_t q2b_advance_state(uint32_t state, uint32_t sym) { + return ((state << Q2B_K) | (sym & (Q2B_ALPHABET - 1u))) & Q2B_STATE_MASK; +} + // --------------------------------------------------------------------------- // cp.async helpers (sm_80+). `valid == false` uses the src-size-0 form, which // zero-fills the destination without touching the source address — that is diff --git a/mistralrs-quant/kernels/qtip/qtip_grouped_gemm.cu b/mistralrs-quant/kernels/qtip/qtip_grouped_gemm.cu index 05731c307..ea9e6978f 100644 --- a/mistralrs-quant/kernels/qtip/qtip_grouped_gemm.cu +++ b/mistralrs-quant/kernels/qtip/qtip_grouped_gemm.cu @@ -78,9 +78,30 @@ // * Epilogue: pack paired bf16 writes into one 32-bit store when both // columns are in range. // -// SM80+ (mma.sync bf16 + cp.async). Gated by `has_qtip_kernels` in build.rs; -// on older arches the kernel compiles to an empty body (it is unreachable -// at runtime because has_qtip_kernels is off below SM80). +// ------------------------------------------------------------------------- +// ARCH PATHS (D16) +// ------------------------------------------------------------------------- +// Two kernels, picked by the device's RUNTIME compute capability in the +// launcher (not by `__CUDA_ARCH__`, which only says what was compiled): +// +// sm_80 / sm_89 `qtip2b_grouped_gemm_kernel_sm80` — 16-pair m-tile, +// per-position window decode straight into mma fragments. +// Unchanged; this is the path with 5/5 hardware parity. +// sm_90 / sm_100 `qtip2b_grouped_gemm_kernel_wg` — 64-pair m-tile, +// contiguous-run decode into a shared-memory weight tile, +// dynamic shared memory at the arch's real budget. +// below sm_80 the launcher returns cudaErrorNotSupported. It does NOT +// launch: both bodies are compiled out there, and an empty +// kernel returns success while writing nothing (D18). +// +// The tuning note below about "TILE_M 16 -> 32/48" is DONE on the SM90+ path +// and is the largest single win in this file — see the header of the +// warpgroup kernel for why m-tile width, not mma throughput, is the lever. +// +// Gated by `has_qtip_kernels` in build.rs (cc >= 80). cudaforge auto-suffixes +// architectures >= 90 to `sm_90a` / `sm_100a`, so the architecture-specific +// MMAs (`wgmma`, `tcgen05`) are reachable from the build when they land — see +// the dispatch comment above the launchers for what is blocking that today. #include #include @@ -110,6 +131,41 @@ constexpr uint32_t QG_INVALID_PAIR = 0xFFFFFFFFu; // Persistent-CTA grid cap (see tuning notes). constexpr int QG_MAX_GRID = 1024; +// --------------------------------------------------------------------------- +// D16 arch schedule. +// +// SM90+ runs the warpgroup path (`qtip2b_grouped_gemm_kernel_wg`), whose +// m-tile is fixed at 64 by `wgmma.mma_async`'s M=64 shape. SM80/SM89 keep the +// Ampere `mma.sync` path at m-tile 16. +// +// The ROUTING kernel bins pairs into m-tiles, so route and GEMM must agree on +// `tile_m`. That is why this is a host-visible query +// (`qtip2b_grouped_query_schedule`) mirrored by `grouped::grouped_tile_m_for_cc` +// rather than a constant: getting it wrong mis-bins the tile map, and the +// symptom would be wrong numbers, not an error. +// --------------------------------------------------------------------------- +constexpr int QGW_TILE_M = 64; // wgmma M — pairs per m-tile on SM90+ +constexpr int QGW_TILE_N = 64; // wgmma N — weight rows per n-tile +constexpr int QGW_TILE_K = QG_TILE_K; // symbols per k-chunk (host gate unchanged) +constexpr int QGW_WARPS = 4; // one warpgroup +constexpr int QGW_THREADS = QGW_WARPS * 32; +// Symbols each thread decodes as ONE contiguous run: 64 rows x 64 symbols / +// 128 threads. Mirrored by `grouped::GROUPED_DECODE_RUN`. +constexpr int QGW_RUN = (QGW_TILE_N * QGW_TILE_K) / QGW_THREADS; // 32 +// Symbols packed in one 32-bit word of the staged stream. A run is read as +// exactly two such words, so the decode loop switches source words once. +constexpr int QGW_SYMS_PER_WORD = 32 / (int)Q2B_K; // 16 +static_assert(QGW_RUN == 2 * QGW_SYMS_PER_WORD, + "the run decoder reads exactly two 32-bit symbol words per run"); +static_assert(QGW_TILE_K % QGW_RUN == 0, + "runs must tile the k-chunk exactly"); +static_assert(QGW_THREADS % QGW_TILE_N == 0, + "every weight row must be covered by a whole number of threads"); + +__host__ __device__ __forceinline__ int qg_tile_m_for_cc(int cc_major) { + return cc_major >= 9 ? QGW_TILE_M : QG_TILE_M; +} + // cp.async helpers (q2b_cp_async_16 / _4 / _commit / _wait) live in // qtip2b_common.cuh — shared with the gen-2 GEMV pipeline. @@ -309,7 +365,7 @@ __device__ __forceinline__ void q2b_stage_chunk( template __global__ void __launch_bounds__(QG_THREADS) -qtip2b_grouped_gemm_kernel( +qtip2b_grouped_gemm_kernel_sm80( const uint8_t* __restrict__ packed, // [E, n_rows, packed_per_row] const float* __restrict__ row_scales, // [E, n_rows] const T* __restrict__ x, // [n_pairs, num_symbols] (rotated) @@ -444,6 +500,489 @@ qtip2b_grouped_gemm_kernel( #endif // __CUDA_ARCH__ >= 800 } +// =========================================================================== +// SM90 (Hopper) / SM100 (Blackwell) path — D16 arch specialisation. +// =========================================================================== +// +// WHAT THIS CHANGES vs the Ampere kernel above, and WHY. +// +// The trellis GEMM is **ALU-issue-bound on decode**, not bandwidth-bound. That +// is not a model, it is the project's own hardware result: the gen-2 GEMV +// sweep measured 98/98 variants and every latency-hiding axis came back +// negative (split-K, staged width, warp specialization -38%), concluding "the +// ceiling is per-symbol trellis decode serialization — decode FEWER SYMBOLS, +// do not hide latency". So the two levers that matter are (1) decode fewer +// weights for the same output, and (2) spend fewer ALU ops per weight decoded. +// +// (1) WIDER M-TILE — 16 -> 64 pairs. +// Total decode work is `(number of m-tiles) x (expert weight count)`: +// every m-tile re-reads and re-decodes its expert's whole matrix. At 64 +// routed pairs per expert that is 4 m-tiles at TILE_M=16 and 1 at 64 — +// **4x less decode AND 4x less HBM traffic** for the same output. Below +// 16 pairs/expert the two schedules issue the same single tile, so this +// is weakly better everywhere and strictly better at batch. It is also +// exactly `wgmma`'s and `tcgen05`'s M=64, so the schedule is what the +// arch-native MMA will want (see the note at the launcher). +// The cost is 4 m-fragments of accumulator (32 f32/thread) — registers, +// not the bottleneck. +// +// (2) SEQUENTIAL RUN DECODE — one window seed per 32 symbols, not per symbol. +// `q2b_state_from_window` is ~11 ALU ops (two smem words, two shifts, an +// OR, __brev, adjacent-bit swap). The Ampere kernel pays that for EVERY +// weight because the mma B-fragment gives each thread runs of length 2. +// Here each thread instead owns a contiguous run of QGW_RUN symbols of +// one weight row, seeds it once with the window identity, and steps the +// trellis recurrence (`q2b_advance_state`, ~2 ops). State reconstruction +// drops from ~11 ops/weight to ~2.3. Runs stay independent — the seed is +// still the random-access identity — so there is no cross-thread state +// chain and no warm-up replay. Property-tested on CPU by +// `grouped::run_states_2b` (D14: that is a unit test, not a measurement). +// +// Decoding a run forces the weights through shared memory rather than +// straight into mma fragments. That is the right trade here (smem is on +// the LSU, which is idle while the ALU saturates) and it is also +// MANDATORY for the arch-native MMA: `wgmma`'s B operand must live in +// shared memory, and `tcgen05` takes no register operands at all. +// +// (3) THE ARCH'S REAL SHARED-MEMORY BUDGET. The tile uses dynamic shared +// memory and the launcher opts in with `cudaFuncSetAttribute`. Without +// that opt-in CUDA caps a block at 48 KB regardless of the 227 KB that +// compute capability 9.0 and 10.x actually allow (8.9 allows only 99 KB — +// hence the arch gate, not a constant). +// +// THE MOAT IS INTACT: the packed 2-bit bytes are what cross HBM, and decode +// happens inside this kernel into a single 64x64 tile of shared memory that is +// overwritten every k-chunk. Full-size weights never materialise anywhere. +// +// mma.sync.m16n8k16.f32.bf16.bf16 requires sm_80 or higher and follows the +// onion-layer model, so it is valid on sm_90a and sm_100a alike — it is the +// one bf16 tensor-core path that spans Ampere/Hopper/Blackwell. `wgmma` is +// sm_90a-EXCLUSIVE (PTX "Requires sm_90a"; a-suffix targets "do not follow the +// onion layer model" and "cannot be run on later generation devices") and +// `tcgen05` is sm_100a+, so those are two further, mutually exclusive +// specialisations layered on this same producer stage. See the launcher for +// exactly what blocks them today. +// --------------------------------------------------------------------------- + +// Padded smem row strides, in ELEMENTS. Both are chosen so the mma fragment +// reads are bank-conflict-free: a lane reads element `stride*(base+g) + 2*tig` +// as one 32-bit word, i.e. bank `(stride/2*g + tig) mod 32`; stride 72 gives +// `(36g + tig) mod 32 = (4g + tig) mod 32`, which is injective over +// g in 0..7, tig in 0..3. The unpadded stride 64 collapses all eight g to the +// same bank (8-way conflict). +constexpr int QGW_X_STRIDE = QGW_TILE_K + 8; // 72 +constexpr int QGW_B_STRIDE = QGW_TILE_K + 8; // 72 +constexpr int QGW_MFRAGS = QGW_TILE_M / 16; // 4 mma m-fragments per CTA + +// Dynamic-smem byte budget for one CTA of the warpgroup kernel. +template +__host__ __device__ __forceinline__ size_t qgw_smem_bytes() { + return (size_t)2 * QGW_TILE_M * QGW_X_STRIDE * sizeof(T) // s_x (double buffered) + + (size_t)2 * QGW_TILE_N * QG_WP_STRIDE // s_wp (double buffered) + + (size_t)QGW_TILE_N * QGW_B_STRIDE * sizeof(T) // s_bw (decoded, single) + + (size_t)QGW_TILE_M * sizeof(uint32_t); // s_pairs +} + +// Stage one k-chunk: QGW_TILE_M activation rows (gathered through the sorted +// pair ids) and QGW_TILE_N packed weight rows (+4-byte history prefix). +// 128 threads issue 4 activation copies and 1 packed copy each. +template +__device__ __forceinline__ void q2bw_stage_chunk( + int c, + T* s_x_buf, // [QGW_TILE_M][QGW_X_STRIDE] + uint8_t* s_wp_buf, // [QGW_TILE_N][QG_WP_STRIDE] + const uint32_t* s_pairs, // [QGW_TILE_M] + const T* __restrict__ x, + const uint8_t* __restrict__ packed, + uint32_t expert, + int n0, + int n_rows, + int num_symbols, + int packed_per_row +) { + const int tid = threadIdx.x; + // Activations: QGW_TILE_M rows x 8 sixteen-byte pieces = 512 copies. + #pragma unroll + for (int rep = 0; rep < (QGW_TILE_M * QGW_TILE_K * (int)sizeof(T) / 16) / QGW_THREADS; ++rep) { + const int cid = tid + rep * QGW_THREADS; + const int slot = cid >> 3; + const int piece = cid & 7; + const uint32_t pair = s_pairs[slot]; + const bool valid = pair != QG_INVALID_PAIR; + void* dst = (uint8_t*)(s_x_buf + (size_t)slot * QGW_X_STRIDE) + piece * 16; + const uint8_t* src = + (const uint8_t*)(x + (size_t)(valid ? pair : 0u) * num_symbols + + (size_t)c * QGW_TILE_K) + + piece * 16; + q2b_cp_async_16(dst, src, valid); + } + // Packed weights: 64 rows; threads 0..63 stage the 4-byte history prefix, + // threads 64..127 this chunk's 16 bytes. Identical contract to the Ampere + // path (chunk 0 has no history and is zero-filled, which matches the + // all-zero initial trellis state). + { + const int r = tid & (QGW_TILE_N - 1); + const int half = tid / QGW_TILE_N; + const int n = n0 + r; + const bool nvalid = n < n_rows; + const uint8_t* row_base = + packed + ((size_t)expert * n_rows + (size_t)(nvalid ? n : 0)) * packed_per_row + + (size_t)c * (QGW_TILE_K / 4); + uint8_t* srow = s_wp_buf + (size_t)r * QG_WP_STRIDE; + if (half == 0) { + const bool pvalid = nvalid && c > 0; + q2b_cp_async_4(srow + 12, pvalid ? (row_base - 4) : row_base, pvalid); + } else { + q2b_cp_async_16(srow + 16, row_base, nvalid); + } + } +} + +// Decode the staged 64x64 packed tile into `s_bw` as 16-bit activations-dtype, +// one contiguous run of QGW_RUN symbols per thread. +// +// Out-of-range weight rows carry scale 0, which zeroes the (finite) decode of +// their cp.async zero-filled bytes — the same tail-guard-free contract the +// Ampere path uses. +template +__device__ __forceinline__ void q2bw_decode_tile( + T* s_bw, + const uint8_t* s_wp_buf, + const float* __restrict__ row_scales, + uint32_t expert, + int n0, + int n_rows, + uint32_t mult +) { + const int tid = threadIdx.x; + const int r = tid / (QGW_THREADS / QGW_TILE_N); // weight row in the n-tile + const int h = tid % (QGW_THREADS / QGW_TILE_N); // which run of the row + const int ts0 = h * QGW_RUN; + + const uint32_t* w = + reinterpret_cast(s_wp_buf + (size_t)r * QG_WP_STRIDE); + + // Seed: the random-access window identity at the first symbol of the run. + const int bit = 2 * ts0 + QG_WIN_BIT_BASE; + const int wi = bit >> 5; + const uint32_t sh = (uint32_t)(bit & 31); + uint32_t win = w[wi] >> sh; + if (sh > 16u) win |= w[wi + 1] << (32u - sh); + uint32_t state = q2b_state_from_window(win); + + // Raw symbols of the run. Chunk-local symbol `ts` sits at staged-buffer + // bit 2*ts + 128 (byte 16 is where this chunk's 16 bytes begin), so a run + // of 32 symbols is exactly two aligned 32-bit words. + const int sw = (2 * ts0 + 128) >> 5; + const uint32_t b_lo = w[sw]; + const uint32_t b_hi = w[sw + 1]; + uint32_t bits = b_lo >> Q2B_K; // low bits = symbol ts0+1 + + const int n = n0 + r; + const float s = + n < n_rows ? __ldg(row_scales + (size_t)expert * n_rows + n) : 0.0f; + + T* out = s_bw + (size_t)r * QGW_B_STRIDE + ts0; + // Both switch conditions are present even though only one can fire at the + // current unroll factor: the guard is on the SYMBOL INDEX being consumed, + // so the loop stays correct if the unroll or QGW_RUN is ever retuned. + // Keying it on loop parity instead would break silently. + #pragma unroll + for (int i = 0; i < QGW_RUN; i += 2) { + const float w0 = q2b_decode(state, mult) * s; + if (i + 1 == QGW_SYMS_PER_WORD) bits = b_hi; + state = q2b_advance_state(state, bits); + bits >>= Q2B_K; + const float w1 = q2b_decode(state, mult) * s; + if (i + 2 == QGW_SYMS_PER_WORD) bits = b_hi; + state = q2b_advance_state(state, bits); + bits >>= Q2B_K; + // Both halves land in one 32-bit store (out + i is 4-byte aligned: + // QGW_B_STRIDE, QGW_RUN and i are all even). + *reinterpret_cast(out + i) = q2b_pack2(w0, w1); + } +} + +template +__global__ void __launch_bounds__(QGW_THREADS) +qtip2b_grouped_gemm_kernel_wg( + const uint8_t* __restrict__ packed, // [E, n_rows, packed_per_row] + const float* __restrict__ row_scales, // [E, n_rows] + const T* __restrict__ x, // [n_pairs, num_symbols] (rotated) + const uint32_t* __restrict__ sorted_pairs, // [n_pairs] grouped by expert + const uint32_t* __restrict__ tile_expert, // [num_tiles] + const uint32_t* __restrict__ tile_row_start, // [num_tiles] + const uint32_t* __restrict__ offsets, // [E+1] + const uint32_t* __restrict__ num_tiles, // [1] (device-side tile count) + T* __restrict__ y, // [n_pairs, n_rows], pre-zeroed + int n_rows, + int packed_per_row, + int num_symbols, + uint32_t mult +) { +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 900 + extern __shared__ __align__(16) uint8_t qgw_smem[]; + T* s_x = reinterpret_cast(qgw_smem); + uint8_t* s_wp = qgw_smem + (size_t)2 * QGW_TILE_M * QGW_X_STRIDE * sizeof(T); + T* s_bw = reinterpret_cast(s_wp + (size_t)2 * QGW_TILE_N * QG_WP_STRIDE); + uint32_t* s_pairs = + reinterpret_cast(s_bw + (size_t)QGW_TILE_N * QGW_B_STRIDE); + + const int tid = threadIdx.x; + const int warp = tid >> 5; + const int lane = tid & 31; + const int g = lane >> 2; // mma groupID + const int tig = lane & 3; // mma threadID-in-group + + const int n_tiles_n = (n_rows + QGW_TILE_N - 1) / QGW_TILE_N; + const int live = (int)*num_tiles * n_tiles_n; + const int n_chunks = num_symbols / QGW_TILE_K; + + for (int flat = blockIdx.x; flat < live; flat += gridDim.x) { + const int mt = flat / n_tiles_n; + const int nt = flat % n_tiles_n; + + const uint32_t expert = tile_expert[mt]; + const int row_start = (int)tile_row_start[mt]; + const int group_end = (int)offsets[expert + 1]; + const int rows_here = min(QGW_TILE_M, group_end - row_start); + const int n0 = nt * QGW_TILE_N; + + __syncthreads(); + if (tid < QGW_TILE_M) { + s_pairs[tid] = + tid < rows_here ? sorted_pairs[row_start + tid] : QG_INVALID_PAIR; + } + __syncthreads(); + + float acc[QGW_MFRAGS][2][4]; + #pragma unroll + for (int m = 0; m < QGW_MFRAGS; ++m) { + #pragma unroll + for (int f = 0; f < 2; ++f) { + #pragma unroll + for (int i = 0; i < 4; ++i) acc[m][f][i] = 0.0f; + } + } + + q2bw_stage_chunk(0, s_x, s_wp, s_pairs, x, packed, expert, n0, n_rows, + num_symbols, packed_per_row); + q2b_cp_commit(); + + for (int c = 0; c < n_chunks; ++c) { + const int buf = c & 1; + T* s_x_cur = s_x + (size_t)buf * QGW_TILE_M * QGW_X_STRIDE; + uint8_t* s_wp_cur = s_wp + (size_t)buf * QGW_TILE_N * QG_WP_STRIDE; + if (c + 1 < n_chunks) { + q2bw_stage_chunk( + c + 1, + s_x + (size_t)(buf ^ 1) * QGW_TILE_M * QGW_X_STRIDE, + s_wp + (size_t)(buf ^ 1) * QGW_TILE_N * QG_WP_STRIDE, + s_pairs, x, packed, expert, n0, n_rows, num_symbols, + packed_per_row); + q2b_cp_commit(); + q2b_cp_wait<1>(); + } else { + q2b_cp_wait<0>(); + } + // Also fences the previous chunk's mma reads of s_bw, which is + // single-buffered (decode is compute, not a copy, so there is + // nothing to overlap it with — warp specialization measured -38% + // on this format). + __syncthreads(); + + q2bw_decode_tile(s_bw, s_wp_cur, row_scales, expert, n0, n_rows, + mult); + __syncthreads(); + + #pragma unroll + for (int kf = 0; kf < QGW_TILE_K / 16; ++kf) { + const int kb = kf * 16; + // B fragments: this warp's 16 weight rows, two 8-wide blocks. + // Each 32-bit load is a (k, k+1) pair of one row — exactly the + // mma B fragment, no shuffling and no re-decode. + uint32_t b[2][2]; + #pragma unroll + for (int f = 0; f < 2; ++f) { + const T* brow = + s_bw + (size_t)(warp * 16 + f * 8 + g) * QGW_B_STRIDE + kb; + b[f][0] = *reinterpret_cast(brow + tig * 2); + b[f][1] = *reinterpret_cast(brow + tig * 2 + 8); + } + #pragma unroll + for (int m = 0; m < QGW_MFRAGS; ++m) { + const int m0 = m * 16; + const T* xr0 = s_x_cur + (size_t)(m0 + g) * QGW_X_STRIDE + kb; + const T* xr8 = s_x_cur + (size_t)(m0 + g + 8) * QGW_X_STRIDE + kb; + uint32_t a[4]; + a[0] = *reinterpret_cast(xr0 + tig * 2); + a[1] = *reinterpret_cast(xr8 + tig * 2); + a[2] = *reinterpret_cast(xr0 + tig * 2 + 8); + a[3] = *reinterpret_cast(xr8 + tig * 2 + 8); + #pragma unroll + for (int f = 0; f < 2; ++f) { + q2b_mma16816::run(acc[m][f], a, b[f][0], b[f][1]); + } + } + } + // REQUIRED: the next iteration issues cp.async into the buffer + // this loop just read, before reaching any other barrier. Without + // this the staging races the mma reads — the corruption would be + // silent (wrong activations, no error). + __syncthreads(); + } + + // Epilogue: same scatter as the Ampere path, over 4 m-fragments. + #pragma unroll + for (int m = 0; m < QGW_MFRAGS; ++m) { + #pragma unroll + for (int f = 0; f < 2; ++f) { + const int nb = n0 + warp * 16 + f * 8 + tig * 2; + #pragma unroll + for (int half = 0; half < 2; ++half) { + const int slot = m * 16 + g + half * 8; + if (slot >= rows_here) continue; + const uint32_t pair = s_pairs[slot]; + T* out = y + (size_t)pair * n_rows + nb; + if (nb < n_rows) out[0] = q2b_from_f32(acc[m][f][half * 2]); + if (nb + 1 < n_rows) out[1] = q2b_from_f32(acc[m][f][half * 2 + 1]); + } + } + } + } +#else + // Compiled for a pre-Hopper target: this entry point is never dispatched + // (the host picks by runtime compute capability and errors out when no + // path matches), so an empty body here is unreachable rather than a silent + // no-op. Reference the parameters so the signature cannot drift. + (void)packed; (void)row_scales; (void)x; (void)sorted_pairs; + (void)tile_expert; (void)tile_row_start; (void)offsets; (void)num_tiles; + (void)y; (void)n_rows; (void)packed_per_row; (void)num_symbols; (void)mult; +#endif // __CUDA_ARCH__ >= 900 +} + +// Claim the arch's real per-block shared-memory budget for the warpgroup +// kernel. Without this opt-in CUDA caps a block at 48 KB regardless of the +// 227 KB that compute capability 9.0 and 10.x allow (8.9 allows only 99 KB) -- +// PR #94 met the live form of this. The call is unconditional rather than +// gated on "> 48 KB" so the path is exercised on every process and a later +// tile-size increase cannot silently hit the cap. +// +// Templates cannot be given C language linkage, so this lives here rather than +// beside the launchers. +template +cudaError_t qgw_prepare_smem() { + static bool done = false; + static cudaError_t status = cudaSuccess; + if (!done) { + status = cudaFuncSetAttribute(qtip2b_grouped_gemm_kernel_wg, + cudaFuncAttributeMaxDynamicSharedMemorySize, + (int)qgw_smem_bytes()); + done = true; + } + return status; +} + +// --------------------------------------------------------------------------- +// THE ARCH WITNESS (D18 #12). +// +// The dispatch below picks the warpgroup path from the DEVICE's compute +// capability, but `qtip2b_grouped_gemm_kernel_wg`'s body is `#if __CUDA_ARCH__ +// >= 900` -- a property of the BINARY, not the device. Those two can disagree, +// and when they do nothing says so: +// +// a build whose archive carries only sm_80 SASS, run on an H200, JITs from +// `compute_80` PTX -- PTX in which the warpgroup body was already compiled +// away. `cc_major` is 9, so dispatch takes the wide branch, launches a kernel +// that writes nothing, and `cudaGetLastError()` returns success. The caller +// receives its `alloc_zeros` buffer as a shape-correct, error-free, ALL-ZERO +// MoE layer. +// +// That is the same shape as the launch-status hole this PR closes, one level +// up: the absence of a signal read as a specific signal. `ARC_CUDA_ARCHS` +// (PR #108) is what makes the sm_90 cubin exist, but a guarantee that depends +// on another PR having landed, on nobody building without the env var, and on +// no branch being cut from an intermediate state is a sequencing hope. This is +// the mechanical version. +// +// The witness is a real kernel carrying the SAME `__CUDA_ARCH__` guard as the +// one it vouches for, compiled in the same translation unit with the same arch +// flags. It reports what the running binary actually contains: 900 if the +// warpgroup body survived compilation, otherwise the arch it was compiled for. +__global__ void qgw_arch_witness_kernel(int* out) { +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 900 + if (threadIdx.x == 0) *out = 900; +#else + if (threadIdx.x == 0) *out = (int)__CUDA_ARCH__; +#endif +} + +struct QgwArchWitness { + cudaError_t status; + int arch; // 900 => the warpgroup body is present in this binary +}; + +// One launch and one sync, on first use. +// +// It deliberately does NOT default `arch` on failure: a witness that could not +// be taken is `status != cudaSuccess`, and the caller must refuse rather than +// assume. "Could not answer" is not "answered yes". +QgwArchWitness qgw_take_arch_witness() { + QgwArchWitness w{cudaSuccess, 0}; + + // If the legacy stream is capturing, a launch on it is APPENDED TO A GRAPH + // rather than executed, and the readback below would return the memset 0 -- + // reported as "this binary has no SM90 device code" and blocking a path + // that works perfectly. A false negative is still a wrong answer, so refuse + // to answer here and let the caller retry outside capture. + cudaStreamCaptureStatus cap = cudaStreamCaptureStatusNone; + const cudaError_t cap_err = cudaStreamIsCapturing(0, &cap); + if (cap_err != cudaSuccess) { + w.status = cap_err; + return w; + } + if (cap != cudaStreamCaptureStatusNone) { + w.status = cudaErrorStreamCaptureUnsupported; + return w; + } + + int* d = nullptr; + w.status = cudaMalloc(&d, sizeof(int)); + if (w.status != cudaSuccess) return w; + w.status = cudaMemset(d, 0, sizeof(int)); + if (w.status == cudaSuccess) { + qgw_arch_witness_kernel<<<1, 1>>>(d); + // Same D18 pairing as every launcher here: a sync returns success for + // a kernel that never launched, so the launch status is read too. + const cudaError_t launch_err = cudaGetLastError(); + const cudaError_t sync_err = cudaStreamSynchronize(0); + w.status = (launch_err != cudaSuccess) ? launch_err : sync_err; + if (w.status == cudaSuccess) { + w.status = cudaMemcpy(&w.arch, d, sizeof(int), cudaMemcpyDeviceToHost); + } + } + cudaFree(d); // the witness result is already on the host + return w; +} + +// Cached only once it has actually ANSWERED. Caching a failure would be the +// same bug in slow motion: a witness that could not be taken during graph +// capture would poison every later call, and the path would stay refused long +// after the condition cleared. So a non-success result is retried. +// +// Two threads racing the first call may both take the witness. That is benign +// and deliberate: the probe is idempotent and its result is a deterministic +// property of the binary, so both writes store identical bytes. The +// alternative -- a lock on a path that runs per MoE layer -- costs more than +// the duplicate. +const QgwArchWitness& qgw_arch_witness() { + static QgwArchWitness w{cudaErrorNotReady, 0}; // "not taken yet" + if (w.status != cudaSuccess) { + w = qgw_take_arch_witness(); + } + return w; +} + } // anonymous namespace // ============================================================================ @@ -452,9 +991,49 @@ qtip2b_grouped_gemm_kernel( extern "C" { +// --------------------------------------------------------------------------- +// D18: every launcher below returns a `cudaError_t` as `int`, and every one of +// them ends with `cudaGetLastError()`. +// +// Before this, the grouped path had NO launch-status check anywhere (the GEMV +// path's `check_gather_gemv_pairs` guard exists precisely because a discarded +// launch status hands the caller the zero-initialized output buffer — a +// silently all-zero MoE layer, shape-correct and error-free). `cudaGetLastError` +// is the call that distinguishes "the kernel ran" from "the kernel never +// launched": `cudaDeviceSynchronize` returns success for a kernel that never +// started. +// --------------------------------------------------------------------------- + +// Which m-tile schedule this device's compute capability selects, and the +// capability itself. Returns a cudaError_t; on failure the out-params are +// untouched and the caller MUST NOT fall back to a default (a wrong `tile_m` +// silently mis-bins the tile map). +// +// Mirrored by `grouped::grouped_tile_m_for_cc` on the Rust side. +int qtip2b_grouped_query_schedule(int* out_cc_major, int* out_cc_minor, int* out_tile_m, + int* out_arch_witness) { + int dev = 0; + cudaError_t err = cudaGetDevice(&dev); + if (err != cudaSuccess) return (int)err; + int major = 0, minor = 0; + err = cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev); + if (err != cudaSuccess) return (int)err; + err = cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev); + if (err != cudaSuccess) return (int)err; + // What the DEVICE is, and separately what the BINARY contains. Reporting + // only the first is what lets an all-zero MoE layer look valid. + const QgwArchWitness& w = qgw_arch_witness(); + if (w.status != cudaSuccess) return (int)w.status; + *out_cc_major = major; + *out_cc_minor = minor; + *out_tile_m = qg_tile_m_for_cc(major); + *out_arch_witness = w.arch; + return (int)cudaSuccess; +} + // Routing: histogram + scans/tile-map + grouped scatter, all on `stream`. // `d_counts` and `d_cursors` must be zero-initialized by the caller. -void launch_qtip2b_moe_route( +int launch_qtip2b_moe_route( const uint32_t* d_indices, // [n_pairs] uint32_t* d_counts, // [E], zeroed uint32_t* d_offsets, // [E+1] @@ -469,6 +1048,12 @@ void launch_qtip2b_moe_route( int tile_m, cudaStream_t stream ) { + // "Nothing to route" is a legitimate state, but "the grid would have been + // empty" is not something the caller can distinguish from "it ran" once + // the zeroed buffers come back, so it is reported. + if (n_pairs <= 0 || num_experts <= 0 || tile_m <= 0) { + return (int)cudaErrorInvalidConfiguration; + } const int blocks_p = (n_pairs + QG_ROUTE_THREADS - 1) / QG_ROUTE_THREADS; qtip2b_moe_histogram_kernel<<>>( d_indices, d_counts, n_pairs, num_experts); @@ -477,10 +1062,45 @@ void launch_qtip2b_moe_route( d_num_tiles, num_experts, tile_m); qtip2b_moe_scatter_kernel<<>>( d_indices, d_offsets, d_cursors, d_sorted_pairs, n_pairs, num_experts); + return (int)cudaGetLastError(); } +// --------------------------------------------------------------------------- +// D16 DISPATCH. Which arch path a launch takes is decided here, from the +// device's runtime compute capability — not from `__CUDA_ARCH__` alone, which +// only says what was compiled. +// +// cc >= 9.0 -> `qtip2b_grouped_gemm_kernel_wg` (64-pair m-tile, run +// decode, dynamic smem at the arch's real budget) +// cc >= 8.0 -> `qtip2b_grouped_gemm_kernel_sm80` (16-pair m-tile) +// cc < 8.0 -> **error**, never a silent no-op. Both kernel bodies are +// `#if`-compiled away below their arch, and a launch of an +// empty kernel writes nothing while returning success — the +// caller would receive its zero-initialized buffer as a valid, +// shape-correct, all-zero MoE output (D18). +// +// WHAT IS NOT HERE YET, stated plainly rather than implied: +// `wgmma.mma_async` (sm_90a) and `tcgen05.mma` (sm_100a/sm_103a) are the +// arch-native MMAs, and this kernel's producer stage — decode into a shared +// -memory tile — is already the shape both of them require (wgmma's B operand +// must be in shared memory; tcgen05 takes no register operands at all). The +// consumer here is still `mma.sync.m16n8k16`, which PTX documents as "requires +// sm_80 or higher" under the onion-layer model and is therefore valid, but not +// native, on sm_90a and sm_100a. +// +// The one thing blocking the swap is the B-operand shared-memory layout: the +// 64-bit matrix descriptor pins swizzle mode, leading-dimension byte offset +// and stride byte offset to a core-matrix layout that the PTX doc specifies +// semantically rather than by a numeric table, and a wrong descriptor produces +// wrong numbers, not an error. That is not something to guess at on the one +// kernel nobody else can adopt. It is settled by a ~30-line descriptor probe +// on any sm_90a box (build a known tile, run one wgmma, diff against this +// kernel's mma.sync result) — minutes of GPU time, and it converts the swap +// from guesswork into a checked change. +// --------------------------------------------------------------------------- + #define Q2B_GROUPED_GEMM_LAUNCHER(NAME, T) \ - void NAME(const uint8_t* d_packed, \ + int NAME(const uint8_t* d_packed, \ const float* d_row_scales, \ const T* d_x_rotated, \ const uint32_t* d_sorted_pairs, \ @@ -495,15 +1115,50 @@ void launch_qtip2b_moe_route( int max_m_tiles, \ uint32_t mult, \ cudaStream_t stream) { \ - const int n_tiles_n = (n_rows + QG_TILE_N - 1) / QG_TILE_N; \ + int dev = 0; \ + cudaError_t err = cudaGetDevice(&dev); \ + if (err != cudaSuccess) return (int)err; \ + int cc_major = 0; \ + err = cudaDeviceGetAttribute( \ + &cc_major, cudaDevAttrComputeCapabilityMajor, dev); \ + if (err != cudaSuccess) return (int)err; \ + /* Pre-Ampere: BOTH kernel bodies are compiled out, so a launch would \ + write nothing and still return success. Report it (D18). */ \ + if (cc_major < 8) return (int)cudaErrorNotSupported; \ + const bool wide = cc_major >= 9; \ + const int tile_n = wide ? QGW_TILE_N : QG_TILE_N; \ + const int n_tiles_n = (n_rows + tile_n - 1) / tile_n; \ const long max_tiles = (long)max_m_tiles * (long)n_tiles_n; \ const int grid = \ (int)(max_tiles < (long)QG_MAX_GRID ? max_tiles : (long)QG_MAX_GRID); \ - if (grid <= 0) return; \ - qtip2b_grouped_gemm_kernel<<>>( \ - d_packed, d_row_scales, d_x_rotated, d_sorted_pairs, \ - d_tile_expert, d_tile_row_start, d_offsets, d_num_tiles, d_y, \ - n_rows, packed_per_row, num_symbols, mult); \ + /* An empty grid means the caller's tile bookkeeping disagrees with \ + ours; that must not be reported as a completed launch (D18). */ \ + if (grid <= 0) return (int)cudaErrorInvalidConfiguration; \ + if (wide) { \ + /* D18 #12: the device is SM90+, but is the SM90 kernel BODY in \ + this binary? If not, the launch below writes nothing and \ + reports success. Defence in depth -- the Rust caller checks the \ + same witness and can say it in words; this keeps the C ABI safe \ + for any other caller. */ \ + const QgwArchWitness& w = qgw_arch_witness(); \ + if (w.status != cudaSuccess) return (int)w.status; \ + if (w.arch < 900) return (int)cudaErrorInvalidDeviceFunction; \ + err = qgw_prepare_smem(); \ + if (err != cudaSuccess) return (int)err; \ + const size_t smem = qgw_smem_bytes(); \ + qtip2b_grouped_gemm_kernel_wg \ + <<>>( \ + d_packed, d_row_scales, d_x_rotated, d_sorted_pairs, \ + d_tile_expert, d_tile_row_start, d_offsets, d_num_tiles, \ + d_y, n_rows, packed_per_row, num_symbols, mult); \ + } else { \ + qtip2b_grouped_gemm_kernel_sm80 \ + <<>>( \ + d_packed, d_row_scales, d_x_rotated, d_sorted_pairs, \ + d_tile_expert, d_tile_row_start, d_offsets, d_num_tiles, \ + d_y, n_rows, packed_per_row, num_symbols, mult); \ + } \ + return (int)cudaGetLastError(); \ } Q2B_GROUPED_GEMM_LAUNCHER(launch_qtip2b_grouped_gemm_bf16, __nv_bfloat16) diff --git a/mistralrs-quant/src/lib.rs b/mistralrs-quant/src/lib.rs index e56ef2336..5b9fa1576 100644 --- a/mistralrs-quant/src/lib.rs +++ b/mistralrs-quant/src/lib.rs @@ -105,6 +105,11 @@ pub use qtip::{ TrellisSearch, QTIP2B_MCG_MULT, QTIP_GATHER_GEMV_MAX_PAIRS, QTIP_GROUPED_TILE_K, QTIP_GROUPED_TILE_M, QTIP_GROUPED_TILE_N, QTIP_ONDEVICE_MOE_MAX_TOKENS_ENV, }; +// D16: the grouped GEMM's m-tile is arch-dependent, so harnesses must query the +// device rather than assume the Ampere constant. The query is a CUDA FFI call, +// so it only exists on the cuda build. +#[cfg(feature = "cuda")] +pub use qtip::qtip_grouped_tile_m; pub use td_moe_factored::TuckerFactoredLayer; pub use unquantized::UnquantLinear; pub use utils::bake_budget::{ diff --git a/mistralrs-quant/src/qtip/cuda_ops.rs b/mistralrs-quant/src/qtip/cuda_ops.rs index 88cb41c9f..601dd3591 100644 --- a/mistralrs-quant/src/qtip/cuda_ops.rs +++ b/mistralrs-quant/src/qtip/cuda_ops.rs @@ -1842,7 +1842,7 @@ pub(crate) fn grouped_gemm_2b_cuda( indices: &Tensor, in_features: usize, ) -> Result { - use super::grouped::{grouped_max_m_tiles, GROUPED_TILE_K, GROUPED_TILE_M}; + use super::grouped::{grouped_max_m_tiles, grouped_tile_m_for_cc, GROUPED_TILE_K}; if !ffi::HAVE_QTIP_KERNELS { candle_core::bail!("qtip2b grouped gemm CUDA: kernels not compiled in"); @@ -1907,7 +1907,48 @@ pub(crate) fn grouped_gemm_2b_cuda( candle_core::bail!("qtip2b grouped gemm CUDA: x_rotated and indices must live on CUDA"); } - let max_m_tiles = grouped_max_m_tiles(n_pairs, num_experts); + // D16: the m-tile schedule is arch-dependent (the SM90+ kernel's m-tile is + // 64, the Ampere kernel's is 16). The routing kernel bins pairs into + // m-tiles, so route, grid bound and GEMM must all use the SAME value -- + // ask the CUDA side which one this device selected rather than assuming. + let (cc_major, cc_minor, tile_m, arch_witness) = { + let (mut major, mut minor, mut tile_m, mut witness) = (0i32, 0i32, 0i32, 0i32); + let rc = unsafe { + ffi::qtip2b_grouped_query_schedule(&mut major, &mut minor, &mut tile_m, &mut witness) + }; + if rc != 0 { + candle_core::bail!( + "qtip2b grouped gemm CUDA: could not query the device compute \ + capability (rc={rc}); refusing to guess a tile schedule" + ); + } + (major, minor, tile_m as usize, witness) + }; + // D18 #12. `cc_major` says what the DEVICE is; `arch_witness` says what + // this BINARY contains. When they disagree the dispatch below still takes + // the SM90 branch, into a kernel whose body was compiled out -- it writes + // nothing, `cudaGetLastError` returns success, and the caller gets the + // `alloc_zeros` buffer as a valid all-zero MoE layer. Refuse instead. + if cc_major >= 9 && arch_witness < 900 { + candle_core::bail!( + "qtip2b grouped gemm CUDA: running on sm_{cc_major}{cc_minor}, but this \ + binary carries no SM90 device code (arch witness reports {arch_witness}) \ + -- the warpgroup kernel body was compiled out, so launching it would \ + have produced an all-zero MoE layer with no error. Rebuild with \ + ARC_CUDA_ARCHS including 90 (see the arc-target crate)." + ); + } + // Mechanical gate against the two definitions drifting apart: a mismatch + // here would mis-bin the tile map and produce wrong numbers, not an error. + if tile_m != grouped_tile_m_for_cc(cc_major) { + candle_core::bail!( + "qtip2b grouped gemm CUDA: tile schedule disagreement on sm_{cc_major}{cc_minor} \ + -- kernel says tile_m={tile_m}, host mirror says {}", + grouped_tile_m_for_cc(cc_major) + ); + } + + let max_m_tiles = grouped_max_m_tiles(n_pairs, num_experts, tile_m); // Routing scratch. `counts` and `cursors` MUST be zeroed (histogram / // scatter accumulate into them); the rest is written before being read. @@ -1970,7 +2011,7 @@ pub(crate) fn grouped_gemm_2b_cuda( // Route: histogram -> scans + ragged tile map -> grouped scatter. All // on the stream; nothing comes back to the host. - unsafe { + let route_rc = unsafe { ffi::launch_qtip2b_moe_route( idx_ptr as *const _, counts_ptr as *mut _, @@ -1983,8 +2024,14 @@ pub(crate) fn grouped_gemm_2b_cuda( sp_ptr as *mut _, n_pairs as i32, num_experts as i32, - GROUPED_TILE_M as i32, + tile_m as i32, dev.cuda_stream().cu_stream(), + ) + }; + if route_rc != 0 { + candle_core::bail!( + "qtip2b grouped gemm CUDA: MoE routing launch failed (rc={route_rc}, \ + n_pairs={n_pairs}, num_experts={num_experts}, tile_m={tile_m})" ); } @@ -1996,7 +2043,7 @@ pub(crate) fn grouped_gemm_2b_cuda( // routing (out-of-range router ids) are never written. let out_buf = dev.alloc_zeros::<$T>(n_pairs * n_rows)?; let (out_ptr, out_guard) = slice_ptr(&out_buf, 0); - unsafe { + let rc = unsafe { $launch( blocks_ptr as *const _, scales_ptr as *const _, @@ -2013,9 +2060,18 @@ pub(crate) fn grouped_gemm_2b_cuda( max_m_tiles as i32, mcg_mult, dev.cuda_stream().cu_stream(), + ) + }; + drop(out_guard); + // Without this the caller receives the `alloc_zeros` buffer as a + // shape-correct, error-free, all-zero MoE output (D18). + if rc != 0 { + candle_core::bail!( + "qtip2b grouped gemm CUDA: kernel launch failed (rc={rc}) on \ + sm_{cc_major}{cc_minor} [n_pairs={n_pairs}, n_rows={n_rows}, \ + num_symbols={num_symbols}, tile_m={tile_m}]" ); } - drop(out_guard); CudaStorage::wrap_cuda_slice(out_buf, dev.clone()) }}; } diff --git a/mistralrs-quant/src/qtip/ffi.rs b/mistralrs-quant/src/qtip/ffi.rs index b873fe557..bfc133bae 100644 --- a/mistralrs-quant/src/qtip/ffi.rs +++ b/mistralrs-quant/src/qtip/ffi.rs @@ -520,11 +520,30 @@ extern "C" { // ----- Grouped GEMM (batched MoE prefill; kernels/qtip/qtip_grouped_gemm.cu) ----- // + // Which m-tile schedule this device's compute capability selects + // (`grouped::grouped_tile_m_for_cc` mirrors the rule). Returns a + // cudaError_t as i32; on failure the out-params are untouched and the + // caller MUST NOT substitute a default -- route and GEMM disagreeing on + // `tile_m` mis-bins the tile map and yields wrong numbers, not an error. + // + // `out_arch_witness` is what the running BINARY contains, which is a + // different question from what the device is: 900 means the SM90 warpgroup + // kernel body survived compilation, anything lower means it was compiled + // out and a "successful" launch of it would write nothing at all (D18 + // #12). It is measured by actually launching a witness kernel carrying the + // same `__CUDA_ARCH__` guard, once per process -- not inferred. + pub(crate) fn qtip2b_grouped_query_schedule( + out_cc_major: *mut i32, + out_cc_minor: *mut i32, + out_tile_m: *mut i32, + out_arch_witness: *mut i32, + ) -> i32; + // On-device routing (histogram + scans/tile-map + grouped scatter) for // the trellis grouped GEMM. ZERO host syncs: `d_num_tiles` stays on the // device and the GEMM grid is sized from the host-side upper bound - // `max_m_tiles = ceil(n_pairs / TILE_M) + num_experts`. `d_counts` and - // `d_cursors` must be zero-initialized. + // `max_m_tiles = ceil(n_pairs / tile_m) + num_experts`. `d_counts` and + // `d_cursors` must be zero-initialized. Returns a cudaError_t as i32. pub(crate) fn launch_qtip2b_moe_route( d_indices: *const u32, d_counts: *mut u32, @@ -539,7 +558,7 @@ extern "C" { num_experts: i32, tile_m: i32, stream: candle_core::cuda::cudarc::driver::sys::CUstream, - ); + ) -> i32; // W2A16 trellis grouped GEMM over expert-sorted (token, slot) pairs. // `d_x_rotated` must already be in the QTIP-rotated frame; `d_y` @@ -563,7 +582,7 @@ extern "C" { max_m_tiles: i32, mult: u32, stream: candle_core::cuda::cudarc::driver::sys::CUstream, - ); + ) -> i32; pub(crate) fn launch_qtip2b_grouped_gemm_f16( d_packed: *const u8, @@ -581,5 +600,5 @@ extern "C" { max_m_tiles: i32, mult: u32, stream: candle_core::cuda::cudarc::driver::sys::CUstream, - ); + ) -> i32; } diff --git a/mistralrs-quant/src/qtip/grouped.rs b/mistralrs-quant/src/qtip/grouped.rs index f0ab306db..e1b588b4b 100644 --- a/mistralrs-quant/src/qtip/grouped.rs +++ b/mistralrs-quant/src/qtip/grouped.rs @@ -38,13 +38,63 @@ pub const GROUPED_TILE_N: usize = 64; /// entirely (DOCTRINE D12). pub const GROUPED_TILE_K: usize = 64; +/// Pairs per m-tile on the SM90 (Hopper) and SM100 (Blackwell) paths. +/// +/// `wgmma.mma_async` and `tcgen05.mma` both fix M at 64, so the arch-native +/// kernels schedule 64-pair m-tiles instead of the Ampere path's 16. Wider +/// m-tiles are close to free here: trellis decode cost is per `(row, k)` and +/// independent of M, and the multiply is not the bottleneck — so a partly +/// empty m-tile wastes tensor-core flops we are not issue-limited on. +pub const GROUPED_TILE_M_WG: usize = 64; + +/// Symbols decoded per thread as one contiguous run on the arch-native paths. +/// +/// The run is seeded once with the (~11 ALU op) window identity and then +/// stepped with the (~2 op) trellis recurrence, so state reconstruction costs +/// `(11 + 2*(R-1)) / R` ops per weight — 11 at R=1 (the Ampere path's +/// per-position decode), ~2.3 at R=32. +pub const GROUPED_DECODE_RUN: usize = 32; + +/// Threads per CTA on the arch-native paths (one warpgroup). Mirrors +/// `QGW_THREADS`. +pub const GROUPED_WG_THREADS: usize = 128; + +// The run length is not free to choose: the kernel gives every thread exactly +// one run, so the weight tile must divide evenly across the warpgroup. If a +// tile constant is ever retuned without the others, this stops the build +// rather than letting the kernel decode a partial tile (D18: a mechanical gate +// beats another fix). +const _: () = assert!( + GROUPED_TILE_N * GROUPED_TILE_K == GROUPED_DECODE_RUN * GROUPED_WG_THREADS, + "GROUPED_DECODE_RUN must equal (tile_n * tile_k) / threads for the SM90+ kernel" +); + +/// Which grouped-GEMM tile schedule a device's compute capability selects. +/// +/// Mirrors `qg_tile_m_for_arch` in `kernels/qtip/qtip_grouped_gemm.cu`. The +/// host needs this before launching because the routing kernel bins pairs into +/// m-tiles and the grid bound is sized from it — route and GEMM must agree on +/// `tile_m` or the tile map is wrong. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +pub(crate) fn grouped_tile_m_for_cc(cc_major: i32) -> usize { + if cc_major >= 9 { + GROUPED_TILE_M_WG + } else { + GROUPED_TILE_M + } +} + /// Host-side upper bound on the grouped GEMM's m-tile count: every expert /// group contributes at most one partial tile beyond its full tiles, so -/// `ceil(P / TILE_M) + E` bounds the ragged total. The kernel grid is sized +/// `ceil(P / tile_m) + E` bounds the ragged total. The kernel grid is sized /// from this bound; the exact count stays on the device (zero host syncs). +/// +/// `tile_m` is arch-dependent ([`grouped_tile_m_for_cc`]) — passing the wrong +/// one under-sizes the grid, so the caller must use the same value it hands +/// the routing kernel. #[cfg_attr(not(feature = "cuda"), allow(dead_code))] // CUDA host wrapper + tests -pub(crate) fn grouped_max_m_tiles(n_pairs: usize, num_experts: usize) -> usize { - n_pairs.div_ceil(GROUPED_TILE_M) + num_experts +pub(crate) fn grouped_max_m_tiles(n_pairs: usize, num_experts: usize, tile_m: usize) -> usize { + n_pairs.div_ceil(tile_m) + num_experts } // --------------------------------------------------------------------------- @@ -163,6 +213,42 @@ pub(crate) fn window_state_2b(packed: &[u8], t: usize) -> u32 { ((r & 0x5555) << 1) | ((r >> 1) & 0x5555) } +/// Host mirror of the kernel's **sequential run** decode +/// (`q2b_advance_state` in `qtip2b_common.cuh`). +/// +/// `window_state_2b` reconstructs an isolated position in ~11 ALU ops. Along a +/// contiguous run the original recurrence costs two: +/// +/// ```text +/// state(t+1) = ((state(t) << 2) | sym[t+1]) & 0xFFFF +/// ``` +/// +/// so the SM90/SM100 grouped kernels seed one window per run of +/// [`GROUPED_DECODE_RUN`] symbols and then step. This mirror exists so the +/// identity `run_states_2b(p, t0, r)[i] == window_state_2b(p, t0 + i)` is +/// property-tested on CPU (D14: a unit test, not a kernel measurement) — if it +/// ever breaks, the GPU kernel decodes the wrong weights. +/// +/// `t0` is the first symbol of the run; `len` its length. Panics if the run +/// leaves `packed`. +#[allow(dead_code)] // kernel mirror; exercised by the colocated tests +pub(crate) fn run_states_2b(packed: &[u8], t0: usize, len: usize) -> Vec { + let mut out = Vec::with_capacity(len); + if len == 0 { + return out; + } + // Seed the run with the (expensive) window identity, exactly as the kernel + // does — this is what keeps runs independent of one another. + let mut state = window_state_2b(packed, t0); + out.push(state); + for t in t0 + 1..t0 + len { + let sym = ((packed[t / 4] >> (2 * (t % 4))) & 0x3) as u32; + state = ((state << 2) | sym) & 0xFFFF; + out.push(state); + } + out +} + #[cfg(test)] mod tests { use super::*; @@ -186,7 +272,14 @@ mod tests { /// into TILE_M-rounded units and respect the ceil(P/TILE_M)+E bound. #[test] fn tile_map_partitions_ragged_groups() { - let tile_m = GROUPED_TILE_M; + // Both schedules: the Ampere path's 16-pair m-tile and the + // wgmma/tcgen05 paths' 64-pair m-tile. + for tile_m in [GROUPED_TILE_M, GROUPED_TILE_M_WG] { + tile_map_partitions_ragged_groups_at(tile_m); + } + } + + fn tile_map_partitions_ragged_groups_at(tile_m: usize) { let cases: Vec> = vec![ // 1 giant expert + many tiny. vec![1000, 1, 2, 3, 0, 1, 5, 7], @@ -215,7 +308,7 @@ mod tests { // The kernel-side upper bound holds. assert!( - tiles.len() <= grouped_max_m_tiles(n_pairs as usize, counts.len()), + tiles.len() <= grouped_max_m_tiles(n_pairs as usize, counts.len(), tile_m), "tile count {} exceeds bound for counts {counts:?}", tiles.len() ); @@ -279,4 +372,105 @@ mod tests { } } } + + /// Deterministic pseudo-random 2-bit symbol stream, packed LSB-first at + /// 4 symbols/byte (the production layout). + fn packed_stream(len: usize, seed: u64) -> Vec { + let mut packed = vec![0u8; len.div_ceil(4)]; + let mut z = seed; + for t in 0..len { + z = z + .wrapping_mul(0x9E37_79B9_7F4A_7C15) + .wrapping_add(0xBF58_476D_1CE4_E5B9); + let s = ((z >> 33) & 0x3) as u8; + packed[t / 4] |= s << (2 * (t % 4)); + } + packed + } + + /// THE load-bearing identity of the SM90/SM100 kernels: seeding one + /// window and then stepping the trellis recurrence must reproduce, at + /// every position of the run, exactly what a per-position window + /// reconstruction would have produced. + /// + /// A run that started from the wrong seed, drifted, or consumed symbols + /// off-by-one would decode different weights — silently, since every + /// state maps to *some* finite codeword. This test is the tripwire. + #[test] + fn run_decode_matches_per_position_window() { + for (len, seed) in [(64usize, 11u64), (256, 12), (1024, 13)] { + let packed = packed_stream(len, seed); + // Every run start that the kernel can produce (chunk-local 0 and + // GROUPED_DECODE_RUN), plus deliberately unaligned starts. + for &run in &[1usize, 2, 8, GROUPED_DECODE_RUN] { + let mut t0 = 0usize; + while t0 + run <= len { + let states = run_states_2b(&packed, t0, run); + assert_eq!(states.len(), run); + for (i, &s) in states.iter().enumerate() { + assert_eq!( + s, + window_state_2b(&packed, t0 + i), + "run(t0={t0}, run={run}) diverged at i={i} (len={len}, seed={seed})" + ); + } + // Step by a non-divisor so unaligned starts are covered. + t0 += if run == 1 { 7 } else { run }; + } + } + } + } + + /// The run decode must also be correct across the warm-up region + /// (`t < 7`), where the 16-bit window reads zeros below stream position 0. + #[test] + fn run_decode_covers_warmup_region() { + let packed = packed_stream(64, 99); + let states = run_states_2b(&packed, 0, GROUPED_DECODE_RUN); + // Sequential recurrence from the all-zero initial state. + let mut state = 0u32; + for (t, s) in states.iter().enumerate().take(GROUPED_DECODE_RUN) { + let sym = ((packed[t / 4] >> (2 * (t % 4))) & 0x3) as u32; + state = ((state << 2) | sym) & 0xFFFF; + assert_eq!(*s, state, "warm-up divergence at t={t}"); + } + } + + /// The arch tile-schedule selector must be total and must never hand the + /// wgmma schedule to a pre-Hopper device (its m-tile would not match the + /// tile map the routing kernel built). + #[test] + fn tile_schedule_selects_by_compute_capability() { + for cc in [7, 8] { + assert_eq!(grouped_tile_m_for_cc(cc), GROUPED_TILE_M); + } + for cc in [9, 10, 11, 12] { + assert_eq!(grouped_tile_m_for_cc(cc), GROUPED_TILE_M_WG); + } + } + + /// The grid bound must dominate the tile map for whichever `tile_m` the + /// arch selected — an under-sized grid silently drops tiles (D18). + #[test] + fn grid_bound_dominates_tile_map_for_every_schedule() { + let cases: Vec> = vec![ + vec![1000, 1, 2, 3, 0, 1, 5, 7], + vec![0, 0, 640, 0], + vec![0, 16, 0, 17, 0, 15, 0], + vec![1], + vec![0, 0, 0], + vec![64, 128, 65], + ]; + for tile_m in [GROUPED_TILE_M, GROUPED_TILE_M_WG] { + for counts in &cases { + let n_pairs: u32 = counts.iter().sum(); + let (_, tiles) = build_group_tile_map(counts, tile_m); + assert!( + tiles.len() <= grouped_max_m_tiles(n_pairs as usize, counts.len(), tile_m), + "bound violated: tile_m={tile_m} counts={counts:?} tiles={}", + tiles.len() + ); + } + } + } } diff --git a/mistralrs-quant/src/qtip/mod.rs b/mistralrs-quant/src/qtip/mod.rs index 01548d4f3..b2a1e2a10 100644 --- a/mistralrs-quant/src/qtip/mod.rs +++ b/mistralrs-quant/src/qtip/mod.rs @@ -89,6 +89,35 @@ pub use grouped::{ ExpertBpwTable, TrellisBpw, GROUPED_TILE_K as QTIP_GROUPED_TILE_K, GROUPED_TILE_M as QTIP_GROUPED_TILE_M, GROUPED_TILE_N as QTIP_GROUPED_TILE_N, }; + +/// The grouped GEMM's m-tile on **this** device. +/// +/// D16 made the schedule arch-dependent (16 pairs per m-tile on Ampere, 64 on +/// SM90+), so any harness that models the kernel's weight traffic has to ask +/// rather than assume: the amortization a benchmark may legitimately claim is +/// bounded by this number, and using the wrong one silently mis-states the +/// grouped path's advantage in whichever direction the guess went. +/// Reports the arch witness too, because a harness that reads `tile_m` off a +/// binary with no SM90 device code would be modelling a kernel that cannot +/// run on this device (D18 #12). +#[cfg(feature = "cuda")] +pub fn qtip_grouped_tile_m() -> candle_core::Result { + let (mut major, mut minor, mut tile_m, mut witness) = (0i32, 0i32, 0i32, 0i32); + let rc = unsafe { + ffi::qtip2b_grouped_query_schedule(&mut major, &mut minor, &mut tile_m, &mut witness) + }; + if rc != 0 { + candle_core::bail!("qtip grouped tile query failed (rc={rc})"); + } + if major >= 9 && witness < 900 { + candle_core::bail!( + "qtip grouped tile query: device is sm_{major}{minor} but this binary carries \ + no SM90 device code (arch witness {witness}); the SM90 m-tile it reports \ + could never be the one that runs. Rebuild with ARC_CUDA_ARCHS including 90." + ); + } + Ok(tile_m as usize) +} #[allow(unused_imports)] pub use viterbi::{ hessian_row_weights, quantize_row, viterbi_quantize_row, TrellisSearch, HESSIAN_SIGMA_REG,