From 96cbce843be00f8dfc179183c7dd915a4b2ab76e Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 4 Oct 2026 18:09:05 +0900 Subject: [PATCH 1/7] update(rocm): port ssm_update_kernel to HIP and read the port table On ROCm every single-token Mamba2 decode step of the hybrid SSM models (granite-4.0-h, falcon-h1, plamo-2, Nemotron-H) ran the ~55-op SSD graph, because ssm_ports() had no .rocm entry and ssm_kernel_available() answered with cu::is_available() off Apple, so a filled slot would have stayed unreachable. - ssm_ports().rocm: SSM_HIP_SOURCE, the CUDA body with __shfl_down(acc, o, 32) and the wave32 preprocessor guard, launched through fast::hip_kernel with the same inputs, outputs, grid and template arguments. - ssm_kernel_available() returns has_kernel_port(ssm_ports()) on every platform. MLXCEL_SSM_KERNEL=0 forces the graph path on every backend; MLXCEL_SSM_CUDA_KERNEL=0 stays as an alias. - On ROCm the launch also keys the hipRTC module on the A_log, B and C dtypes. Nemotron-H stores A_log in f32 next to bf16 activations, granite in bf16; Metal and CUDA kernel names are unchanged. - ssm_update_parity_tests compares the kernel with an f32 MLX-op reference of the single-token SSD step at the granite-4.0-h-tiny and Nemotron-H shapes, f32 and bf16, plus the A_log dtype key and both kill switches. Each test was checked to fail with the fold starting at 8 or the dtype keys removed. Refs #2067 --- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 4 +- src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp | 132 +++++- src/lib/mlxcel-core/src/lib.rs | 16 +- .../src/ssm_update_parity_tests.rs | 375 ++++++++++++++++++ 4 files changed, 511 insertions(+), 16 deletions(-) create mode 100644 src/lib/mlxcel-core/src/ssm_update_parity_tests.rs diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index 292977790..f6770673f 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1957,7 +1957,9 @@ std::unique_ptr fused_moe_forward( int32_t bits ); -// Check if SSM Metal kernel is available (Metal GPU only) +// Whether the fused single-token SSM update kernel can run: the GPU backend +// has a port in ssm_ports() (Metal, CUDA, ROCm) and MLXCEL_SSM_KERNEL=0 (or +// its older alias MLXCEL_SSM_CUDA_KERNEL=0) is not set. bool ssm_kernel_available(); // Mamba1 selective scan fused over the sequence (Jamba, issue #2005). diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp index 845c554ba..2064b76f4 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp @@ -633,6 +633,99 @@ namespace { return holder; } + // ROCm port (issue #2067). Same computation, thread mapping, inputs, + // outputs, grid and template arguments as the CUDA port above; the body + // differs in the lane reduction only. `__shfl_down_sync` is a HIP + // compatibility shim that ignores its mask, so the native + // `__shfl_down(var, delta, width)` is used with the width stated. The fold + // starts at 16 because one 32-lane wavefront owns one (head, row) pair: + // the launch's threadgroup is (32, 8, 1), so each `threadIdx.y` row is one + // wave on a wave32 target. On a wave64 target (CDNA: gfx90a, gfx942) a + // 64-lane wave would span two rows and this fold would mix them, which is + // why the source refuses to compile there. + static const char* SSM_HIP_SOURCE = R"( + // A fold that starts at 16 is correct only for a 32-lane wavefront; + // on wave64 it would return a finite, plausible, wrong sum. Checked + // with the preprocessor because `static_assert(warpSize == 32)` does + // not compile in HIP (`warpSize` is not a constant expression). Both + // macro spellings are checked; the trailing-underscore one is older. + #if defined(__AMDGCN_WAVEFRONT_SIZE__) && __AMDGCN_WAVEFRONT_SIZE__ != 32 + #error "ssm_kernel_hip assumes a 32-lane wavefront" + #endif + #if defined(__AMDGCN_WAVEFRONT_SIZE) && __AMDGCN_WAVEFRONT_SIZE != 32 + #error "ssm_kernel_hip assumes a 32-lane wavefront" + #endif + uint32_t n = blockIdx.z; // batch * heads + uint32_t h_idx = n % H; + uint32_t g_idx = n / G; + constexpr int n_per_t = Ds / 32; + + // Wave-uniform: all 32 lanes of a wave share threadIdx.y. + uint32_t d_idx = blockIdx.y * blockDim.y + threadIdx.y; + if (d_idx >= (uint32_t)Dh) return; + + auto x = X + n * Dh; + out += n * Dh; + auto i_state = state_in + n * Dh * Ds; + auto o_state = state_out + n * Dh * Ds; + + auto C_ = C + g_idx * Ds; + auto B_ = B + g_idx * Ds; + + uint32_t ds_idx = threadIdx.x; // lane 0..31 + + float dt_ = (float)dt[n]; + float A = -expf((float)A_log[h_idx]); + float dA = expf(A * dt_); + + float acc = 0.0f; + float x_ = (float)x[d_idx]; + + for (int i = 0; i < n_per_t; ++i) { + int s_idx = n_per_t * (int)ds_idx + i; + int idx = (int)d_idx * Ds + s_idx; + float dB_by_x = x_ * dt_ * (float)B_[s_idx]; + float state = dA * (float)i_state[idx] + dB_by_x; + o_state[idx] = (U)state; + acc += state * (float)C_[s_idx]; + } + #pragma unroll + for (int o = 16; o > 0; o >>= 1) { + acc += __shfl_down(acc, o, 32); + } + if (ds_idx == 0u) { + out[d_idx] = (T)(acc + x_ * (float)D[h_idx]); + } + )"; + + struct SsmKernelHolderHip { + std::optional kernel; + bool initialized = false; + + mlx::core::fast::CustomKernelFunction& get() { + if (!initialized) { +#ifdef MLXCEL_BRIDGE_ROCM_BACKEND + kernel = mlx::core::fast::hip_kernel( + "ssm_kernel_hip", + {"X", "A_log", "B", "C", "D", "dt", "state_in"}, + {"out", "state_out"}, + SSM_HIP_SOURCE + ); +#else + throw std::runtime_error( + "[ssm_update_kernel] this build has no ROCm backend"); +#endif + initialized = true; + } + return *kernel; + } + }; + + static SsmKernelHolderHip& get_ssm_kernel_hip() { + static SsmKernelHolderHip holder; + return holder; + } + // This kernel's ports, in one place. `has_kernel_port` and `select_kernel_port` // both read it, so a support predicate and the dispatch cannot answer // differently (#1801). @@ -644,8 +737,9 @@ const mlxcel::KernelPorts& ssm_ports() { .cuda = +[]() -> mlx::core::fast::CustomKernelFunction& { return get_ssm_kernel_cuda().get(); }, - // No HIP port yet (#1814). - .rocm = nullptr, + .rocm = +[]() -> mlx::core::fast::CustomKernelFunction& { + return get_ssm_kernel_hip().get(); + }, }; return ports; } @@ -1574,19 +1668,21 @@ std::unique_ptr fused_moe_forward( } bool ssm_kernel_available() { -#ifdef __APPLE__ - return mlx::core::metal::is_available(); -#else - // CUDA port of the fused SSM decode kernel (#631); previously the hybrid - // SSM models fell back to the ~55-op graph path on every decode step. - // MLXCEL_SSM_CUDA_KERNEL=0 forces that graph path (debug / A-B bench). - if (const char* e = std::getenv("MLXCEL_SSM_CUDA_KERNEL")) { - if (e[0] == '0' && e[1] == '\0') { + // MLXCEL_SSM_KERNEL=0 forces the ~55-op ssm_step graph path on every + // backend (debug / A-B bench). MLXCEL_SSM_CUDA_KERNEL=0 is the name the + // switch had while the CUDA port (#631) was the only one off Apple; it is + // kept as an alias so existing scripts still work (#2067). + for (const char* name : {"MLXCEL_SSM_KERNEL", "MLXCEL_SSM_CUDA_KERNEL"}) { + const char* e = std::getenv(name); + if (e != nullptr && e[0] == '0' && e[1] == '\0') { return false; } } - return mlx::core::cu::is_available(); -#endif + // Read from the same table the dispatch in ssm_update_kernel selects + // from, so the Rust gates (`seq_len == 1 && ssm_kernel_available()`) can + // `expect` the launch: a true here means select_kernel_port will not + // refuse. + return mlxcel::has_kernel_port(ssm_ports()); } void ssm_update_kernel( @@ -1636,6 +1732,18 @@ void ssm_update_kernel( {"H", h}, {"G", g}, }; + // The hipRTC module cache is keyed on the name plus this list, while the + // generated signature takes every input's runtime dtype, so an input whose + // dtype is not named here could reuse a module compiled for another dtype + // (scripts/ci/check_kernel_dtype_keys.py). T and U cover X, D (cast to T + // below) and state_in, and dt is always float32 from compute_dt; A_log, B + // and C are not tied to T, so the HIP port names them too. Appended on + // ROCm only, which leaves the Metal and CUDA kernel names unchanged. + if (mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Rocm) { + template_args.push_back({"TA", A_log.inner.dtype()}); + template_args.push_back({"TB", B.inner.dtype()}); + template_args.push_back({"TC", C.inner.dtype()}); + } std::vector inputs = { hidden_states.inner, A_log.inner, B.inner, C.inner, diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index daa18b0cd..bb17baf15 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -1840,12 +1840,15 @@ mod ffi { ) -> Result>; // SSM (Mamba2) fused Metal kernel. - /// Check if SSM Metal kernel is available + /// Whether the fused single-token SSM update kernel can run: the + /// resolved GPU backend has a port in `ssm_ports()` (Metal, CUDA, + /// ROCm) and neither `MLXCEL_SSM_KERNEL=0` nor its older alias + /// `MLXCEL_SSM_CUDA_KERNEL=0` is set. fn ssm_kernel_available() -> bool; /// Fused SSM update kernel for single-token decode - /// Replaces ~55 individual ops with a single Metal kernel call - /// Used by: NemotronH, NemotronNAS, Mamba2 + /// Replaces ~55 individual ops with a single kernel launch + /// Used by: GraniteMoeHybrid, Falcon-H1, PLaMo 2, NemotronH fn ssm_update_kernel( hidden_states: &MlxArray, a_log: &MlxArray, @@ -3970,6 +3973,13 @@ mod fused_norm_parity_tests; #[path = "mamba1_scan_parity_tests.rs"] mod mamba1_scan_parity_tests; +// Kernel-against-graph parity for the fused single-token Mamba2 SSM update +// (#2067) at the granite-4.0-h-tiny and Nemotron-H shapes, f32 and bf16, plus +// the `MLXCEL_SSM_KERNEL` kill switch. GPU-only; they skip on CPU-only builds. +#[cfg(test)] +#[path = "ssm_update_parity_tests.rs"] +mod ssm_update_parity_tests; + // Numeric regression tests for `fast::rms_norm` on the small-axis CUDA dispatch // band (#830/#831): the deleted pre-#3792 overlay read past its shared scratch // for 16-bit axes in (256, 512] (f32: (128, 256]), which is exactly the diff --git a/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs b/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs new file mode 100644 index 000000000..1136d4c8d --- /dev/null +++ b/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs @@ -0,0 +1,375 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Parity tests for the fused single-token Mamba2 SSM update kernel (#2067). +//! +//! `ssm_update_kernel` replaces the ~55-op SSD graph (`ssm_step` in +//! `src/models/granitemoehybrid.rs` and its falcon-h1, plamo-2 and Nemotron-H +//! siblings) for every single-token decode step that has a state. For one +//! token that graph reduces to the recurrence +//! +//! ```text +//! dt' = clip(softplus(dt + dt_bias), t_min, t_max) +//! s' = exp(-exp(A_log) * dt') * s + dt' * x * B +//! y = s' . C + D * x +//! ``` +//! +//! with B and C shared by the heads of a group. The reference below composes +//! exactly that from MLX ops in float32 (the graph path also promotes to +//! float32), and the kernel's output and new state are compared with it at the +//! granite-4.0-h-tiny and Nemotron-H mixer shapes, in f32 and bf16. +//! +//! Tolerances are normalized RMS and normalized max deviation (both divided by +//! the reference's RMS): f32 1e-5 / 1e-4 and bf16 1.6e-2 / 7e-2, the bf16 +//! budget of `fused_norm_parity_tests.rs`. The state is carried in float32 in +//! every case, as the models carry it, so it is held to the f32 budget even +//! when the activations are bf16. +//! +//! These tests return early only on a build with no GPU backend, or when the +//! `MLXCEL_SSM_KERNEL=0` / `MLXCEL_SSM_CUDA_KERNEL=0` kill switch is set. On +//! Metal, CUDA and ROCm a false `ssm_kernel_available()` is itself a defect +//! (each has a port), so the test fails there instead of skipping. +//! +//! Run on ROCm: +//! cargo test --release --features rocm -p mlxcel-core --lib \ +//! ssm_update_parity_tests -- --test-threads=1 + +use super::*; +use crate::hardware::{GpuBackendKind, gpu_backend_kind}; + +const KILL_SWITCHES: [&str; 2] = ["MLXCEL_SSM_KERNEL", "MLXCEL_SSM_CUDA_KERNEL"]; + +/// Normalized (RMS, max) deviation budget for values stored in `dtype`. +fn tolerance_for(dtype: i32) -> (f64, f64) { + if dtype == dtype::FLOAT32 { + (1e-5, 1e-4) + } else { + (1.6e-2, 7e-2) + } +} + +fn kill_switch_set() -> bool { + KILL_SWITCHES + .iter() + .any(|name| std::env::var(name).is_ok_and(|v| v == "0")) +} + +/// True when the calling test should return early, after saying why. +fn skip() -> bool { + if ssm_kernel_available() { + return false; + } + let backend = gpu_backend_kind(); + if backend == GpuBackendKind::None { + eprintln!("skipping ssm_update_parity_tests: no GPU backend in this build"); + return true; + } + if kill_switch_set() { + eprintln!("skipping ssm_update_parity_tests: SSM kernel kill switch is set"); + return true; + } + panic!( + "ssm_kernel_available() is false on {backend:?}, which has a port in ssm_ports(); \ + the predicate and the port table disagree" + ); +} + +fn flatten_f32(arr: &MlxArray) -> Vec { + let a = astype(arr, dtype::FLOAT32); + eval(&a); + array_to_raw_bytes(&a) + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect() +} + +/// RMS and max of `a - b`, both divided by the RMS of `b`. +fn normalized_deviation(a: &[f32], b: &[f32]) -> (f64, f64) { + assert_eq!(a.len(), b.len(), "length mismatch in deviation check"); + let mut diff_sq = 0f64; + let mut ref_sq = 0f64; + let mut max_abs = 0f64; + for (x, y) in a.iter().zip(b.iter()) { + assert!(x.is_finite(), "non-finite kernel output {x}"); + let d = f64::from(*x) - f64::from(*y); + diff_sq += d * d; + ref_sq += f64::from(*y) * f64::from(*y); + max_abs = max_abs.max(d.abs()); + } + let ref_rms = (ref_sq / b.len() as f64).sqrt().max(1e-20); + ( + (diff_sq / a.len() as f64).sqrt() / ref_rms, + max_abs / ref_rms, + ) +} + +fn assert_within(label: &str, got: &MlxArray, want: &MlxArray, dtype: i32) { + let (nrms, nmax) = normalized_deviation(&flatten_f32(got), &flatten_f32(want)); + let (rms_tol, max_tol) = tolerance_for(dtype); + assert!( + nrms < rms_tol && nmax < max_tol, + "{label}: normalized rms {nrms:.3e} (tol {rms_tol:.1e}), \ + normalized max {nmax:.3e} (tol {max_tol:.1e})" + ); +} + +/// One mixer's single-token step. +struct Shape { + name: &'static str, + batch: i32, + heads: i32, + head_dim: i32, + groups: i32, + state_dim: i32, + /// `time_step_limit` of the model. + dt_limits: (f32, f32), +} + +/// granite-4.0-h-tiny: 48 heads of 64, one group, state 128, no dt clip. +const GRANITE_TINY: Shape = Shape { + name: "granite-4.0-h-tiny", + batch: 1, + heads: 48, + head_dim: 64, + groups: 1, + state_dim: 128, + dt_limits: (0.0, f32::INFINITY), +}; + +/// Nemotron-3-Nano-30B-A3B: 64 heads of 64 in 8 groups, state 128. Run at +/// batch 2 with a clipping dt window so the group/batch indexing and the clip +/// are both exercised. +const NEMOTRON_H: Shape = Shape { + name: "nemotron-h", + batch: 2, + heads: 64, + head_dim: 64, + groups: 8, + state_dim: 128, + dt_limits: (1e-3, 0.1), +}; + +struct Case { + x: UniquePtr, + a_log: UniquePtr, + b: UniquePtr, + c: UniquePtr, + d: UniquePtr, + dt: UniquePtr, + dt_bias: UniquePtr, + state: UniquePtr, +} + +fn normal(shape: &[i32]) -> UniquePtr { + unsafe { random_normal(shape, dtype::FLOAT32, std::ptr::null()) } +} + +fn uniform(low: f32, high: f32, shape: &[i32]) -> UniquePtr { + unsafe { random_uniform(low, high, shape, dtype::FLOAT32, std::ptr::null()) } +} + +/// Random single-token inputs. Activations and the per-head parameters are in +/// `act_dtype`, except `A_log`, which is in `a_log_dtype` (Nemotron-H stores it +/// in f32 next to bf16 activations, granite in bf16). The state is f32. +/// Everything is evaluated up front so the kernel and the reference read the +/// same bytes. +fn make_case(s: &Shape, seed: u64, act_dtype: i32, a_log_dtype: i32) -> Case { + random_seed(seed); + let cast = |a: UniquePtr, dt: i32| { + let out = astype(&a, dt); + eval(&out); + out + }; + // A = exp(A_log) in [1, 16] and dt' around softplus(-2) ~ 0.13 keep the + // decay exp(-A dt') well inside (0, 1), so the carried state contributes. + let a_log = cast(log(&uniform(1.0, 16.0, &[s.heads])), a_log_dtype); + let dt_bias = cast(uniform(-3.0, -1.0, &[s.heads]), act_dtype); + Case { + x: cast(normal(&[s.batch, 1, s.heads, s.head_dim]), act_dtype), + a_log, + b: cast(normal(&[s.batch, 1, s.groups, s.state_dim]), act_dtype), + c: cast(normal(&[s.batch, 1, s.groups, s.state_dim]), act_dtype), + d: cast(uniform(0.5, 1.5, &[s.heads]), act_dtype), + dt: cast(normal(&[s.batch, 1, s.heads]), act_dtype), + dt_bias, + state: cast( + normal(&[s.batch, s.heads, s.head_dim, s.state_dim]), + dtype::FLOAT32, + ), + } +} + +fn run_kernel(s: &Shape, k: &Case) -> (UniquePtr, UniquePtr) { + let mut out = UniquePtr::null(); + let mut state = UniquePtr::null(); + ssm_update_kernel( + &k.x, + &k.a_log, + &k.b, + &k.c, + &k.d, + &k.dt, + &k.dt_bias, + &k.state, + s.dt_limits.0, + s.dt_limits.1, + &mut out, + &mut state, + ) + .expect("ssm_kernel_available() is true, so the launcher must not refuse"); + eval(&out); + eval(&state); + (out, state) +} + +/// The single-token SSD step from MLX ops, in float32. Returns `y` as +/// `[batch, 1, heads, head_dim]` and the new state as +/// `[batch, heads, head_dim, state_dim]`. +fn reference(s: &Shape, k: &Case) -> (UniquePtr, UniquePtr) { + let f32_of = |a: &MlxArray| astype(a, dtype::FLOAT32); + let (b, h, dh, g, ds) = (s.batch, s.heads, s.head_dim, s.groups, s.state_dim); + + let lo = full_f32(&[1], s.dt_limits.0, dtype::FLOAT32); + let hi = full_f32(&[1], s.dt_limits.1, dtype::FLOAT32); + let dt = clip( + &softplus(&add(&f32_of(&k.dt), &f32_of(&k.dt_bias))), + &lo, + &hi, + ); + let dt = reshape(&dt, &[b, h, 1, 1]); + + let a = negative(&exp(&f32_of(&k.a_log))); + let decay = exp(&multiply(&dt, &reshape(&a, &[1, h, 1, 1]))); + + // B and C per head: [b, 1, g, ds] -> [b, g, ds] -> repeat heads/group -> [b, h, 1, ds]. + let per_head = |m: &MlxArray| { + let m = reshape(&f32_of(m), &[b, g, ds]); + let m = repeat(&m, h / g, 1); + reshape(&m, &[b, h, 1, ds]) + }; + let b_h = per_head(&k.b); + let c_h = per_head(&k.c); + + let x = reshape(&f32_of(&k.x), &[b, h, dh, 1]); + let dbx = multiply(&multiply(&dt, &x), &b_h); + let new_state = add(&multiply(&decay, &k.state), &dbx); + + let y = sum_axis(&multiply(&new_state, &c_h), -1, false); + let d = reshape(&f32_of(&k.d), &[1, h, 1]); + let x_flat = reshape(&x, &[b, h, dh]); + let y = add(&y, &multiply(&x_flat, &d)); + let y = reshape(&y, &[b, 1, h, dh]); + eval(&y); + eval(&new_state); + (y, new_state) +} + +fn check(s: &Shape, seed: u64, act_dtype: i32, a_log_dtype: i32) { + let case = make_case(s, seed, act_dtype, a_log_dtype); + let (out, state) = run_kernel(s, &case); + let (want_out, want_state) = reference(s, &case); + let backend = gpu_backend_kind(); + let label = format!( + "{} act dtype {act_dtype} A_log dtype {a_log_dtype} on {backend:?}", + s.name + ); + assert_eq!(array_dtype(&out), act_dtype, "{label}: output dtype"); + assert_eq!( + array_dtype(&state), + dtype::FLOAT32, + "{label}: state dtype follows state_in" + ); + assert_within(&format!("{label}: y"), &out, &want_out, act_dtype); + assert_within( + &format!("{label}: new state"), + &state, + &want_state, + dtype::FLOAT32, + ); +} + +#[test] +fn ssm_update_kernel_matches_graph_step_f32() { + let _guard = crate::test_support::env_lock::env_lock(); + if skip() { + return; + } + check(&GRANITE_TINY, 2067, dtype::FLOAT32, dtype::FLOAT32); + check(&NEMOTRON_H, 2068, dtype::FLOAT32, dtype::FLOAT32); +} + +#[test] +fn ssm_update_kernel_matches_graph_step_bf16() { + let _guard = crate::test_support::env_lock::env_lock(); + if skip() { + return; + } + check(&GRANITE_TINY, 2069, dtype::BFLOAT16, dtype::BFLOAT16); + check(&NEMOTRON_H, 2070, dtype::BFLOAT16, dtype::FLOAT32); +} + +/// The same shape and activation dtype with `A_log` first in bf16 and then in +/// f32, in one process. The kernel's generated signature takes each input's +/// runtime dtype; on a backend whose JIT cache key omits `A_log`'s dtype the +/// second launch would reuse the module compiled for the first and read f32 +/// bytes as bf16. +#[test] +fn ssm_update_kernel_keys_on_a_log_dtype() { + let _guard = crate::test_support::env_lock::env_lock(); + if skip() { + return; + } + check(&GRANITE_TINY, 2071, dtype::BFLOAT16, dtype::BFLOAT16); + check(&GRANITE_TINY, 2072, dtype::BFLOAT16, dtype::FLOAT32); +} + +/// `MLXCEL_SSM_KERNEL=0` and its older alias `MLXCEL_SSM_CUDA_KERNEL=0` each +/// turn the predicate off, on every backend with a port. +#[test] +fn ssm_kernel_kill_switches_turn_predicate_off() { + let _guard = crate::test_support::env_lock::env_lock(); + let saved: Vec<(&str, Option)> = KILL_SWITCHES + .iter() + .map(|name| (*name, std::env::var(name).ok())) + .collect(); + for name in KILL_SWITCHES { + unsafe { std::env::remove_var(name) }; + } + // Checked before restoring the environment, so a failure leaves it as the + // test found it. + let result = std::panic::catch_unwind(|| { + if skip() { + return; + } + for name in KILL_SWITCHES { + unsafe { std::env::set_var(name, "0") }; + assert!( + !ssm_kernel_available(), + "{name}=0 must force the graph path" + ); + unsafe { std::env::set_var(name, "1") }; + assert!(ssm_kernel_available(), "{name}=1 must leave the kernel on"); + unsafe { std::env::remove_var(name) }; + } + }); + for (name, value) in saved { + match value { + Some(v) => unsafe { std::env::set_var(name, v) }, + None => unsafe { std::env::remove_var(name) }, + } + } + if let Err(panic) = result { + std::panic::resume_unwind(panic); + } +} From 844bd94c2a43bfe5d10f595b14377269332b14bc Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 4 Oct 2026 19:01:14 +0900 Subject: [PATCH 2/7] fix(kernels): key the SSM launch on A_log, B and C dtypes everywhere The first commit appended the three dtype keys only when gpu_kernel_backend() was ROCm, which make verify-kernel-port-dispatch rejects as port selection by backend kind. Name them on every backend instead. The kernels never reference TA, TB or TC, so the arithmetic and the launch are unchanged; only the JIT module names differ. CUDA had the same exposure as HIP (its cache key is the name plus the template arguments), so a process that ran a bf16-A_log model and then an f32-A_log one at the same shapes could reuse the wrong module there too. Also adds the ROCm w1ctx512 and w1 logit traces for granite-4.0-h-tiny and Nemotron-H with the HIP kernel, plus MLXCEL_SSM_KERNEL=0 graph traces from the same binary. Refs #2067 --- .../rocm_gfx1151_96cbce84/METADATA.txt | 15 ++ .../rocm_gfx1151_96cbce84/README.md | 18 +++ .../rocm_gfx1151_96cbce84/RUNS.txt | 6 + .../rocm_gfx1151_96cbce84/SHA256SUMS | 6 + ...96cbce84_granite-4.0-h-tiny_default_w1.tsv | 135 +++++++++++++++++ ...84_granite-4.0-h-tiny_default_w1ctx512.tsv | 135 +++++++++++++++++ ...granite-4.0-h-tiny_ssmkernel0_w1ctx512.tsv | 135 +++++++++++++++++ ...e84_nemotron-3-nano-30b-a3b_default_w1.tsv | 140 ++++++++++++++++++ ...motron-3-nano-30b-a3b_default_w1ctx512.tsv | 140 ++++++++++++++++++ ...ron-3-nano-30b-a3b_ssmkernel0_w1ctx512.tsv | 140 ++++++++++++++++++ src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp | 23 +-- 11 files changed, 882 insertions(+), 11 deletions(-) create mode 100644 benchmarks/logit_traces/rocm_gfx1151_96cbce84/METADATA.txt create mode 100644 benchmarks/logit_traces/rocm_gfx1151_96cbce84/README.md create mode 100644 benchmarks/logit_traces/rocm_gfx1151_96cbce84/RUNS.txt create mode 100644 benchmarks/logit_traces/rocm_gfx1151_96cbce84/SHA256SUMS create mode 100644 benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1.tsv create mode 100644 benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1ctx512.tsv create mode 100644 benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_ssmkernel0_w1ctx512.tsv create mode 100644 benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1.tsv create mode 100644 benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1ctx512.tsv create mode 100644 benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_ssmkernel0_w1ctx512.tsv diff --git a/benchmarks/logit_traces/rocm_gfx1151_96cbce84/METADATA.txt b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/METADATA.txt new file mode 100644 index 000000000..2bf65c116 --- /dev/null +++ b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/METADATA.txt @@ -0,0 +1,15 @@ +host: AMD Ryzen AI MAX+ 395 with Radeon 8060S (gfx1151, RDNA 3.5), 96 GiB VRAM carve-out, 31 GiB host RAM +os: Debian GNU/Linux 13 (trixie), kernel 6.18.12+deb13-amd64 +rocm: 10.0.0 packages (HIP 7.15.26333), ROCM_PATH=/opt/rocm +mlxcel_commit: 96cbce84 (branch update/issue-2067-rocm-ssm-update-port, on origin/main c0b71344), the first commit of the #2067 PR. The PR's next commit names the A_log, B and C dtypes in the SSM launch's template arguments on every backend instead of on ROCm only; on ROCm the launch, its arguments and its cache key are the same in both commits, so these traces stand for the PR head as well. +mlx_pin: 81ba1c6a0e50a9268b931579c2d4f1158b9aab5a +rocm_overlay_fork: NripeshN/mlx rocm-support @ 75915908dfe5028335d318b10340313744fd3a8d (src/lib/mlx-cpp/patches-rocm/UPSTREAM), plus the local fixes in LOCAL_FIXES.md +logit_trace_sha256: 4c1ecae4dd823f08c6e9aa997cbfcfe535a3aef3a9c85dca63669327a6ee078e +build: cargo build --release --features rocm --example logit_trace +corpus: tests/fixtures/wikitext2_excerpt.txt sha256 16721fc1057d69efc8493718fcad562de9e4554ce05369876cd06dcba431394e +checkpoints: the same local directories as ../rocm_gfx1151_3c9edea0/ (granite-4.0-h-tiny-4bit@02c8783da0cb171942e30a7dde8acce11fe82e0d, NVIDIA-Nemotron-3-Nano-30B-A3B-4bit@832f602eba5d22436c258c1462bdedc5afddb42b) +args: w1ctx512='1 128 8 512' with MLXCEL_TRACE_START_TOKEN=512; w1='1 128 8 0' (CHUNK_TOKENS MAX_CHUNKS TOPK PREFILL) +tags: 'default' = no override, so each w1ctx512 single-token forward meets the prefilled SSM state and takes the HIP ssm_update_kernel (granite through ssm_step_kernel, Nemotron-H through fused_mamba2_forward). 'ssmkernel0' = MLXCEL_SSM_KERNEL=0, the same binary forced onto the SSD graph (ssm_step). w1 chunks have no prefill, so they never meet a state and run the graph on both tags. +env: LD_LIBRARY_PATH unset; MLXCEL_KEEP_BF16 unset; MLXCEL_FUSED_MOE unset; MLXCEL_DEVICE unset (trace header reads '# device GPU') +gpu: correctness traces, not timings; a parallel development unit may have used the GPU during a run +metal_reference: ../metal_m5_3c9edea0/ (w1ctx512); graph references: ../rocm_gfx1151_3c9edea0/ (w1ctx512), ../rocm_gfx1151_c5fe9a16/ (w1) diff --git a/benchmarks/logit_traces/rocm_gfx1151_96cbce84/README.md b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/README.md new file mode 100644 index 000000000..a1cbc6fd2 --- /dev/null +++ b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/README.md @@ -0,0 +1,18 @@ +# ROCm logit traces with the HIP SSM update kernel (lablup/mlxcel#2067) + +Teacher-forced traces from `examples/logit_trace` on the Radeon 8060S (`gfx1151`) after `ssm_update_kernel` gained a HIP port. Before it, ROCm ran the Mamba2 SSD graph for every single-token step, so the `w1ctx512` rows in `../rocm_gfx1151_3c9edea0/` compared Metal's fused kernel with ROCm's graph. These compare kernel with kernel, and the `ssmkernel0` rows give the graph on the same binary for a same-backend A/B. `METADATA.txt` records the host, versions, commit and binary hash; `RUNS.txt` the exit status and row count of every run; `SHA256SUMS` covers every trace. The comparison is in `docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md`. + +```bash +cargo build --release --features rocm --example logit_trace +T=./target/release/examples/logit_trace +C=tests/fixtures/wikitext2_excerpt.txt +P=rocm_gfx1151_96cbce84 +for pair in granite-4.0-h-tiny:granite-4.0-h-tiny-4bit nemotron-3-nano-30b-a3b:NVIDIA-Nemotron-3-Nano-30B-A3B-4bit; do + tag=${pair%%:*}; dir=${pair##*:} + MLXCEL_TRACE_START_TOKEN=512 $T models/mlx/$dir $C 1 128 8 512 > ${P}_${tag}_default_w1ctx512.tsv + MLXCEL_SSM_KERNEL=0 MLXCEL_TRACE_START_TOKEN=512 $T models/mlx/$dir $C 1 128 8 512 > ${P}_${tag}_ssmkernel0_w1ctx512.tsv + $T models/mlx/$dir $C 1 128 8 0 > ${P}_${tag}_default_w1.tsv +done +``` + +The Nemotron-H traces carry five `[NemotronH]` loader lines on stdout; filter them with `grep -v '^\[NemotronH\] '` on both sides before `scripts/compare_logit_traces.py`. diff --git a/benchmarks/logit_traces/rocm_gfx1151_96cbce84/RUNS.txt b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/RUNS.txt new file mode 100644 index 000000000..aa8119acb --- /dev/null +++ b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/RUNS.txt @@ -0,0 +1,6 @@ +rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1ctx512 exit=0 rows=135 elapsed=120s +rocm_gfx1151_96cbce84_granite-4.0-h-tiny_ssmkernel0_w1ctx512 exit=0 rows=135 elapsed=115s +rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1 exit=0 rows=135 elapsed=4s +rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1ctx512 exit=0 rows=140 elapsed=235s +rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_ssmkernel0_w1ctx512 exit=0 rows=140 elapsed=223s +rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1 exit=0 rows=140 elapsed=6s diff --git a/benchmarks/logit_traces/rocm_gfx1151_96cbce84/SHA256SUMS b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/SHA256SUMS new file mode 100644 index 000000000..96f469087 --- /dev/null +++ b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/SHA256SUMS @@ -0,0 +1,6 @@ +dea067abe86dde12683271aff5683e1ca231a23a99c3e6d5cd1c5411e0c60063 rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1ctx512.tsv +a6caad4d60d0bca640d0b39065ade10f1e412b36f8204fae56510ad9e3c9dbbe rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1.tsv +70b1b5e7ff6ec79cd61e2437bd3fb26b377f88570554c963d1010a1c5dada116 rocm_gfx1151_96cbce84_granite-4.0-h-tiny_ssmkernel0_w1ctx512.tsv +7cdbb6ad3ac999fc379771060b296e17fb062e6c64b9976007df1a84e37e1a7c rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1ctx512.tsv +9326e8bde7aaad14c7e57aabc48f4eadfecd5a9b746f82da67dc9b26a2ccdf98 rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1.tsv +40bb05103bc5c5dffba821328756f346432ff05a73775e4526444f40b062322b rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_ssmkernel0_w1ctx512.tsv diff --git a/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1.tsv b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1.tsv new file mode 100644 index 000000000..06ab5717e --- /dev/null +++ b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1.tsv @@ -0,0 +1,135 @@ +# model models/mlx/granite-4.0-h-tiny-4bit +# device GPU +# corpus tests/fixtures/wikitext2_excerpt.txt tokens 289403 +# chunks 128 chunk_tokens 1 topk 8 prefill 0 +# columns chunk pos target nll top_ids top_logits +# start_token 0 +# contiguous false +0 0 43205 18.296762 198,220,279,271,720,58,16,311 18.875000,15.500000,15.000000,14.437500,14.062500,14.000000,13.750000,13.687500 +1 0 275 9.115309 198,271,279,220,311,791,40,264 17.375000,14.750000,14.562500,14.312500,13.750000,13.562500,13.312500,13.250000 +2 0 428 13.753754 198,279,220,271,311,264,58,720 18.125000,15.062500,15.000000,14.500000,14.250000,13.750000,13.437500,13.312500 +3 0 12 6.231743 198,220,279,271,58,311,720,264 18.500000,15.250000,15.000000,14.812500,14.562500,14.250000,14.000000,13.687500 +4 0 17 5.710806 198,220,279,271,720,16,264,311 18.375000,15.750000,15.500000,14.437500,14.312500,14.187500,14.187500,14.062500 +5 0 3475 14.424252 198,220,279,271,58,311,264,720 18.500000,15.187500,14.937500,14.750000,14.000000,13.937500,13.750000,13.687500 +6 0 76834 14.907146 198,279,220,271,311,264,720,58 18.000000,16.125000,16.000000,15.187500,14.750000,14.687500,14.187500,14.000000 +7 0 80583 14.663322 198,220,271,279,311,264,58,459 18.250000,15.750000,15.750000,15.562500,14.812500,14.250000,13.562500,13.562500 +8 0 3692 12.923090 198,220,720,271,58,279,311,334 19.375000,16.125000,14.937500,14.750000,14.562500,14.437500,13.500000,13.437500 +9 0 257 9.496799 198,271,220,279,58,311,264,512 18.875000,15.937500,15.562500,15.437500,14.812500,14.437500,14.187500,14.187500 +10 0 551 9.481077 198,220,279,720,271,311,264,58 18.625000,16.625000,15.562500,15.187500,14.750000,14.562500,14.500000,14.312500 +11 0 3788 12.794767 198,220,279,271,311,264,58,720 18.750000,15.812500,15.687500,14.562500,14.312500,14.250000,13.750000,13.750000 +12 0 1129 14.763803 198,220,271,279,311,58,720,264 19.750000,15.937500,15.687500,15.437500,14.687500,14.562500,14.437500,13.812500 +13 0 71 7.514406 198,279,220,271,311,264,1148,720 18.250000,15.250000,15.062500,14.750000,14.000000,13.500000,13.437500,13.062500 +14 0 36368 16.065123 198,220,279,271,311,58,264,720 18.500000,15.000000,14.937500,14.687500,14.000000,13.937500,13.750000,13.750000 +15 0 1594 12.394127 198,220,279,264,271,311,459,720 17.625000,16.000000,16.000000,14.750000,14.750000,14.687500,13.937500,13.812500 +16 0 6973 18.462471 198,279,220,271,311,58,264,720 19.000000,15.500000,15.312500,15.312500,14.562500,14.312500,14.062500,13.812500 +17 0 3529 15.583487 198,220,271,279,720,58,311,264 18.750000,15.250000,15.187500,14.750000,14.312500,13.812500,13.812500,13.250000 +18 0 77749 17.923609 198,220,279,271,58,311,264,720 19.500000,15.750000,15.750000,15.562500,14.562500,14.562500,14.187500,14.187500 +19 0 11628 17.739681 198,271,220,279,58,311,264,11 19.500000,16.500000,15.687500,15.500000,15.312500,14.812500,14.187500,13.937500 +20 0 3916 16.248640 198,220,271,279,311,58,264,720 19.125000,15.750000,15.437500,15.437500,14.250000,13.937500,13.937500,13.937500 +21 0 9009 12.491665 198,279,220,271,311,264,58,1148 18.875000,15.562500,15.500000,15.187500,14.750000,14.250000,14.062500,13.750000 +22 0 6458 15.024443 198,279,220,311,264,271,720,459 17.875000,16.000000,15.312500,14.937500,14.750000,14.750000,14.062500,13.562500 +23 0 61503 17.485357 198,220,279,271,311,264,720,1148 18.500000,15.687500,15.500000,14.750000,14.500000,14.000000,13.562500,13.500000 +24 0 428 14.177526 198,220,279,311,271,264,720,459 16.750000,15.187500,14.312500,14.000000,13.937500,13.187500,13.062500,12.500000 +25 0 198 0.419243 198,220,279,271,58,311,720,264 18.500000,15.250000,15.000000,14.812500,14.562500,14.250000,14.000000,13.687500 +26 0 71684 16.023062 198,271,279,220,58,311,40,264 19.875000,15.687500,15.500000,15.312500,14.562500,14.437500,14.250000,14.187500 +27 0 257 8.707249 198,220,271,279,58,264,311,720 17.375000,14.937500,14.562500,14.562500,13.687500,13.562500,13.562500,13.500000 +28 0 551 9.481077 198,220,279,720,271,311,264,58 18.625000,16.625000,15.562500,15.187500,14.750000,14.562500,14.500000,14.312500 +29 0 289 9.763517 198,220,279,271,311,264,58,720 18.750000,15.812500,15.687500,14.562500,14.312500,14.250000,13.750000,13.750000 +30 0 61503 17.130968 198,279,220,271,311,264,58,1148 18.375000,15.187500,14.937500,14.812500,14.187500,13.812500,13.687500,13.250000 +31 0 428 14.177526 198,220,279,311,271,264,720,459 16.750000,15.187500,14.312500,14.000000,13.937500,13.187500,13.062500,12.500000 +32 0 12 6.231743 198,220,279,271,58,311,720,264 18.500000,15.250000,15.000000,14.812500,14.562500,14.250000,14.000000,13.687500 +33 0 17 5.710806 198,220,279,271,720,16,264,311 18.375000,15.750000,15.500000,14.437500,14.312500,14.187500,14.187500,14.062500 +34 0 12 6.799252 198,220,279,271,58,311,264,720 18.500000,15.187500,14.937500,14.750000,14.000000,13.937500,13.750000,13.687500 +35 0 1059 15.882681 198,220,279,271,720,16,264,311 18.375000,15.750000,15.500000,14.437500,14.312500,14.187500,14.187500,14.062500 +36 0 8437 15.759243 198,220,279,271,58,311,720,998 18.125000,15.000000,14.437500,14.312500,13.937500,13.687500,13.500000,13.062500 +37 0 16 5.860119 198,220,279,271,311,720,58,264 18.750000,15.687500,15.187500,14.500000,14.250000,14.250000,14.000000,13.937500 +38 0 198 0.264400 198,220,279,271,58,311,264,16 19.000000,15.250000,15.062500,14.937500,14.250000,14.187500,13.937500,13.750000 +39 0 20805 15.241812 198,271,279,220,58,311,40,264 19.875000,15.687500,15.500000,15.312500,14.562500,14.437500,14.250000,14.187500 +40 0 415 10.797653 198,220,279,271,58,311,264,16 18.750000,15.437500,15.437500,15.000000,14.687500,14.187500,14.062500,13.437500 +41 0 551 9.485817 198,220,279,720,271,264,311,58 19.000000,16.375000,15.812500,15.000000,14.750000,14.562500,14.500000,13.937500 +42 0 1296 13.107267 198,220,279,271,311,264,58,720 18.750000,15.812500,15.687500,14.562500,14.312500,14.250000,13.750000,13.750000 +43 0 320 7.715372 198,271,279,220,311,58,264,40 19.125000,16.125000,16.125000,15.437500,15.062500,15.000000,14.687500,14.562500 +44 0 19305 13.075552 198,279,220,271,40,311,264,1148 18.750000,15.312500,15.000000,14.750000,14.062500,14.062500,13.937500,13.812500 +45 0 23 7.261166 198,220,279,271,311,720,58,264 18.375000,15.562500,14.937500,14.687500,14.062500,13.750000,13.687500,13.687500 +46 0 7123 16.862946 198,271,220,279,311,58,264,1148 18.625000,15.437500,15.312500,15.250000,14.312500,14.187500,13.937500,13.437500 +47 0 340 9.029768 198,220,279,271,311,58,264,720 19.375000,16.125000,16.125000,15.250000,15.062500,14.687500,14.687500,14.250000 +48 0 10028 13.109947 198,271,220,279,58,720,311,264 19.625000,15.687500,15.562500,15.312500,14.937500,14.500000,14.312500,14.000000 +49 0 262 7.594938 198,220,271,279,311,720,264,512 18.625000,16.625000,15.562500,15.562500,15.000000,14.562500,14.312500,13.812500 +50 0 551 9.784988 198,220,279,271,720,264,311,58 18.750000,16.125000,15.812500,15.562500,14.750000,14.562500,14.562500,13.687500 +51 0 13844 13.732267 198,220,279,271,311,264,58,720 18.750000,15.812500,15.687500,14.562500,14.312500,14.250000,13.750000,13.750000 +52 0 7866 14.005470 198,279,220,271,311,264,720,459 18.250000,16.000000,15.562500,15.062500,14.937500,14.500000,13.750000,13.687500 +53 0 93314 14.376535 198,220,279,271,264,311,720,1148 18.375000,16.625000,16.500000,14.937500,14.812500,14.562500,14.250000,14.187500 +54 0 220 4.056431 198,220,271,279,311,264,720,58 19.500000,15.687500,15.687500,15.562500,14.562500,14.187500,14.187500,14.000000 +55 0 18 6.230272 198,279,220,271,311,16,264,58 18.625000,15.250000,15.187500,14.562500,14.062500,13.937500,13.937500,13.687500 +56 0 13 8.507011 198,220,271,279,58,311,264,720 18.875000,15.250000,15.187500,15.062500,14.250000,14.250000,13.937500,13.500000 +57 0 15 7.375696 198,279,220,271,40,791,264,58 19.250000,15.687500,15.562500,15.562500,14.562500,14.562500,14.437500,14.312500 +58 0 320 8.259653 198,220,271,279,58,311,264,720 19.375000,15.250000,15.187500,15.187500,14.250000,14.250000,13.812500,13.687500 +59 0 2485 13.294302 198,279,220,271,40,311,264,1148 18.750000,15.312500,15.000000,14.750000,14.062500,14.062500,13.937500,13.812500 +60 0 1129 13.280151 198,220,279,720,271,311,58,998 18.125000,15.312500,14.187500,14.187500,14.000000,13.750000,13.062500,12.687500 +61 0 86331 16.155031 198,279,220,271,311,264,1148,720 18.250000,15.250000,15.062500,14.750000,14.000000,13.500000,13.437500,13.062500 +62 0 2726 12.325586 198,220,279,271,311,720,264,459 18.625000,15.687500,15.000000,14.812500,14.250000,13.812500,13.437500,13.250000 +63 0 7116 15.607037 198,220,279,271,311,720,264,58 18.875000,16.000000,15.562500,15.437500,14.687500,14.437500,14.000000,13.750000 +64 0 68833 17.580374 198,220,279,271,720,311,264,459 19.375000,16.375000,16.000000,15.687500,15.062500,14.812500,14.062500,13.562500 +65 0 1355 13.232002 198,220,279,271,311,264,720,1148 18.750000,16.125000,16.125000,14.750000,14.562500,14.312500,14.000000,13.812500 +66 0 64 5.806099 198,279,220,271,311,264,58,998 18.875000,15.500000,15.000000,14.812500,14.437500,14.000000,13.812500,13.687500 +67 0 14 10.409391 198,220,279,271,58,311,264,720 18.500000,15.187500,14.937500,14.500000,14.000000,13.937500,13.812500,13.750000 +68 0 18 6.453409 198,279,220,271,311,264,1148,791 18.875000,15.750000,15.062500,15.000000,14.500000,14.250000,14.187500,14.062500 +69 0 13 8.507011 198,220,271,279,58,311,264,720 18.875000,15.250000,15.187500,15.062500,14.250000,14.250000,13.937500,13.500000 +70 0 15 7.375696 198,279,220,271,40,791,264,58 19.250000,15.687500,15.562500,15.562500,14.562500,14.562500,14.437500,14.312500 +71 0 54660 16.931528 198,220,271,279,58,311,264,720 19.375000,15.250000,15.187500,15.187500,14.250000,14.250000,13.812500,13.687500 +72 0 20241 13.700876 198,271,279,220,311,58,264,720 20.500000,16.750000,16.375000,16.250000,15.312500,14.750000,14.750000,14.687500 +73 0 1516 15.071934 198,220,279,271,311,58,264,998 17.250000,14.812500,14.062500,13.812500,13.687500,13.250000,12.937500,12.937500 +74 0 25 8.579760 198,220,279,271,311,264,58,720 18.750000,15.812500,15.562500,15.312500,14.812500,14.437500,14.312500,14.062500 +75 0 27685 13.977633 198,220,279,264,271,311,720,1148 18.250000,16.250000,15.937500,14.750000,14.562500,14.500000,14.312500,14.062500 +76 0 20965 15.450951 198,279,220,271,311,264,459,1148 18.125000,16.875000,16.375000,16.000000,15.437500,15.312500,14.750000,14.562500 +77 0 11 6.225031 198,279,220,271,311,264,1148,459 19.375000,16.250000,15.937500,15.812500,15.562500,14.812500,14.437500,14.000000 +78 0 82035 17.809288 198,279,220,271,311,264,1148,40 19.125000,16.000000,15.500000,15.062500,14.687500,14.562500,14.437500,14.062500 +79 0 15592 11.026390 198,279,271,220,311,264,459,1148 17.625000,16.000000,15.562500,15.312500,15.000000,14.687500,14.187500,13.937500 +80 0 8483 13.499424 198,279,311,220,271,264,459,720 16.750000,15.437500,14.750000,14.687500,14.500000,14.000000,13.437500,13.062500 +81 0 198 0.537973 198,271,279,220,311,264,1148,791 19.000000,16.375000,16.125000,15.812500,15.062500,14.750000,14.562500,14.312500 +82 0 1835 13.210562 198,271,279,220,58,311,40,264 19.875000,15.687500,15.500000,15.312500,14.562500,14.437500,14.250000,14.187500 +83 0 320 8.260883 198,220,279,271,720,58,311,264 19.375000,16.000000,15.437500,14.812500,14.687500,14.250000,14.250000,14.187500 +84 0 50 7.763052 198,279,220,271,40,311,264,1148 18.750000,15.312500,15.000000,14.750000,14.062500,14.062500,13.937500,13.812500 +85 0 13 8.431351 198,271,220,279,58,720,311,512 19.125000,15.250000,15.000000,14.812500,14.437500,13.937500,13.812500,13.750000 +86 0 8930 15.344446 198,279,220,271,40,791,264,58 19.250000,15.687500,15.562500,15.562500,14.562500,14.562500,14.437500,14.312500 +87 0 488 13.400963 198,279,220,271,311,264,791,720 18.000000,14.750000,14.687500,14.562500,13.687500,13.312500,13.187500,12.937500 +88 0 11 6.137121 198,271,279,220,311,58,264,791 18.625000,15.500000,15.437500,15.250000,14.500000,14.187500,14.187500,13.750000 +89 0 452 9.817101 198,279,220,271,311,264,1148,40 19.125000,16.000000,15.500000,15.062500,14.687500,14.562500,14.437500,14.062500 +90 0 13 7.948751 198,279,220,271,311,720,264,58 18.250000,15.250000,15.187500,15.000000,14.062500,14.000000,13.937500,13.500000 +91 0 50356 17.883509 198,279,220,271,40,791,264,58 19.250000,15.687500,15.562500,15.562500,14.562500,14.562500,14.437500,14.312500 +92 0 819 13.374430 198,271,279,220,791,40,311,3923 17.875000,14.437500,14.187500,13.812500,13.812500,13.562500,13.062500,13.062500 +93 0 84077 17.042751 198,279,220,271,311,264,1148,58 18.000000,15.312500,15.062500,14.687500,14.187500,14.062500,13.437500,13.250000 +94 0 29234 12.184427 198,220,279,271,311,264,1148,791 17.375000,14.812500,14.500000,14.187500,13.687500,13.250000,12.812500,12.750000 +95 0 11 5.538054 198,220,279,271,311,264,58,720 18.125000,15.437500,15.062500,14.750000,14.312500,13.937500,13.812500,13.437500 +96 0 432 9.567101 198,279,220,271,311,264,1148,40 19.125000,16.000000,15.500000,15.062500,14.687500,14.562500,14.437500,14.062500 +97 0 13 8.027153 198,220,279,271,720,311,264,58 18.375000,15.312500,15.312500,15.062500,14.437500,14.250000,14.062500,13.750000 +98 0 2100 9.875696 198,279,220,271,40,791,264,58 19.250000,15.687500,15.562500,15.562500,14.562500,14.562500,14.437500,14.312500 +99 0 9211 14.479223 198,279,220,271,264,311,1148,720 18.875000,16.250000,15.937500,15.187500,14.750000,14.750000,14.562500,14.437500 +100 0 2001 10.646482 198,220,279,311,271,264,720,459 17.000000,15.312500,15.187500,14.437500,13.937500,13.812500,13.312500,13.000000 +101 0 330 6.227028 198,220,279,311,264,271,720,1148 18.875000,16.375000,16.125000,14.750000,14.562500,14.562500,14.562500,14.312500 +102 0 9261 11.700617 198,279,271,220,40,791,1148,311 19.125000,15.500000,15.250000,14.750000,14.687500,14.437500,14.312500,14.250000 +103 0 57893 17.266718 198,271,220,279,58,311,512,720 18.875000,15.437500,15.250000,14.750000,14.687500,14.000000,13.937500,13.750000 +104 0 386 10.106499 198,271,279,220,58,311,264,40 18.750000,15.812500,15.312500,15.250000,14.500000,14.437500,14.062500,13.812500 +105 0 13025 17.719200 198,279,220,271,311,720,264,58 18.625000,15.437500,15.250000,15.250000,14.250000,14.187500,14.062500,13.812500 +106 0 27972 15.818333 198,220,271,279,58,311,264,720 18.125000,15.437500,15.000000,15.000000,14.000000,14.000000,13.937500,13.812500 +107 0 498 12.901283 198,220,279,271,311,264,720,58 18.875000,16.375000,15.750000,15.562500,15.187500,14.562500,14.562500,14.437500 +108 0 220 3.528170 198,220,279,271,264,720,311,1148 19.625000,16.375000,16.375000,15.000000,14.937500,14.937500,14.812500,14.562500 +109 0 679 8.605272 198,279,220,271,311,16,264,58 18.625000,15.250000,15.187500,14.562500,14.062500,13.937500,13.937500,13.687500 +110 0 21 6.675917 198,220,279,271,311,720,264,58 18.625000,15.812500,15.062500,14.687500,14.187500,13.937500,13.750000,13.562500 +111 0 340 9.080715 198,220,271,279,311,58,264,1148 18.875000,15.312500,15.250000,15.187500,14.187500,14.062500,13.937500,13.500000 +112 0 1835 12.109947 198,271,220,279,58,720,311,264 19.625000,15.687500,15.562500,15.312500,14.937500,14.500000,14.312500,14.000000 +113 0 40283 15.870258 198,220,279,271,720,58,311,264 19.375000,16.000000,15.437500,14.812500,14.687500,14.250000,14.250000,14.187500 +114 0 3786 15.619562 198,271,220,279,311,58,264,1148 19.125000,16.375000,15.812500,15.750000,14.812500,14.750000,14.500000,13.687500 +115 0 25 9.455753 198,279,220,271,311,264,58,1148 19.125000,15.812500,15.687500,15.687500,15.000000,14.437500,14.250000,14.187500 +116 0 3788 12.102633 198,220,279,264,271,311,720,1148 18.250000,16.250000,15.937500,14.750000,14.562500,14.500000,14.312500,14.062500 +117 0 1129 14.763803 198,220,271,279,311,58,720,264 19.750000,15.937500,15.687500,15.437500,14.687500,14.562500,14.437500,13.812500 +118 0 71 7.514406 198,279,220,271,311,264,1148,720 18.250000,15.250000,15.062500,14.750000,14.000000,13.500000,13.437500,13.062500 +119 0 36368 16.065123 198,220,279,271,311,58,264,720 18.500000,15.000000,14.937500,14.687500,14.000000,13.937500,13.750000,13.750000 +120 0 1594 12.394127 198,220,279,264,271,311,459,720 17.625000,16.000000,16.000000,14.750000,14.750000,14.687500,13.937500,13.812500 +121 0 6973 18.462471 198,279,220,271,311,58,264,720 19.000000,15.500000,15.312500,15.312500,14.562500,14.312500,14.062500,13.812500 +122 0 3529 15.583487 198,220,271,279,720,58,311,264 18.750000,15.250000,15.187500,14.750000,14.312500,13.812500,13.812500,13.250000 +123 0 77749 17.923609 198,220,279,271,58,311,264,720 19.500000,15.750000,15.750000,15.562500,14.562500,14.562500,14.187500,14.187500 +124 0 11628 17.739681 198,271,220,279,58,311,264,11 19.500000,16.500000,15.687500,15.500000,15.312500,14.812500,14.187500,13.937500 +125 0 3916 16.248640 198,220,271,279,311,58,264,720 19.125000,15.750000,15.437500,15.437500,14.250000,13.937500,13.937500,13.937500 +126 0 9009 12.491665 198,279,220,271,311,264,58,1148 18.875000,15.562500,15.500000,15.187500,14.750000,14.250000,14.062500,13.750000 +127 0 6458 15.024443 198,279,220,311,264,271,720,459 17.875000,16.000000,15.312500,14.937500,14.750000,14.750000,14.062500,13.562500 diff --git a/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1ctx512.tsv b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1ctx512.tsv new file mode 100644 index 000000000..40fabac81 --- /dev/null +++ b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_default_w1ctx512.tsv @@ -0,0 +1,135 @@ +# model models/mlx/granite-4.0-h-tiny-4bit +# device GPU +# corpus tests/fixtures/wikitext2_excerpt.txt tokens 289403 +# chunks 128 chunk_tokens 1 topk 8 prefill 512 +# columns chunk pos target nll top_ids top_logits +# start_token 512 +# contiguous false +0 0 323 0.123800 323,612,1174,662,304,13,1405,1990 25.125000,22.375000,21.375000,21.125000,20.000000,19.375000,19.125000,18.875000 +1 0 48550 0.174885 48550,574,6476,520,37778,4410,10455,48043 22.625000,20.125000,19.250000,18.000000,17.875000,17.625000,17.500000,17.125000 +2 0 5721 0.000135 5721,71,1494,4317,13279,8350,6798,65 29.625000,20.625000,17.000000,16.500000,16.500000,15.937500,15.187500,14.812500 +3 0 662 0.898527 662,323,1174,13,304,1990,2652,505 23.625000,22.625000,22.625000,21.625000,21.625000,20.875000,20.500000,20.375000 +4 0 426 2.172865 1283,763,426,578,720,5414,8563,1102 22.625000,21.625000,21.125000,19.875000,19.875000,19.375000,19.375000,18.375000 +5 0 11206 0.000046 11206,619,5049,3102,3514,1656,13,980 26.875000,15.250000,14.437500,14.250000,13.812500,13.750000,13.562500,13.250000 +6 0 466 0.000175 466,5153,783,998,2641,4253,376,977 26.625000,16.750000,16.625000,15.062500,15.062500,14.687500,14.437500,14.312500 +7 0 59335 1.908243 574,6476,59335,1101,9922,1243,45877,1047 22.375000,22.375000,22.375000,22.125000,21.625000,21.125000,20.875000,20.750000 +8 0 304 1.112289 16662,304,439,14329,449,389,1828,1578 25.375000,24.875000,23.625000,21.500000,20.750000,20.125000,20.125000,19.875000 +9 0 1403 4.728743 279,264,420,578,2650,459,3892,1403 23.375000,20.875000,20.250000,20.125000,20.125000,19.625000,19.375000,19.125000 +10 0 12631 2.209673 54248,11335,12707,12631,18243,220,1023,6007 22.625000,22.500000,22.500000,22.375000,21.875000,21.625000,21.000000,20.875000 +11 0 304 0.753664 304,15910,6004,11,1174,430,9124,369 23.875000,22.625000,22.000000,21.375000,21.250000,21.125000,20.500000,20.375000 +12 0 220 0.088245 220,279,4216,3389,4062,11295,430,902 26.500000,23.500000,22.000000,21.375000,20.625000,20.375000,20.125000,20.000000 +13 0 1049 0.063057 1049,679,2550,2366,508,17,3753,4468 24.625000,21.875000,17.250000,16.000000,15.187500,14.812500,14.562500,14.250000 +14 0 23 0.923016 22,23,24,21,20,19,18,15 22.125000,22.000000,20.625000,19.875000,17.125000,15.437500,14.937500,14.687500 +15 0 1174 1.020199 1174,2652,662,551,13,25,26,2001 23.125000,22.625000,22.125000,21.875000,20.375000,20.125000,20.000000,20.000000 +16 0 6187 8.655167 279,578,5737,15910,832,2225,25891,40500 20.000000,19.875000,19.625000,18.875000,18.125000,18.125000,18.125000,17.625000 +17 0 4238 0.978745 4238,10573,315,369,9137,1712,18811,8828 20.875000,20.000000,19.875000,19.250000,18.875000,18.625000,18.125000,17.375000 +18 0 4997 3.385035 323,320,1174,15371,52483,4997,49538,46046 19.750000,18.875000,18.625000,18.125000,18.125000,17.625000,17.625000,16.875000 +19 0 31304 -0.000000 31304,655,1941,7278,88,7006,5515,727 37.250000,22.875000,22.500000,21.250000,19.625000,18.375000,17.750000,17.375000 +20 0 555 5.772234 323,15910,1174,320,16662,14329,902,5439 24.125000,23.375000,22.500000,22.125000,21.625000,21.125000,20.875000,20.625000 +21 0 57337 6.956081 7957,12131,8563,6941,11291,17520,15469,11519 19.375000,18.875000,18.750000,18.625000,18.625000,18.625000,18.500000,18.375000 +22 0 12366 10.043489 6941,7505,7957,11517,15469,29335,12131,13929 18.375000,18.375000,18.375000,17.750000,17.750000,17.625000,17.500000,17.500000 +23 0 2009 8.241274 69012,362,81431,1676,47790,22791,11258,6938 18.625000,16.125000,15.437500,15.250000,14.937500,14.750000,13.937500,13.812500 +24 0 546 3.491472 288,546,688,613,589,64,377,65 20.000000,16.625000,16.375000,16.125000,14.812500,14.562500,14.500000,14.312500 +25 0 72 0.000050 72,64044,285,437,288,99355,88,290 30.625000,19.875000,19.375000,18.000000,17.875000,17.875000,17.750000,17.375000 +26 0 1174 1.058737 323,1174,449,320,40500,304,662,11 24.375000,23.875000,20.875000,20.500000,20.375000,20.250000,19.500000,19.250000 +27 0 323 0.678768 323,902,304,889,15910,1405,369,8272 24.625000,24.125000,22.125000,21.500000,21.375000,21.125000,20.375000,20.375000 +28 0 4418 8.501568 578,279,304,52483,362,800,3842,16290 20.875000,20.500000,20.000000,18.625000,18.375000,17.625000,17.625000,17.625000 +29 0 798 2.045752 956,798,3489,3919,90970,364,29868,12565 17.000000,15.812500,14.562500,14.250000,14.062500,13.937500,13.750000,13.687500 +30 0 54275 0.126122 54275,79,96363,445,21004,11617,596,364 21.250000,17.375000,17.375000,16.750000,16.625000,16.250000,16.000000,15.937500 +31 0 15910 1.125801 15910,555,1174,902,320,662,13,11 23.375000,23.000000,22.500000,21.500000,21.375000,21.375000,21.250000,21.125000 +32 0 555 0.001562 555,323,1578,304,520,7043,3131,8096 28.875000,21.625000,20.125000,19.250000,19.125000,18.625000,18.375000,18.125000 +33 0 12225 6.292309 7505,54424,30411,11291,16768,18984,19475,29656 18.000000,15.500000,14.812500,14.750000,14.750000,14.250000,14.250000,14.187500 +34 0 398 0.602987 398,29306,6885,5657,402,2642,23797,6246 19.875000,19.625000,15.250000,15.000000,14.000000,14.000000,13.562500,13.437500 +35 0 89453 1.980915 63592,5348,89453,32103,35934,13566,23162,46881 16.625000,16.000000,15.812500,15.687500,15.000000,14.062500,13.750000,13.437500 +36 0 662 0.375700 662,1174,323,13,2652,16853,320,449 23.875000,21.875000,21.625000,21.375000,20.750000,19.500000,19.375000,19.375000 +37 0 763 2.132114 1283,763,426,578,11995,720,5414,6187 22.625000,21.125000,20.875000,20.375000,20.125000,19.875000,19.625000,19.000000 +38 0 3297 3.852106 220,2225,279,6187,5369,3297,6287,5651 23.875000,22.500000,22.125000,21.875000,20.750000,20.625000,19.875000,19.500000 +39 0 220 0.178417 220,315,1174,11,323,568,279,430 24.375000,21.875000,21.875000,19.625000,19.375000,18.875000,18.375000,18.375000 +40 0 1049 0.203209 1049,679,2366,2550,508,220,3753,17 23.500000,22.000000,16.875000,15.312500,13.437500,13.312500,13.187500,13.062500 +41 0 23 0.702232 23,24,22,21,20,19,18,15 23.125000,23.125000,18.625000,17.625000,16.875000,15.312500,14.812500,14.562500 +42 0 1174 0.763935 1174,568,11,426,279,8563,578,433 24.000000,23.500000,22.375000,22.375000,21.125000,19.500000,19.125000,18.875000 +43 0 426 0.905916 568,426,279,8563,578,433,813,12360 23.625000,23.500000,21.250000,20.375000,19.875000,18.875000,18.750000,18.750000 +44 0 11206 0.000027 11206,619,1656,980,3102,3514,5049,1286 28.500000,17.375000,15.062500,14.562500,14.562500,14.500000,14.312500,13.812500 +45 0 466 0.000214 466,783,668,998,5153,655,2641,2002 27.625000,18.875000,15.812500,15.812500,15.437500,15.312500,15.062500,14.812500 +46 0 1903 3.445528 574,59335,9922,6476,10887,1903,4036,11096 22.625000,22.625000,22.500000,22.000000,21.625000,20.875000,20.625000,20.625000 +47 0 264 2.152884 813,264,459,3925,279,2500,31186,12707 26.500000,24.625000,24.500000,21.125000,20.875000,20.875000,20.875000,20.750000 +48 0 8810 0.328165 8810,471,3361,10015,6566,9024,17755,95156 24.875000,22.875000,21.500000,21.500000,21.125000,20.750000,20.500000,20.375000 +49 0 11341 0.071541 11341,571,6917,40500,17755,471,5178,7858 21.500000,17.875000,17.250000,16.875000,15.562500,15.500000,15.250000,15.187500 +50 0 389 0.360193 389,304,439,520,369,5737,16662,315 27.500000,26.125000,25.500000,23.875000,22.375000,22.125000,21.625000,20.625000 +51 0 264 3.843517 279,459,264,18588,6566,578,12707,6460 23.500000,20.500000,19.875000,19.625000,19.125000,19.000000,18.875000,18.750000 +52 0 1403 5.506533 12707,220,4101,18588,3361,18629,6007,9417 21.375000,20.375000,19.125000,19.125000,18.875000,18.875000,18.500000,18.375000 +53 0 571 4.731678 29137,482,961,12,9417,571,58098,21925 20.875000,20.500000,20.125000,19.750000,18.500000,17.125000,16.375000,16.125000 +54 0 12 0.003632 12,59562,35394,51442,12898,28751,39743,52656 25.875000,20.125000,16.625000,15.187500,15.062500,15.000000,15.000000,14.937500 +55 0 31 0.000040 31,19908,571,9164,19741,555,5,63899 28.125000,16.500000,15.750000,15.500000,15.187500,14.562500,14.250000,14.250000 +56 0 961 0.008816 961,9417,1370,6596,2046,1938,9568,3814 29.000000,23.500000,23.000000,22.125000,21.375000,20.875000,20.125000,19.875000 +57 0 9417 0.296082 9417,12707,3361,6007,4101,18588,3446,1514 23.125000,21.375000,19.500000,19.375000,19.125000,18.625000,18.375000,18.250000 +58 0 15952 14.955324 315,330,25891,320,304,20458,11,1174 27.375000,21.875000,21.625000,21.500000,20.875000,20.375000,20.125000,20.125000 +59 0 315 0.064276 315,389,304,25891,369,320,330,20458 26.375000,22.500000,22.000000,21.500000,21.375000,20.500000,20.250000,20.125000 +60 0 279 0.132408 279,578,6460,16071,49922,18588,19150,54943 23.625000,19.750000,19.375000,19.125000,18.875000,18.750000,18.250000,17.875000 +61 0 12707 0.426838 12707,18588,8013,5526,13673,94455,4101,13740 23.375000,21.125000,20.750000,20.125000,20.125000,19.625000,19.375000,19.250000 +62 0 4101 0.045280 4101,6275,20156,1501,9977,23160,27883,18629 24.875000,20.375000,20.125000,19.375000,19.375000,19.250000,19.125000,18.875000 +63 0 468 6.593527 578,11,16071,1174,49922,3165,6460,86550 18.875000,18.125000,18.125000,17.875000,17.500000,17.000000,17.000000,16.875000 +64 0 1802 1.519329 6310,1802,16,1721,1924,55889,26286,19011 18.625000,17.500000,15.812500,14.000000,14.000000,14.000000,13.562500,12.937500 +65 0 279 0.055069 279,578,791,11,220,16454,1820,555 27.375000,24.500000,17.375000,16.125000,15.750000,15.562500,15.437500,15.250000 +66 0 15371 0.000093 15371,5710,29420,84088,79169,81834,33160,83942 31.125000,21.500000,20.375000,18.125000,16.750000,16.750000,16.625000,15.562500 +67 0 1174 0.988344 1174,662,439,304,2652,902,323,15910 22.875000,22.500000,21.375000,21.000000,20.875000,20.250000,19.875000,19.500000 +68 0 8272 5.409058 902,5737,25891,304,15910,323,1405,5439 23.500000,23.500000,22.500000,22.125000,22.125000,21.875000,21.500000,21.000000 +69 0 555 0.053879 555,304,3010,279,20193,449,264,1828 29.000000,25.500000,24.875000,22.625000,22.000000,21.875000,21.500000,21.500000 +70 0 459 3.171417 264,459,40500,1403,2500,27351,813,8810 25.375000,22.500000,22.500000,22.125000,21.625000,21.375000,21.250000,21.250000 +71 0 11341 0.122332 11341,9417,8223,14529,15718,11838,220,653 23.500000,20.625000,19.500000,18.000000,18.000000,17.750000,17.625000,17.625000 +72 0 389 0.577507 389,304,439,3010,279,520,14329,16662 27.500000,27.000000,25.500000,22.875000,22.125000,22.125000,22.125000,21.625000 +73 0 279 0.915653 279,264,459,49922,578,16071,2500,52483 23.125000,22.875000,21.250000,20.625000,20.500000,19.875000,19.500000,19.375000 +74 0 12707 0.317753 12707,18588,8013,220,4948,4101,9417,23160 23.875000,20.625000,20.250000,20.125000,19.875000,19.625000,19.625000,19.375000 +75 0 4101 0.038681 4101,20156,1332,6275,18629,1501,2068,23160 25.125000,19.875000,19.750000,19.750000,19.500000,19.250000,19.250000,19.125000 +76 0 29878 9.245026 16071,49922,578,86550,11,14894,6460,11301 19.625000,19.250000,19.125000,18.375000,17.875000,17.875000,17.375000,17.125000 +77 0 86493 0.123402 86493,2299,535,344,6632,2270,6672,10126 27.500000,25.375000,23.000000,20.250000,17.375000,16.125000,15.500000,15.187500 +78 0 304 0.435938 304,662,1174,3010,389,13,279,439 25.125000,22.625000,22.375000,22.375000,22.125000,21.875000,21.625000,21.625000 +79 0 6841 3.264097 220,279,5651,5587,6664,6841,3297,7552 24.625000,22.875000,22.500000,22.000000,22.000000,22.000000,21.875000,21.625000 +80 0 220 0.401440 220,315,430,662,13,279,1174,323 25.125000,24.000000,22.125000,22.000000,21.125000,21.000000,20.750000,20.500000 +81 0 1049 0.007944 1049,679,220,2550,17,508,2366,23 24.000000,19.125000,15.312500,13.437500,13.250000,13.062500,12.500000,11.562500 +82 0 23 0.048805 23,24,22,21,15,20,16,2421 24.875000,21.875000,16.125000,13.687500,13.250000,13.062500,12.187500,11.937500 +83 0 662 0.192989 662,13,1174,439,323,2652,6905,16853 25.000000,22.125000,22.125000,21.125000,20.500000,20.500000,20.500000,20.375000 +84 0 1283 0.665136 1283,763,426,720,5414,578,4815,1952 23.375000,22.625000,21.625000,20.875000,20.750000,19.375000,19.250000,18.875000 +85 0 1047 2.652273 1101,574,6476,9922,59335,1047,1243,27212 22.625000,22.500000,22.500000,22.500000,22.500000,21.875000,21.500000,21.500000 +86 0 264 0.260509 264,40500,2678,8810,13073,813,1403,12899 26.375000,23.375000,22.500000,22.500000,22.500000,22.375000,22.375000,22.375000 +87 0 46350 2.233368 40500,3560,46350,8810,2678,5912,6522,12899 23.875000,23.625000,22.625000,22.125000,21.500000,21.000000,20.500000,20.375000 +88 0 3560 0.055948 3560,8810,40500,961,11341,3063,3752,9099 25.375000,22.000000,19.875000,19.375000,19.125000,18.875000,18.625000,18.375000 +89 0 304 1.133781 439,304,389,505,315,5737,369,2391 26.125000,25.875000,25.375000,23.125000,22.625000,22.125000,21.875000,21.875000 +90 0 5899 12.972918 279,49922,578,264,16071,6460,18588,94455 24.875000,19.875000,19.750000,19.625000,19.500000,19.250000,19.250000,19.125000 +91 0 18243 0.118702 18243,571,315,5596,220,24871,4101,961 23.375000,21.125000,17.375000,17.250000,17.125000,17.125000,16.625000,16.500000 +92 0 315 0.153917 315,389,439,304,369,2391,505,5737 27.500000,24.625000,24.625000,24.125000,21.625000,21.625000,21.500000,21.000000 +93 0 279 0.112495 279,578,49922,16071,6460,18588,7290,264 24.875000,21.000000,20.250000,19.625000,19.000000,18.875000,18.500000,18.375000 +94 0 12707 0.474403 12707,4948,18588,2132,8013,94455,13740,8198 23.625000,21.125000,21.000000,20.375000,20.375000,20.125000,19.750000,19.625000 +95 0 4101 0.081318 4101,20156,23160,6275,1332,9977,27883,81466 24.875000,21.125000,20.750000,20.625000,20.125000,19.500000,18.750000,18.500000 +96 0 49922 2.900558 578,16071,49922,14894,11,18449,86550,6460 19.375000,18.625000,18.375000,18.125000,17.750000,17.125000,17.125000,16.875000 +97 0 1919 0.000177 1919,29646,488,2676,14221,6834,10984,89226 29.000000,20.250000,16.625000,15.437500,15.312500,14.937500,14.937500,14.937500 +98 0 304 2.335238 439,1990,304,505,1174,5737,7314,902 23.375000,23.000000,22.375000,22.375000,22.375000,21.875000,21.625000,21.500000 +99 0 220 0.085257 220,279,902,4216,3389,5587,7552,6790 27.125000,24.000000,22.875000,22.125000,21.625000,21.125000,20.875000,20.500000 +100 0 679 1.502308 1049,679,2550,17,508,2366,220,3753 24.625000,23.375000,16.125000,15.187500,14.812500,14.750000,13.312500,13.187500 +101 0 15 0.366486 15,16,17,18,20,19,21,22 23.625000,22.125000,21.250000,20.750000,20.000000,19.875000,19.125000,18.125000 +102 0 1174 0.909258 1174,662,439,323,5737,13,2652,304 24.000000,23.375000,23.125000,22.125000,21.875000,20.375000,20.125000,19.875000 +103 0 439 3.880621 5737,323,8272,439,902,1603,1405,40500 24.625000,23.125000,22.000000,21.250000,21.125000,21.125000,20.875000,20.875000 +104 0 330 4.679625 6812,1664,2999,56379,65644,279,16943,11162 21.375000,21.125000,20.625000,20.125000,20.125000,19.750000,19.750000,19.000000 +105 0 735 6.803286 6812,7762,11517,7043,44632,7505,7957,25459 18.500000,16.875000,16.375000,16.125000,16.125000,16.000000,16.000000,16.000000 +106 0 40608 2.275879 82481,5230,40608,575,16414,1122,1037,5420 17.875000,17.500000,16.375000,14.812500,14.312500,13.750000,13.312500,12.500000 +107 0 69168 5.903984 330,31038,4410,16844,7957,23074,56872,34131 19.125000,17.750000,17.125000,17.125000,17.000000,17.000000,17.000000,16.875000 +108 0 330 0.008001 330,662,1174,6058,3755,1,42448,34921 25.125000,18.875000,18.875000,18.625000,18.500000,16.375000,16.000000,15.812500 +109 0 662 0.357847 662,1174,2652,6905,323,16853,279,264 24.125000,22.625000,21.125000,21.125000,20.125000,19.875000,19.500000,18.875000 +110 0 426 2.528252 1283,763,426,720,5414,578,1115,4815 22.500000,21.500000,20.625000,20.125000,20.125000,19.375000,18.875000,18.875000 +111 0 11206 0.000019 11206,3102,7379,619,1098,13,1656,5049 28.375000,15.312500,15.250000,14.500000,14.500000,14.312500,14.312500,13.812500 +112 0 466 0.000168 466,783,5153,998,2641,977,4253,668 27.375000,17.750000,17.125000,16.125000,16.125000,15.000000,15.000000,14.750000 +113 0 59335 2.037708 6476,59335,9922,574,1101,1047,45877,27212 22.000000,21.625000,21.500000,21.375000,21.000000,20.625000,20.375000,20.250000 +114 0 304 0.826206 304,16662,439,14329,389,449,520,1828 26.000000,25.875000,24.875000,22.625000,21.500000,21.375000,20.750000,19.875000 +115 0 279 0.799534 279,1403,264,2380,578,459,3892,330 23.125000,21.625000,21.375000,20.875000,20.250000,20.125000,20.000000,18.875000 +116 0 220 0.773724 220,1514,6566,4632,4410,34596,12707,3063 23.625000,22.125000,21.500000,21.375000,21.125000,20.875000,20.750000,20.625000 +117 0 679 0.031321 679,1049,2550,3753,17,4468,5162,2366 25.375000,21.375000,19.875000,18.625000,18.000000,17.875000,17.875000,17.625000 +118 0 16 0.604622 16,15,17,18,19,20,21,22 22.375000,21.375000,21.000000,20.125000,19.125000,19.000000,18.125000,17.625000 +119 0 4632 1.605379 34596,4632,1514,4410,6566,12707,5788,13673 22.375000,22.000000,21.625000,21.125000,20.625000,20.625000,20.125000,19.250000 +120 0 25145 10.747141 578,34185,11,362,5348,800,2373,36082 19.625000,18.125000,17.500000,17.500000,17.500000,17.250000,17.250000,17.125000 +121 0 85375 1.255869 85375,72,2172,47228,73241,332,552,5092 17.625000,17.125000,16.875000,16.875000,16.875000,15.000000,15.000000,14.562500 +122 0 15910 3.082453 220,1174,551,25,15910,16662,369,323 23.000000,21.500000,20.875000,20.625000,20.625000,20.500000,20.125000,20.000000 +123 0 555 0.000980 555,323,304,369,3842,7043,8563,7957 29.125000,21.500000,19.625000,18.500000,18.500000,18.125000,18.125000,18.000000 +124 0 12366 10.622049 7957,15341,15469,11517,3842,17520,8563,25459 19.000000,18.625000,18.375000,18.250000,18.000000,18.000000,17.750000,17.750000 +125 0 2009 0.034842 2009,1676,30982,445,22791,323,362,514 22.000000,17.000000,16.375000,16.125000,16.125000,16.000000,15.437500,14.937500 +126 0 546 0.000013 546,6167,519,688,2159,354,64,53377 30.625000,18.125000,17.625000,17.625000,17.500000,16.875000,16.500000,15.562500 +127 0 72 0.000126 72,285,88,288,99355,648,283,6870 27.125000,17.375000,16.125000,16.125000,15.500000,15.437500,13.687500,13.250000 diff --git a/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_ssmkernel0_w1ctx512.tsv b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_ssmkernel0_w1ctx512.tsv new file mode 100644 index 000000000..0fdd5fdcb --- /dev/null +++ b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_granite-4.0-h-tiny_ssmkernel0_w1ctx512.tsv @@ -0,0 +1,135 @@ +# model models/mlx/granite-4.0-h-tiny-4bit +# device GPU +# corpus tests/fixtures/wikitext2_excerpt.txt tokens 289403 +# chunks 128 chunk_tokens 1 topk 8 prefill 512 +# columns chunk pos target nll top_ids top_logits +# start_token 512 +# contiguous false +0 0 323 0.138105 323,612,1174,662,304,13,1405,1990 25.000000,22.375000,21.375000,21.125000,20.000000,19.375000,19.125000,18.875000 +1 0 48550 0.203667 48550,574,6476,520,37778,10455,48043,27616 22.375000,20.125000,19.125000,17.875000,17.500000,17.000000,17.000000,16.875000 +2 0 5721 0.000172 5721,71,1494,4317,13279,8350,6798,65 29.500000,20.750000,16.875000,16.500000,16.500000,16.000000,15.250000,14.812500 +3 0 662 0.899998 662,323,1174,13,304,1990,2652,505 23.625000,22.625000,22.625000,21.625000,21.625000,20.875000,20.500000,20.375000 +4 0 426 2.113909 1283,763,426,720,578,8563,5414,1115 22.500000,21.625000,21.125000,20.000000,19.875000,19.375000,19.250000,18.250000 +5 0 11206 0.000044 11206,619,5049,3102,1656,3514,13,980 27.000000,15.312500,14.562500,14.312500,13.812500,13.812500,13.750000,13.437500 +6 0 466 0.000166 466,5153,783,2641,998,4253,376,977 26.875000,17.000000,16.875000,15.437500,15.250000,14.750000,14.437500,14.437500 +7 0 59335 2.077045 574,1101,6476,59335,9922,1243,45877,1047 22.375000,22.125000,22.125000,22.125000,21.500000,21.125000,20.875000,20.750000 +8 0 304 1.282543 16662,304,439,14329,449,1828,389,1578 25.375000,24.625000,23.625000,21.375000,20.625000,20.125000,19.875000,19.875000 +9 0 1403 4.612274 279,264,420,578,2650,459,3892,1403 23.375000,20.875000,20.250000,20.125000,20.125000,19.625000,19.375000,19.250000 +10 0 12631 2.151197 54248,11335,12631,12707,18243,220,1023,6007 22.875000,22.500000,22.500000,22.500000,22.000000,21.625000,21.000000,20.875000 +11 0 304 0.862793 304,15910,6004,11,430,1174,9124,2391 23.625000,22.625000,22.000000,21.250000,21.125000,21.125000,20.500000,20.250000 +12 0 220 0.086637 220,279,4216,3389,4062,11295,430,902 26.500000,23.500000,22.000000,21.250000,20.625000,20.375000,20.000000,20.000000 +13 0 1049 0.063063 1049,679,2550,2366,508,17,3753,4468 24.625000,21.875000,17.250000,16.000000,15.250000,14.812500,14.562500,14.250000 +14 0 23 0.936684 22,23,24,21,20,19,18,15 22.000000,21.875000,20.625000,19.750000,17.125000,15.250000,14.812500,14.562500 +15 0 1174 1.030802 1174,2652,662,551,13,25,26,2001 23.125000,22.625000,22.125000,21.875000,20.500000,20.125000,20.125000,20.000000 +16 0 6187 8.695541 279,578,5737,15910,832,2225,25891,40500 20.000000,19.875000,19.500000,18.875000,18.125000,18.125000,18.000000,17.625000 +17 0 4238 0.789576 4238,315,10573,369,9137,1712,18811,17586 21.375000,20.125000,20.125000,19.375000,19.125000,18.625000,18.375000,17.750000 +18 0 4997 3.509829 323,320,1174,15371,52483,4997,21603,49538 19.625000,18.625000,18.500000,18.125000,18.125000,17.375000,17.125000,17.125000 +19 0 31304 -0.000000 31304,655,1941,7278,88,7006,5515,727 37.500000,23.000000,22.875000,21.500000,19.625000,18.625000,17.875000,17.500000 +20 0 555 5.902033 323,15910,1174,320,16662,14329,902,5439 24.375000,23.375000,22.500000,22.125000,21.625000,21.125000,20.875000,20.625000 +21 0 57337 6.902946 7957,12131,6941,8563,11291,15469,17520,8096 19.250000,18.750000,18.625000,18.625000,18.625000,18.500000,18.500000,18.375000 +22 0 12366 10.049431 6941,7957,7505,11517,15469,12131,13929,29335 18.375000,18.375000,18.250000,17.750000,17.750000,17.625000,17.625000,17.625000 +23 0 2009 8.207556 69012,362,1676,81431,47790,22791,6938,32817 18.750000,16.375000,15.250000,15.250000,14.812500,14.687500,14.000000,13.750000 +24 0 546 4.434725 288,546,688,613,589,64,1441,81378 20.625000,16.250000,16.250000,16.000000,15.000000,14.812500,14.500000,14.500000 +25 0 72 0.000059 72,64044,285,288,437,88,99355,290 30.375000,19.875000,19.125000,18.000000,17.875000,17.625000,17.625000,17.125000 +26 0 1174 1.059917 323,1174,449,40500,304,320,662,2652 24.375000,23.875000,20.875000,20.500000,20.375000,20.375000,19.500000,19.250000 +27 0 323 0.679844 323,902,304,889,15910,1405,369,8272 24.625000,24.125000,22.125000,21.500000,21.375000,21.125000,20.500000,20.375000 +28 0 4418 8.536655 578,279,304,362,52483,800,3842,16290 20.875000,20.500000,20.125000,18.500000,18.500000,17.875000,17.625000,17.625000 +29 0 798 2.094664 956,798,3489,3919,90970,364,29868,12565 17.000000,15.750000,14.500000,14.250000,14.062500,14.000000,13.812500,13.750000 +30 0 54275 0.139969 54275,96363,79,445,21004,11617,596,364 21.125000,17.375000,17.125000,16.875000,16.625000,16.375000,16.125000,15.937500 +31 0 15910 1.052723 15910,555,1174,662,902,320,13,11 23.500000,23.000000,22.500000,21.500000,21.500000,21.375000,21.250000,21.125000 +32 0 555 0.001820 555,323,1578,304,520,7043,3131,7957 28.875000,21.875000,20.125000,19.375000,19.375000,18.750000,18.375000,18.250000 +33 0 12225 6.235201 7505,54424,30411,11291,16768,18984,19475,31421 18.000000,15.562500,14.812500,14.750000,14.750000,14.312500,14.250000,14.187500 +34 0 398 0.496895 398,29306,6885,5657,2642,402,6246,23797 20.125000,19.625000,15.062500,15.000000,14.187500,14.000000,13.562500,13.500000 +35 0 89453 2.098206 63592,5348,32103,89453,35934,13566,23162,12404 16.750000,16.000000,15.750000,15.750000,15.062500,14.250000,13.812500,13.562500 +36 0 662 0.399191 662,1174,323,13,2652,449,16853,320 23.625000,21.625000,21.500000,21.125000,20.625000,19.375000,19.375000,19.250000 +37 0 763 2.172619 1283,763,426,578,11995,720,5414,6187 22.875000,21.250000,20.875000,20.375000,20.250000,19.875000,19.625000,19.125000 +38 0 3297 3.853315 220,2225,279,6187,5369,3297,6287,5651 23.875000,22.500000,22.125000,21.875000,20.750000,20.625000,19.875000,19.500000 +39 0 220 0.177145 220,315,1174,11,323,568,430,279 24.375000,21.875000,21.875000,19.500000,19.375000,18.875000,18.375000,18.250000 +40 0 1049 0.203182 1049,679,2366,2550,508,220,3753,17 23.500000,22.000000,16.875000,15.312500,13.312500,13.250000,13.187500,13.062500 +41 0 23 0.583931 23,24,22,21,20,19,18,15 23.375000,23.125000,18.625000,17.625000,16.875000,15.312500,14.937500,14.687500 +42 0 1174 0.744474 1174,568,426,11,279,8563,578,433 24.000000,23.500000,22.375000,22.125000,21.125000,19.625000,19.000000,18.875000 +43 0 426 0.914038 568,426,279,8563,578,433,12360,813 23.500000,23.375000,21.125000,20.375000,19.875000,18.875000,18.750000,18.625000 +44 0 11206 0.000027 11206,619,1656,980,3102,3514,5049,1286 28.500000,17.375000,15.000000,14.562500,14.562500,14.562500,14.250000,13.750000 +45 0 466 0.000242 466,783,998,668,5153,655,2002,2641 27.500000,18.875000,16.000000,15.750000,15.312500,15.250000,15.000000,15.000000 +46 0 1903 3.534035 574,9922,59335,6476,10887,1903,4036,11096 22.625000,22.500000,22.500000,22.000000,21.625000,20.750000,20.625000,20.625000 +47 0 264 2.153769 813,264,459,3925,2500,279,31186,12707 26.500000,24.625000,24.500000,21.125000,21.000000,20.875000,20.875000,20.750000 +48 0 8810 0.447971 8810,471,3361,10015,6566,9024,17755,95156 24.625000,23.125000,21.625000,21.500000,21.125000,20.875000,20.625000,20.625000 +49 0 11341 0.068512 11341,571,6917,40500,17755,471,7858,5178 21.500000,18.000000,17.000000,16.750000,15.687500,15.437500,15.000000,14.750000 +50 0 389 0.359701 389,304,439,520,369,5737,16662,315 27.375000,26.000000,25.375000,23.625000,22.375000,22.125000,21.625000,20.500000 +51 0 264 3.836206 279,459,264,18588,6566,578,12707,6460 23.625000,20.625000,20.000000,19.625000,19.375000,19.125000,18.875000,18.750000 +52 0 1403 5.823925 12707,220,4101,18588,18629,3361,6007,9417 21.500000,20.375000,19.125000,19.125000,19.000000,18.875000,18.625000,18.500000 +53 0 571 4.855267 29137,482,961,12,9417,571,58098,21925 21.000000,20.625000,20.375000,19.750000,18.375000,17.125000,16.500000,16.125000 +54 0 12 0.002274 12,59562,35394,39743,12898,28751,52656,51442 26.375000,20.125000,16.875000,15.687500,15.437500,15.437500,15.250000,15.187500 +55 0 31 0.000038 31,19908,571,9164,19741,555,63899,5 28.125000,16.500000,15.687500,15.437500,15.062500,14.500000,14.250000,14.187500 +56 0 961 0.010572 961,9417,1370,6596,2046,1938,9568,3814 29.000000,23.875000,23.000000,22.000000,21.500000,20.875000,20.000000,19.750000 +57 0 9417 0.282137 9417,12707,3361,6007,4101,18588,3446,8013 23.125000,21.250000,19.625000,19.375000,19.125000,18.625000,18.375000,18.125000 +58 0 15952 14.392565 315,330,25891,320,304,20458,11,1174 27.375000,21.875000,21.625000,21.375000,20.875000,20.375000,20.125000,20.125000 +59 0 315 0.054888 315,389,304,25891,369,320,20458,330 26.375000,22.125000,21.875000,21.625000,21.125000,20.375000,20.250000,20.125000 +60 0 279 0.145798 279,578,6460,16071,49922,18588,19150,54943 23.625000,19.875000,19.375000,19.250000,19.000000,18.875000,18.375000,18.000000 +61 0 12707 0.431559 12707,18588,8013,5526,13673,94455,13740,4101 23.375000,21.250000,20.750000,20.125000,20.125000,19.625000,19.375000,19.125000 +62 0 4101 0.042583 4101,6275,20156,1501,9977,23160,27883,18629 24.875000,20.250000,20.125000,19.375000,19.250000,19.125000,19.125000,18.750000 +63 0 468 6.607880 578,11,16071,1174,49922,3165,6460,86550 18.875000,18.125000,18.125000,17.875000,17.625000,17.000000,17.000000,16.875000 +64 0 1802 1.623247 6310,1802,16,1721,1924,55889,26286,19011 18.625000,17.375000,15.937500,14.000000,14.000000,13.812500,13.562500,13.000000 +65 0 279 0.048759 279,578,791,11,220,16454,1820,3146 27.500000,24.500000,17.375000,16.125000,15.812500,15.562500,15.500000,15.437500 +66 0 15371 0.000093 15371,5710,29420,84088,79169,81834,33160,83942 31.125000,21.500000,20.375000,18.250000,16.875000,16.750000,16.625000,15.437500 +67 0 1174 1.006168 1174,662,439,304,2652,902,323,1405 22.875000,22.500000,21.375000,21.125000,20.875000,20.375000,19.875000,19.625000 +68 0 8272 5.314432 902,5737,25891,304,15910,323,1405,5439 23.500000,23.500000,22.625000,22.125000,22.125000,21.875000,21.625000,21.125000 +69 0 555 0.054171 555,304,3010,279,20193,449,264,1828 29.000000,25.625000,24.625000,22.625000,22.000000,21.875000,21.625000,21.375000 +70 0 459 3.168373 264,459,40500,1403,2500,27351,813,8810 25.500000,22.625000,22.500000,22.375000,21.875000,21.500000,21.375000,21.375000 +71 0 11341 0.119780 11341,9417,8223,14529,15718,11838,220,653 23.500000,20.625000,19.375000,18.000000,18.000000,17.750000,17.625000,17.625000 +72 0 389 0.577129 389,304,439,3010,279,520,14329,16662 27.500000,27.000000,25.500000,22.875000,22.125000,22.125000,22.125000,21.500000 +73 0 279 0.936365 279,264,459,578,49922,16071,2500,52483 23.125000,23.000000,21.125000,20.625000,20.375000,19.750000,19.375000,19.250000 +74 0 12707 0.293806 12707,18588,8013,220,4948,4101,9417,23160 24.000000,20.625000,20.250000,20.125000,19.875000,19.625000,19.625000,19.375000 +75 0 4101 0.041119 4101,1332,6275,20156,18629,2068,1501,23160 25.125000,19.875000,19.875000,19.875000,19.500000,19.375000,19.250000,19.125000 +76 0 29878 9.077332 16071,49922,578,86550,14894,11,6460,11301 19.500000,19.250000,19.125000,18.250000,17.875000,17.625000,17.375000,17.125000 +77 0 86493 0.137547 86493,2299,535,344,6632,2270,6672,10126 27.500000,25.500000,23.000000,20.500000,17.625000,16.375000,15.312500,15.187500 +78 0 304 0.460680 304,662,1174,3010,389,13,439,5737 25.125000,22.875000,22.375000,22.375000,22.125000,22.000000,21.625000,21.500000 +79 0 6841 3.285042 220,279,5651,5587,3297,6664,6841,6790 24.500000,22.625000,22.500000,22.000000,21.875000,21.875000,21.875000,21.625000 +80 0 220 0.434717 220,315,430,662,13,279,1174,323 25.125000,24.125000,22.125000,22.125000,21.125000,21.000000,20.750000,20.625000 +81 0 1049 0.007975 1049,679,220,2550,17,508,2366,23 24.000000,19.125000,15.437500,13.500000,13.250000,13.062500,12.500000,11.562500 +82 0 23 0.048866 23,24,22,21,15,20,16,2421 24.625000,21.625000,16.125000,13.687500,13.250000,13.062500,12.187500,11.937500 +83 0 662 0.175493 662,13,1174,439,2652,323,6905,16853 25.125000,22.125000,22.125000,21.250000,20.625000,20.500000,20.500000,20.375000 +84 0 1283 0.648254 1283,763,426,720,5414,578,4815,1952 23.375000,22.625000,21.500000,20.875000,20.625000,19.375000,19.125000,18.875000 +85 0 1047 2.844477 574,1101,6476,9922,59335,1047,1243,27212 22.500000,22.500000,22.500000,22.500000,22.375000,21.625000,21.500000,21.375000 +86 0 264 0.246944 264,40500,2678,8810,13073,459,813,1403 26.375000,23.125000,22.500000,22.500000,22.500000,22.375000,22.375000,22.375000 +87 0 46350 2.238781 40500,3560,46350,8810,2678,5912,6522,12899 23.875000,23.625000,22.625000,22.125000,21.625000,21.125000,20.500000,20.375000 +88 0 3560 0.066244 3560,8810,40500,961,11341,3063,3752,9099 25.125000,21.875000,19.875000,19.250000,19.125000,18.875000,18.625000,18.375000 +89 0 304 1.274349 439,304,389,505,315,5737,369,2391 26.375000,25.875000,25.500000,23.125000,22.875000,22.375000,22.000000,21.875000 +90 0 5899 12.913292 279,49922,578,264,16071,18588,6460,94455 24.875000,19.875000,19.750000,19.625000,19.500000,19.375000,19.250000,19.250000 +91 0 18243 0.119099 18243,571,315,5596,24871,220,4101,961 23.375000,21.125000,17.375000,17.250000,17.250000,17.000000,16.750000,16.500000 +92 0 315 0.154236 315,389,439,304,369,2391,505,5737 27.500000,24.625000,24.625000,24.125000,21.625000,21.625000,21.500000,21.000000 +93 0 279 0.142405 279,578,49922,16071,6460,18588,7290,19150 24.625000,21.000000,20.250000,19.625000,19.000000,18.875000,18.500000,18.375000 +94 0 12707 0.470779 12707,4948,18588,2132,8013,94455,13740,8198 23.625000,21.125000,20.875000,20.375000,20.375000,20.125000,19.750000,19.625000 +95 0 4101 0.075832 4101,20156,23160,6275,1332,9977,27883,81466 24.875000,21.000000,20.750000,20.625000,19.875000,19.375000,18.625000,18.500000 +96 0 49922 2.871273 578,16071,49922,14894,11,18449,86550,6460 19.375000,18.625000,18.375000,18.125000,17.750000,17.125000,17.125000,16.875000 +97 0 1919 0.000198 1919,29646,488,2676,6834,10984,14221,89226 29.000000,20.375000,16.625000,15.500000,15.187500,15.187500,15.187500,14.687500 +98 0 304 2.239655 439,1990,304,505,1174,5737,7314,902 23.125000,23.000000,22.375000,22.375000,22.125000,21.875000,21.625000,21.375000 +99 0 220 0.088287 220,279,902,4216,3389,5587,7552,6790 27.000000,23.875000,22.875000,22.000000,21.500000,21.000000,20.750000,20.375000 +100 0 679 1.502298 1049,679,2550,17,2366,508,220,3753 24.625000,23.375000,16.125000,15.062500,14.812500,14.750000,13.250000,13.187500 +101 0 15 0.337111 15,16,17,18,19,20,21,22 23.625000,22.000000,21.125000,20.750000,19.875000,19.875000,19.000000,18.125000 +102 0 1174 0.879066 1174,662,439,323,5737,13,2652,12 23.875000,23.125000,23.000000,22.000000,21.625000,20.375000,20.125000,19.750000 +103 0 439 3.766798 5737,323,8272,439,902,1603,1405,26449 24.625000,23.125000,22.000000,21.375000,21.250000,21.125000,20.875000,20.875000 +104 0 330 4.718849 6812,1664,2999,56379,65644,16943,279,11162 21.500000,21.125000,20.625000,20.125000,20.125000,19.875000,19.750000,19.125000 +105 0 735 6.660444 6812,7762,11517,2999,7505,7957,44632,7043 18.375000,16.750000,16.250000,16.000000,16.000000,16.000000,16.000000,15.937500 +106 0 40608 2.236141 82481,5230,40608,575,16414,1122,1037,1631 17.875000,17.375000,16.375000,14.750000,14.250000,13.687500,13.312500,12.500000 +107 0 69168 5.890888 330,31038,16844,4410,56872,23074,30563,34131 18.875000,17.750000,17.125000,17.000000,17.000000,16.875000,16.875000,16.875000 +108 0 330 0.008997 330,662,1174,6058,3755,1,42448,34921 24.875000,18.625000,18.625000,18.625000,18.500000,16.250000,16.000000,15.750000 +109 0 662 0.354065 662,1174,2652,6905,323,16853,279,264 24.125000,22.625000,21.125000,21.125000,20.125000,19.750000,19.375000,18.875000 +110 0 426 2.529295 1283,763,426,720,5414,578,1115,4815 22.500000,21.500000,20.625000,20.125000,20.125000,19.375000,18.875000,18.875000 +111 0 11206 0.000019 11206,3102,7379,619,1098,1656,13,3514 28.375000,15.312500,15.250000,14.500000,14.500000,14.437500,14.250000,14.062500 +112 0 466 0.000162 466,783,5153,998,2641,977,4253,668 27.375000,17.625000,17.125000,16.250000,16.125000,15.000000,15.000000,14.687500 +113 0 59335 2.044844 6476,59335,9922,574,1101,1047,45877,27212 22.000000,21.625000,21.500000,21.375000,21.000000,20.625000,20.500000,20.250000 +114 0 304 0.737911 304,16662,439,14329,389,449,520,1828 26.000000,25.625000,24.875000,22.625000,21.500000,21.500000,20.625000,20.125000 +115 0 279 0.797514 279,1403,264,2380,578,459,3892,220 23.125000,21.625000,21.500000,20.875000,20.250000,20.125000,20.000000,19.000000 +116 0 220 0.781500 220,1514,6566,4632,4410,12707,34596,3063 23.625000,22.125000,21.500000,21.375000,21.125000,20.875000,20.875000,20.625000 +117 0 679 0.034386 679,1049,2550,3753,17,4468,5162,2366 25.375000,21.500000,19.875000,18.625000,18.125000,17.875000,17.875000,17.750000 +118 0 16 0.607916 16,15,17,18,19,20,21,22 22.500000,21.500000,21.125000,20.250000,19.375000,19.125000,18.375000,17.625000 +119 0 4632 1.662218 34596,4632,1514,4410,6566,12707,5788,13673 22.375000,22.000000,21.875000,21.125000,20.750000,20.625000,20.125000,19.375000 +120 0 25145 10.788450 578,34185,11,362,5348,800,2373,52483 19.625000,17.875000,17.625000,17.625000,17.500000,17.375000,17.250000,17.125000 +121 0 85375 1.224890 85375,72,47228,73241,2172,552,332,519 17.750000,17.125000,17.125000,17.000000,16.875000,15.000000,14.937500,14.562500 +122 0 15910 3.535872 220,1174,16662,25,551,369,15910,323 22.875000,21.500000,20.875000,20.625000,20.500000,20.125000,20.125000,20.000000 +123 0 555 0.001101 555,323,304,369,3842,7043,7957,8563 29.000000,21.500000,19.625000,18.500000,18.375000,18.125000,18.000000,18.000000 +124 0 12366 10.474194 7957,15341,15469,11517,3842,17520,11291,25459 19.125000,18.750000,18.375000,18.125000,17.875000,17.875000,17.750000,17.750000 +125 0 2009 0.039101 2009,1676,30982,445,22791,323,362,514 21.875000,17.125000,16.250000,16.125000,16.125000,15.937500,15.500000,15.000000 +126 0 546 0.000011 546,6167,519,688,2159,354,64,53377 30.875000,18.125000,17.625000,17.625000,17.625000,16.875000,16.625000,15.750000 +127 0 72 0.000132 72,285,88,288,99355,648,283,6870 27.125000,17.500000,16.125000,16.125000,15.562500,15.250000,13.500000,13.437500 diff --git a/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1.tsv b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1.tsv new file mode 100644 index 000000000..461e2f3ce --- /dev/null +++ b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1.tsv @@ -0,0 +1,140 @@ +[NemotronH] Loading config... +[NemotronH] Config loaded: 52 layers (23 mamba, 6 attention, 0 mlp, 23 moe) +[NemotronH] Loading weights from safetensors... +[NemotronH] Building model... +[NemotronH] Model loaded successfully +# model models/mlx/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit +# device GPU +# corpus tests/fixtures/wikitext2_excerpt.txt tokens 301749 +# chunks 128 chunk_tokens 1 topk 8 prefill 0 +# columns chunk pos target nll top_ids top_logits +# start_token 0 +# contiguous false +0 0 20211 14.692390 59199,1256,1032,2251,1603,1319,1050,1317 12.687500,12.625000,12.250000,12.062500,11.937500,11.875000,11.812500,11.562500 +1 0 1752 12.361443 109431,1044,1046,1034,2251,59199,1041,1338 14.187500,13.187500,13.187500,12.937500,12.625000,12.500000,12.125000,12.062500 +2 0 1637 13.157224 1044,109431,1046,1034,59199,1338,1048,1041 12.937500,12.687500,12.625000,12.312500,12.250000,11.750000,11.687500,11.625000 +3 0 1045 4.639105 109431,1044,59199,1034,1046,1048,1319,1053 14.000000,13.312500,13.187500,12.750000,12.750000,12.687500,12.062500,12.000000 +4 0 1050 4.206770 1256,59199,1438,1032,11,1603,2251,1050 12.937500,12.687500,12.062500,11.937500,11.875000,11.812500,11.812500,11.750000 +5 0 7637 9.865001 1256,1034,59199,1438,1044,11,1338,109431 12.312500,12.250000,12.250000,12.187500,11.812500,11.750000,11.750000,11.750000 +6 0 103920 11.968136 109431,1044,1046,1034,1317,1338,2251,59199 13.125000,12.812500,12.750000,12.687500,12.250000,12.000000,11.812500,11.812500 +7 0 1376 7.997871 1044,1034,109431,1046,59199,1319,1438,2251 13.937500,13.562500,13.437500,13.375000,13.062500,12.750000,12.625000,12.562500 +8 0 1267 7.165167 1256,59199,1032,1603,2251,1319,1364,1050 12.625000,12.437500,12.250000,12.062500,11.750000,11.625000,11.562500,11.250000 +9 0 8263 14.221610 59199,1034,1256,1438,1032,2251,1044,1338 12.437500,12.312500,12.312500,12.250000,11.750000,11.687500,11.562500,11.562500 +10 0 1260 10.358089 1034,109431,1044,59199,1046,1438,1042,1319 13.937500,13.812500,13.562500,13.125000,13.062500,13.000000,12.437500,12.437500 +11 0 1737 8.949326 1034,1256,1438,11,59199,1032,1044,1042 12.937500,12.750000,12.750000,12.500000,12.312500,12.187500,12.125000,11.937500 +12 0 7155 11.079857 1256,59199,1032,1364,11,1603,2251,1429 13.375000,13.250000,12.875000,12.562500,12.500000,12.500000,12.250000,12.000000 +13 0 2345 12.569784 1044,1046,1034,1338,1317,2251,109431,1438 13.250000,13.187500,12.875000,12.625000,12.562500,12.437500,12.375000,11.937500 +14 0 1104 9.137060 1034,1046,109431,1044,2251,1438,59199,1338 13.312500,13.000000,12.625000,12.562500,12.500000,12.312500,12.312500,12.250000 +15 0 1936 11.194901 1034,1044,1438,1046,59199,109431,1338,2251 13.062500,12.437500,12.250000,11.937500,11.937500,11.937500,11.750000,11.750000 +16 0 7366 12.676600 109431,1044,1046,59199,1034,1041,1338,2251 13.750000,13.312500,13.000000,13.000000,12.875000,12.000000,12.000000,12.000000 +17 0 3136 15.951866 1034,1044,1278,1046,1261,1317,1438,1042 12.937500,12.812500,12.812500,12.375000,12.375000,12.187500,12.187500,11.937500 +18 0 20524 12.136600 1034,109431,1044,1046,1438,1042,59199,1041 13.625000,13.625000,13.250000,12.687500,12.500000,12.437500,12.000000,11.875000 +19 0 6406 9.213186 1034,109431,1044,1438,2251,1046,59199,1414 13.937500,13.875000,13.312500,12.750000,12.750000,12.687500,12.687500,12.250000 +20 0 71129 16.441151 1034,109431,1046,1044,1317,1438,1338,59199 12.937500,12.937500,12.750000,12.375000,12.250000,12.062500,12.000000,11.812500 +21 0 14327 11.122322 1044,1034,1438,2251,1046,1338,1317,1319 13.500000,13.375000,12.937500,12.687500,12.625000,12.062500,12.000000,11.875000 +22 0 3289 13.973761 109431,1046,1034,1044,59199,1338,1438,1319 13.437500,13.250000,13.062500,12.875000,12.500000,12.250000,12.125000,12.000000 +23 0 17271 16.483624 1044,1046,109431,59199,1034,1048,1319,1338 13.125000,12.625000,12.500000,12.312500,12.187500,11.875000,11.750000,11.750000 +24 0 13001 13.272747 1034,109431,1044,1438,1046,1414,1338,59199 13.750000,13.562500,13.187500,12.875000,12.750000,12.062500,12.000000,11.937500 +25 0 1552 12.855757 1034,109431,1046,1044,1438,1317,1338,59199 13.500000,13.437500,13.000000,12.937500,12.312500,12.250000,12.187500,12.125000 +26 0 1752 13.207749 109431,1044,1046,59199,1034,1319,1338,1438 13.437500,13.187500,12.875000,12.687500,12.562500,12.000000,11.937500,11.812500 +27 0 1637 13.157224 1044,109431,1046,1034,59199,1338,1048,1041 12.937500,12.687500,12.625000,12.312500,12.250000,11.750000,11.687500,11.625000 +28 0 1010 5.701605 109431,1044,59199,1034,1046,1048,1319,1053 14.000000,13.312500,13.187500,12.750000,12.750000,12.687500,12.062500,12.000000 +29 0 7324 13.366667 1256,59199,11,1032,1438,1050,2251,1603 13.375000,12.750000,12.625000,12.312500,12.250000,11.812500,11.812500,11.750000 +30 0 2683 11.947104 109431,1034,59199,1046,1044,1438,1317,1319 14.812500,13.500000,13.062500,12.875000,12.750000,12.562500,12.500000,12.125000 +31 0 1260 10.168397 1034,1044,109431,1046,1317,1438,1338,1307 12.812500,12.562500,12.562500,12.125000,12.062500,12.062500,11.750000,11.687500 +32 0 1737 8.949326 1034,1256,1438,11,59199,1032,1044,1042 12.937500,12.750000,12.750000,12.500000,12.312500,12.187500,12.125000,11.937500 +33 0 1285 10.173607 1256,59199,1032,1364,11,1603,2251,1429 13.375000,13.250000,12.875000,12.562500,12.500000,12.500000,12.250000,12.000000 +34 0 1552 12.493006 109431,1034,1044,1317,1046,1338,1408,1438 13.062500,13.000000,12.687500,12.625000,12.437500,12.250000,11.812500,11.812500 +35 0 1752 13.207749 109431,1044,1046,59199,1034,1319,1338,1438 13.437500,13.187500,12.875000,12.687500,12.562500,12.000000,11.937500,11.812500 +36 0 1637 13.157224 1044,109431,1046,1034,59199,1338,1048,1041 12.937500,12.687500,12.625000,12.312500,12.250000,11.750000,11.687500,11.625000 +37 0 1045 4.639105 109431,1044,59199,1034,1046,1048,1319,1053 14.000000,13.312500,13.187500,12.750000,12.750000,12.687500,12.062500,12.000000 +38 0 1050 4.206770 1256,59199,1438,1032,11,1603,2251,1050 12.937500,12.687500,12.062500,11.937500,11.875000,11.812500,11.812500,11.750000 +39 0 1045 6.146251 1256,1034,59199,1438,1044,11,1338,109431 12.312500,12.250000,12.250000,12.187500,11.812500,11.750000,11.750000,11.750000 +40 0 3416 13.863020 1256,59199,1438,1032,11,1603,2251,1050 12.937500,12.687500,12.062500,11.937500,11.875000,11.812500,11.812500,11.750000 +41 0 9332 11.203976 109431,1034,1044,1046,1438,59199,1278,1041 13.937500,13.250000,13.000000,12.625000,12.187500,12.187500,12.125000,11.937500 +42 0 1049 5.237698 109431,1034,1044,1046,1438,59199,1338,1041 13.562500,13.187500,12.812500,12.687500,12.125000,12.125000,11.937500,11.875000 +43 0 1010 4.959213 2251,1034,1256,1338,1317,1044,1319,59199 12.187500,11.875000,11.875000,11.875000,11.812500,11.750000,11.687500,11.625000 +44 0 44149 14.210417 1256,59199,11,1032,1438,1050,2251,1603 13.375000,12.750000,12.625000,12.312500,12.250000,11.812500,11.812500,11.750000 +45 0 1652 9.741051 109431,1034,1044,59199,1046,1317,1438,1414 14.750000,13.875000,13.187500,13.187500,12.937500,12.875000,12.812500,12.562500 +46 0 1737 9.188922 1034,59199,1438,1256,2251,11,1032,1317 12.812500,12.500000,12.375000,12.125000,12.000000,11.875000,11.875000,11.750000 +47 0 2688 9.423607 1256,59199,1032,1364,11,1603,2251,1429 13.375000,13.250000,12.875000,12.562500,12.500000,12.500000,12.250000,12.000000 +48 0 1319 4.287510 109431,1044,1046,1034,1338,1438,1319,59199 13.187500,13.062500,12.937500,12.312500,12.250000,11.625000,11.500000,11.500000 +49 0 1052 6.399452 1256,59199,1034,1438,2251,1032,1317,1603 12.375000,12.125000,11.625000,11.625000,11.562500,11.500000,11.250000,11.250000 +50 0 1051 5.820013 1034,109431,1044,1046,59199,1438,1338,1319 13.062500,12.500000,12.437500,12.312500,12.250000,12.187500,11.937500,11.625000 +51 0 1053 5.632983 1034,1044,1046,109431,1438,59199,1338,1319 12.812500,12.437500,12.375000,12.375000,12.125000,12.062500,11.937500,11.687500 +52 0 1056 6.208900 1034,109431,1044,1046,59199,1438,1338,1319 12.875000,12.500000,12.437500,12.375000,12.250000,12.125000,11.812500,11.625000 +53 0 12964 8.176165 1034,1044,109431,1046,59199,1438,1338,1319 12.750000,12.312500,12.312500,12.187500,12.187500,12.000000,11.750000,11.562500 +54 0 1456 6.738632 1044,1046,1338,1034,109431,1041,1319,1438 13.437500,12.875000,12.375000,12.312500,12.250000,11.875000,11.875000,11.750000 +55 0 93552 13.887463 1256,59199,11,1438,1032,1034,1050,2251 13.750000,12.875000,12.812500,12.625000,12.562500,12.250000,12.062500,12.062500 +56 0 1293 9.793375 1034,109431,1044,1046,59199,1438,1414,1041 14.187500,14.125000,13.750000,13.750000,13.062500,12.875000,12.625000,12.500000 +57 0 1737 9.046120 59199,1256,1034,1438,1032,11,2251,1317 12.750000,12.500000,12.250000,12.187500,12.000000,11.875000,11.812500,11.687500 +58 0 25079 11.454857 1256,59199,1032,1364,11,1603,2251,1429 13.375000,13.250000,12.875000,12.562500,12.500000,12.500000,12.250000,12.000000 +59 0 20992 13.437098 109431,1034,1044,59199,1046,1438,2251,1317 13.500000,13.125000,13.125000,12.937500,12.750000,12.750000,12.625000,12.250000 +60 0 5624 9.491405 1034,1044,1046,109431,1438,1338,1317,1278 13.062500,12.750000,12.625000,12.625000,12.437500,12.250000,12.125000,11.812500 +61 0 1065 7.659161 109431,1034,59199,1044,1046,1438,1319,1050 13.312500,13.000000,12.562500,12.500000,12.500000,12.375000,11.875000,11.687500 +62 0 1032 4.489594 1034,1438,59199,1044,1046,1319,109431,1338 12.625000,12.375000,12.187500,12.125000,11.812500,11.687500,11.687500,11.625000 +63 0 1051 4.946739 1256,59199,1032,1603,2251,1050,11,1319 13.500000,12.687500,12.437500,12.250000,12.250000,12.125000,12.000000,11.937500 +64 0 1046 3.320483 1034,1044,1046,109431,1438,59199,1338,1319 12.812500,12.437500,12.375000,12.375000,12.125000,12.062500,11.937500,11.687500 +65 0 1048 5.284225 1256,59199,1032,1034,1438,2251,1603,1319 12.562500,12.250000,11.812500,11.750000,11.750000,11.687500,11.562500,11.375000 +66 0 1319 4.117126 1034,1438,1044,109431,59199,1317,1338,1046 12.625000,12.125000,12.000000,11.750000,11.687500,11.625000,11.625000,11.562500 +67 0 3299 13.907265 1256,59199,1034,1438,2251,1032,1317,1603 12.375000,12.125000,11.625000,11.625000,11.562500,11.500000,11.250000,11.250000 +68 0 2345 11.959261 1034,1044,1046,109431,1438,1042,59199,2251 13.687500,13.437500,13.375000,13.250000,13.062500,12.625000,12.625000,12.500000 +69 0 1860 15.562841 1034,1046,109431,1044,2251,1438,59199,1338 13.312500,13.000000,12.625000,12.562500,12.500000,12.312500,12.312500,12.250000 +70 0 2277 12.313381 109431,1034,1044,1046,59199,1438,1319,1338 13.875000,13.562500,13.062500,12.937500,12.750000,12.437500,12.062500,11.937500 +71 0 1730 12.178952 1044,1034,1046,109431,1338,1438,1317,1319 12.812500,12.562500,12.500000,12.125000,11.875000,11.562500,11.312500,11.312500 +72 0 15807 14.816801 109431,1034,1044,1046,59199,1438,1041,1317 14.125000,13.312500,12.937500,12.937500,12.750000,12.437500,12.250000,12.062500 +73 0 6766 10.781712 1044,1034,1046,109431,1438,1338,2251,1042 13.562500,13.250000,13.187500,12.562500,12.500000,12.250000,12.250000,12.187500 +74 0 1047 4.609680 1044,1046,109431,59199,1034,1041,1048,1319 13.937500,13.562500,13.187500,13.125000,13.062500,12.625000,12.562500,12.500000 +75 0 1676 11.387224 1256,59199,1032,2251,1319,1050,1603,11 12.875000,12.500000,12.000000,11.937500,11.875000,11.812500,11.687500,11.500000 +76 0 12204 15.724436 1044,109431,1046,59199,1048,1053,1034,1319 13.437500,13.125000,12.875000,12.687500,12.562500,12.562500,12.375000,12.187500 +77 0 1047 6.500931 1034,1044,1046,1438,1317,1338,1042,1278 12.375000,12.187500,11.875000,11.750000,11.437500,11.437500,11.375000,11.375000 +78 0 3326 12.293474 1256,59199,1032,2251,1319,1050,1603,11 12.875000,12.500000,12.000000,11.937500,11.875000,11.812500,11.687500,11.500000 +79 0 2553 7.707102 109431,1034,1046,1044,1338,1278,1438,59199 12.937500,12.812500,12.625000,12.500000,11.875000,11.812500,11.812500,11.750000 +80 0 1097 9.519714 1034,109431,1044,1046,1438,59199,1338,1042 13.187500,12.937500,12.312500,12.312500,12.312500,12.000000,11.875000,11.750000 +81 0 1047 6.238598 1034,1044,1046,1438,1338,109431,1319,59199 12.437500,12.250000,11.875000,11.875000,11.687500,11.625000,11.562500,11.500000 +82 0 1051 5.012224 1256,59199,1032,2251,1319,1050,1603,11 12.875000,12.500000,12.000000,11.937500,11.875000,11.812500,11.687500,11.500000 +83 0 1046 3.320483 1034,1044,1046,109431,1438,59199,1338,1319 12.812500,12.437500,12.375000,12.375000,12.125000,12.062500,11.937500,11.687500 +84 0 1048 5.284225 1256,59199,1032,1034,1438,2251,1603,1319 12.562500,12.250000,11.812500,11.750000,11.750000,11.687500,11.562500,11.375000 +85 0 85501 14.124938 1034,1438,1044,109431,59199,1317,1338,1046 12.625000,12.125000,12.000000,11.750000,11.687500,11.625000,11.625000,11.562500 +86 0 5780 13.242945 1034,1438,59199,1042,1046,2251,1044,109431 13.375000,12.937500,12.562500,12.437500,12.375000,12.312500,12.250000,11.875000 +87 0 5604 15.629936 1034,109431,59199,1317,1438,1044,1046,2251 13.937500,13.937500,13.312500,13.187500,13.062500,12.937500,12.562500,12.437500 +88 0 1058 4.389065 1044,1046,109431,1034,1438,1041,59199,1338 13.125000,12.750000,12.750000,12.562500,12.125000,12.000000,12.000000,11.937500 +89 0 40377 12.184471 1256,59199,1032,1603,1364,2251,11,1429 13.437500,13.000000,12.750000,12.687500,12.187500,12.125000,11.937500,11.937500 +90 0 79858 14.404896 1044,1046,1034,2251,1338,109431,1317,1438 13.500000,13.375000,12.937500,12.562500,12.500000,12.375000,11.875000,11.875000 +91 0 1044 2.319795 1044,1046,1034,1338,2251,1317,1438,1041 13.500000,13.250000,12.687500,12.687500,12.125000,12.000000,11.875000,11.812500 +92 0 37910 9.320695 1256,59199,1032,1034,1438,1603,2251,1364 12.062500,11.812500,11.437500,11.375000,11.375000,11.312500,11.312500,11.125000 +93 0 17271 14.020786 109431,1044,1034,1046,1438,2251,59199,1317 13.500000,13.437500,13.250000,12.875000,12.812500,12.812500,12.625000,12.312500 +94 0 26554 8.053997 1034,109431,1044,1438,1046,1414,1338,59199 13.750000,13.562500,13.187500,12.875000,12.750000,12.062500,12.000000,11.937500 +95 0 11071 9.286315 1044,1046,1034,109431,1317,2251,1438,1338 13.125000,12.625000,12.562500,12.312500,12.187500,12.062500,11.937500,11.875000 +96 0 1010 6.097094 1044,1034,109431,1046,59199,1438,2251,1317 13.250000,13.062500,12.937500,12.812500,12.500000,12.375000,12.312500,12.000000 +97 0 6359 9.351042 1256,59199,11,1032,1438,1050,2251,1603 13.375000,12.750000,12.625000,12.312500,12.250000,11.812500,11.812500,11.750000 +98 0 1319 4.573673 1034,1438,1256,11,1044,59199,1032,1338 12.937500,12.625000,12.562500,12.375000,12.187500,12.187500,12.125000,12.125000 +99 0 1083 10.024452 1256,59199,1034,1438,2251,1032,1317,1603 12.375000,12.125000,11.625000,11.625000,11.562500,11.500000,11.250000,11.250000 +100 0 1046 3.839457 1034,1438,59199,1044,109431,2251,1046,1319 13.062500,12.500000,12.500000,12.437500,12.312500,12.000000,11.937500,11.812500 +101 0 10007 9.815475 1256,59199,1032,1034,1438,2251,1603,1319 12.562500,12.250000,11.812500,11.750000,11.750000,11.687500,11.562500,11.375000 +102 0 1532 12.341724 109431,1034,1046,1044,11,59199,1032,1317 13.875000,13.312500,13.187500,12.937500,12.750000,12.562500,12.375000,12.375000 +103 0 1044 2.644709 1044,109431,1046,59199,1034,1048,1041,1338 12.875000,12.625000,12.562500,12.187500,11.937500,11.937500,11.875000,11.687500 +104 0 1464 6.758195 1256,59199,1032,1034,1438,1603,2251,1364 12.062500,11.812500,11.437500,11.375000,11.375000,11.312500,11.312500,11.125000 +105 0 1046 3.339869 109431,1034,1044,1046,59199,1338,1317,1319 13.000000,12.750000,12.625000,12.500000,12.312500,12.187500,12.062500,11.937500 +106 0 60018 16.272507 1256,59199,1032,1034,1438,2251,1603,1319 12.562500,12.250000,11.812500,11.750000,11.750000,11.687500,11.562500,11.375000 +107 0 2140 10.455551 1034,1044,109431,1046,59199,1338,1050,1319 13.562500,13.062500,13.062500,12.937500,12.562500,11.875000,11.812500,11.812500 +108 0 1550 7.886243 1044,109431,1046,1034,59199,1338,1438,1041 13.312500,13.000000,12.937500,12.750000,12.375000,12.000000,11.937500,11.750000 +109 0 12336 13.765064 109431,1044,1034,1046,59199,1338,1319,1317 12.812500,12.562500,12.500000,12.250000,12.187500,11.937500,11.750000,11.687500 +110 0 1277 11.276552 109431,1044,1046,1034,59199,1438,1319,1338 13.875000,13.625000,13.250000,12.812500,12.812500,12.375000,12.250000,12.187500 +111 0 1044 3.201593 1034,1044,109431,1046,1438,1338,1278,59199 12.812500,12.437500,12.375000,12.250000,11.875000,11.625000,11.562500,11.500000 +112 0 1434 7.851945 1256,59199,1032,1034,1438,1603,2251,1364 12.062500,11.812500,11.437500,11.375000,11.375000,11.312500,11.312500,11.125000 +113 0 1046 3.785347 1034,109431,59199,1044,1438,1046,1338,1319 12.625000,12.375000,12.312500,12.250000,12.187500,11.875000,11.812500,11.687500 +114 0 4493 6.409225 1256,59199,1032,1034,1438,2251,1603,1319 12.562500,12.250000,11.812500,11.750000,11.750000,11.687500,11.562500,11.375000 +115 0 5757 15.742081 59199,1034,1317,1256,1032,1338,1438,2251 12.562500,12.437500,12.375000,12.250000,12.187500,12.187500,12.187500,12.125000 +116 0 2251 4.486694 1034,1044,1046,109431,1338,59199,1261,1438 13.312500,13.187500,12.937500,12.937500,12.187500,12.125000,11.937500,11.937500 +117 0 1429 4.304319 1256,1364,59199,1603,1032,2251,2129,1429 12.687500,12.250000,12.250000,12.187500,12.000000,12.000000,11.937500,11.875000 +118 0 28218 15.138948 1256,2251,59199,1034,1032,1317,1438,1319 12.250000,12.000000,12.000000,11.875000,11.625000,11.625000,11.562500,11.500000 +119 0 1335 8.701824 109431,1034,1044,1046,1438,59199,1041,1319 14.187500,14.062500,13.812500,13.312500,13.312500,13.187500,12.562500,12.500000 +120 0 102817 17.271538 1034,1044,109431,1317,59199,1046,1338,1438 12.437500,12.312500,12.187500,12.125000,12.062500,11.937500,11.812500,11.812500 +121 0 13064 12.671898 1034,1044,109431,1438,1046,59199,1319,2251 13.625000,13.625000,13.375000,13.062500,12.687500,12.500000,12.125000,12.125000 +122 0 10773 17.019188 109431,1046,1034,1044,59199,1317,1041,1338 13.937500,13.000000,12.875000,12.750000,12.375000,12.125000,12.062500,11.937500 +123 0 52928 13.578262 1044,1046,1034,109431,1438,1338,1307,1041 12.937500,12.562500,12.375000,12.312500,11.875000,11.812500,11.750000,11.500000 +124 0 1897 6.055256 1044,1046,1034,2251,1338,1438,1317,1319 13.687500,13.250000,12.875000,12.875000,12.625000,12.437500,12.375000,12.250000 +125 0 1032 3.921146 1034,59199,1438,2251,1032,1319,1317,1256 12.687500,12.375000,11.937500,11.875000,11.687500,11.625000,11.500000,11.375000 +126 0 1050 3.884239 1256,59199,1032,1603,2251,1050,11,1319 13.500000,12.687500,12.437500,12.250000,12.250000,12.125000,12.000000,11.937500 +127 0 1048 4.896251 1256,1034,59199,1438,1044,11,1338,109431 12.312500,12.250000,12.250000,12.187500,11.812500,11.750000,11.750000,11.750000 diff --git a/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1ctx512.tsv b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1ctx512.tsv new file mode 100644 index 000000000..9a004419d --- /dev/null +++ b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_default_w1ctx512.tsv @@ -0,0 +1,140 @@ +[NemotronH] Loading config... +[NemotronH] Config loaded: 52 layers (23 mamba, 6 attention, 0 mlp, 23 moe) +[NemotronH] Loading weights from safetensors... +[NemotronH] Building model... +[NemotronH] Model loaded successfully +# model models/mlx/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit +# device GPU +# corpus tests/fixtures/wikitext2_excerpt.txt tokens 301749 +# chunks 128 chunk_tokens 1 topk 8 prefill 512 +# columns chunk pos target nll top_ids top_logits +# start_token 512 +# contiguous false +0 0 1048 0.040632 1048,24926,1052,1054,1805,1462,1053,1056 24.500000,20.000000,19.500000,19.000000,19.000000,18.375000,18.000000,17.750000 +1 0 1048 0.289999 1048,1049,1050,1462,1411,2545,40438,2126 25.375000,24.250000,20.500000,18.125000,17.750000,17.625000,17.125000,16.750000 +2 0 1054 1.194082 1055,1054,1056,1053,1057,1052,1051,1050 24.125000,24.000000,23.375000,22.750000,22.500000,20.875000,20.500000,20.375000 +3 0 15011 0.522497 15011,19895,41500,22876,14696,8281,17058,5395 20.125000,18.750000,17.750000,17.625000,17.000000,16.375000,16.375000,15.875000 +4 0 1307 0.006359 1307,2126,1319,1408,29818,1429,1576,28140 26.000000,19.250000,18.750000,18.625000,18.375000,18.000000,18.000000,17.875000 +5 0 1278 0.334827 1278,1531,51288,19895,109492,22876,2813,1429 20.000000,17.000000,16.500000,16.000000,16.000000,15.937500,15.687500,15.625000 +6 0 14696 0.269133 14696,19895,22876,7520,8281,81873,41500,9371 20.875000,18.625000,18.250000,17.250000,16.875000,16.250000,16.125000,15.687500 +7 0 5395 0.121286 5395,25865,18823,41500,2629,3467,18944,14593 22.750000,19.875000,19.250000,18.125000,17.875000,17.750000,17.000000,16.750000 +8 0 3897 6.093834 51288,1531,2813,43193,15087,109492,1488,1373 17.375000,16.750000,16.625000,16.125000,15.687500,15.562500,15.375000,15.000000 +9 0 43193 4.231699 2813,1531,51288,3140,1799,19148,43193,109492 15.812500,15.375000,14.875000,14.500000,14.437500,13.875000,13.812500,13.812500 +10 0 1839 0.000210 1839,1606,1270,1428,2398,6094,1479,69555 27.750000,18.750000,16.500000,16.125000,15.937500,15.937500,15.687500,15.687500 +11 0 3897 1.062742 2247,3897,1435,18988,1294,12440,4295,11159 22.250000,22.250000,20.625000,20.250000,19.875000,19.625000,19.500000,19.375000 +12 0 7125 9.866854 12440,1435,1294,22594,4418,28140,2478,1321 23.000000,21.250000,20.375000,19.625000,19.125000,19.125000,19.000000,18.625000 +13 0 1536 0.033222 1536,1294,29598,1278,2015,1454,1261,10912 25.875000,21.375000,20.625000,20.125000,20.125000,19.875000,19.125000,18.750000 +14 0 1261 0.420521 1261,1420,3866,1278,24836,1794,15720,12440 21.875000,19.875000,19.250000,19.125000,19.125000,18.875000,18.250000,17.500000 +15 0 4934 1.405819 4934,24836,17317,47715,96432,1032,68393,3709 19.375000,19.250000,18.375000,18.125000,18.000000,17.500000,17.500000,17.375000 +16 0 1294 0.269489 1294,1435,1408,22594,12440,1307,3897,15019 24.875000,23.375000,21.250000,21.000000,20.750000,19.625000,19.000000,18.625000 +17 0 1278 0.084139 1278,1032,1261,1531,22876,1420,19895,81873 21.250000,17.000000,16.750000,16.375000,16.375000,16.000000,15.250000,14.875000 +18 0 1032 1.851337 3354,1032,4076,14696,19895,4958,22876,37848 19.625000,19.000000,18.750000,18.500000,18.500000,17.250000,17.125000,16.500000 +19 0 1050 0.003174 1050,1049,1055,1010,1051,1053,1052,1054 27.625000,20.750000,19.875000,19.500000,19.375000,19.250000,19.125000,19.000000 +20 0 1048 0.000631 1048,1049,1051,1298,1050,1052,1053,1055 26.250000,17.250000,16.375000,16.250000,16.000000,16.000000,16.000000,15.937500 +21 0 1048 0.021187 1048,1049,1050,1411,1055,1462,1051,2247 27.500000,23.625000,19.250000,19.125000,16.750000,16.500000,15.750000,15.125000 +22 0 1055 0.310822 1055,1056,1054,1057,1053,1052,1051,1050 25.750000,24.125000,23.625000,22.500000,20.875000,18.625000,18.000000,17.375000 +23 0 37848 2.827694 4076,14696,19895,3354,37848,22876,7326,15011 19.125000,18.500000,17.875000,17.750000,17.375000,16.625000,16.250000,16.250000 +24 0 5466 0.061932 5466,53458,3354,63239,6267,25349,2629,13079 25.000000,21.625000,21.000000,18.750000,18.250000,18.250000,17.750000,17.625000 +25 0 1307 0.056883 1307,3897,1531,1576,1294,2247,29818,1429 23.125000,18.875000,17.750000,17.625000,16.625000,16.500000,16.375000,16.250000 +26 0 3075 8.187262 1278,1531,1429,22511,15723,1349,6959,19420 19.000000,17.500000,15.250000,15.250000,15.187500,14.875000,14.437500,14.187500 +27 0 1317 2.699766 1278,1362,1531,1317,3870,9595,2247,1576 18.000000,17.000000,16.625000,16.000000,14.687500,14.500000,14.312500,13.687500 +28 0 18275 1.607203 40659,18275,3094,64561,93616,19008,13516,107178 18.000000,17.375000,16.875000,16.125000,15.625000,15.500000,15.000000,14.812500 +29 0 1415 0.106592 1415,3108,1419,1938,39664,86674,2392,3741 20.625000,17.875000,17.000000,15.750000,15.187500,14.687500,14.500000,12.812500 +30 0 15179 16.170532 1294,1278,1454,2247,1261,1513,1321,1737 21.750000,19.375000,19.250000,19.000000,18.875000,18.375000,17.875000,17.500000 +31 0 1536 0.011869 1536,2790,8360,1321,2126,1513,6421,5785 25.750000,20.000000,19.500000,19.125000,18.750000,18.500000,18.500000,17.500000 +32 0 8233 4.641790 15723,26004,26009,16566,6031,3986,16136,7456 16.750000,16.625000,16.000000,15.500000,15.437500,15.250000,15.250000,15.187500 +33 0 1386 0.000410 1386,3390,16926,97044,33151,2755,8101,1704 23.500000,13.875000,13.312500,13.312500,13.250000,13.125000,13.000000,12.500000 +34 0 1434 0.232655 1434,10307,1757,1535,1550,7456,14857,16006 18.250000,15.312500,15.062500,14.375000,14.000000,13.750000,13.312500,13.312500 +35 0 1479 1.508276 10917,1479,1539,25763,2765,4940,15392,47111 15.312500,14.687500,14.187500,12.000000,11.937500,11.187500,11.062500,10.875000 +36 0 1965 0.000164 1965,1107,2247,1121,1985,1046,12623,1837 24.875000,15.375000,13.750000,13.375000,13.062500,12.187500,12.187500,12.000000 +37 0 2247 0.446108 2247,1513,1321,1294,3897,18988,1454,11159 24.500000,23.000000,22.500000,21.500000,21.500000,20.750000,20.500000,20.000000 +38 0 3075 7.058146 2182,1032,1656,1531,2409,8052,4973,3264 20.250000,19.500000,19.375000,18.750000,18.125000,17.625000,17.250000,17.125000 +39 0 1317 0.043365 1317,3870,1278,9327,6793,6421,4058,2247 22.375000,18.750000,16.125000,15.437500,15.375000,15.250000,15.125000,15.000000 +40 0 18275 0.013184 18275,1359,121786,56669,83955,6421,44477,43060 25.750000,20.000000,19.875000,18.625000,18.500000,18.375000,17.875000,17.250000 +41 0 1415 0.000036 1415,1314,1672,1778,2247,1419,3108,32571 27.000000,15.562500,14.500000,13.937500,13.875000,13.500000,13.437500,13.375000 +42 0 1486 0.685030 1486,1395,56671,20209,3897,18166,5537,57062 21.375000,20.500000,18.250000,18.125000,18.000000,18.000000,17.500000,17.375000 +43 0 5537 0.335567 5537,8360,2158,1261,1420,4057,2095,8460 22.125000,20.375000,18.875000,18.375000,18.125000,17.750000,17.625000,17.625000 +44 0 1513 0.178732 1513,1294,1435,22594,1394,1317,1536,1562 25.250000,23.000000,21.500000,20.750000,20.500000,20.000000,19.750000,19.750000 +45 0 27640 9.916958 1278,2960,22410,1531,1335,1261,34614,7893 23.250000,19.000000,18.500000,17.625000,17.500000,17.000000,16.000000,15.750000 +46 0 22410 0.008699 22410,10786,27677,2247,19404,1109,35229,28869 23.250000,17.125000,17.000000,15.875000,15.750000,15.562500,15.500000,15.375000 +47 0 1294 1.285908 1321,1294,2247,3897,1454,2478,4295,1435 22.625000,22.500000,22.500000,20.375000,19.875000,19.125000,19.125000,18.875000 +48 0 1278 10.376852 7893,1032,36541,11250,75728,1278,2247,5663 25.625000,18.875000,16.625000,15.687500,15.500000,15.250000,15.062500,14.937500 +49 0 7893 2.150557 5046,7893,3354,1032,14719,2158,6509,10485 18.000000,16.625000,16.375000,15.750000,14.875000,14.500000,14.437500,14.375000 +50 0 81947 0.212025 81947,13375,2247,70574,3714,37848,5046,4457 20.500000,17.750000,17.625000,17.250000,16.250000,16.000000,15.250000,15.000000 +51 0 1307 0.004097 1307,1281,2247,9214,6421,18988,11159,35549 26.375000,19.750000,19.125000,19.125000,18.000000,17.750000,17.375000,16.500000 +52 0 38199 3.253658 108492,53875,38199,3316,16210,7893,45261,9581 21.500000,21.125000,18.875000,18.250000,17.500000,16.750000,16.750000,16.375000 +53 0 1428 0.000006 1428,1920,52636,8315,49887,1966,2247,58401 30.750000,18.125000,17.125000,16.125000,15.312500,14.625000,14.625000,14.437500 +54 0 52829 0.000191 52829,23825,10786,6989,1425,11564,1109,2710 25.750000,15.937500,15.187500,14.750000,14.375000,14.000000,13.625000,13.500000 +55 0 1321 0.023052 1321,2247,1748,18988,11159,4295,3897,6421 25.625000,21.375000,19.625000,19.500000,18.625000,18.125000,18.000000,17.375000 +56 0 118799 0.008919 118799,75728,38199,10330,44841,21374,39185,1486 22.625000,16.625000,16.500000,15.375000,14.812500,14.687500,14.437500,14.312500 +57 0 2247 0.312902 2247,18988,1321,11159,1294,3897,1454,4295 23.750000,21.625000,21.500000,20.375000,20.125000,20.000000,19.250000,19.000000 +58 0 89656 4.675072 1032,1656,2182,1531,8052,2409,2157,6304 20.125000,19.750000,19.625000,18.500000,17.375000,17.250000,17.000000,17.000000 +59 0 1397 0.001940 1397,3701,5032,1400,1868,1781,21985,2683 23.375000,15.437500,15.375000,15.312500,14.250000,13.562500,13.500000,13.375000 +60 0 56671 3.322027 8308,10354,1880,2095,1576,1486,56671,5537 20.125000,18.750000,18.500000,18.500000,18.375000,18.000000,17.875000,17.750000 +61 0 1294 2.070908 22594,1294,1435,1454,15019,1408,2790,1513 25.375000,23.500000,22.250000,20.125000,19.750000,19.000000,18.875000,18.500000 +62 0 2295 5.014128 1278,1261,1420,1531,1032,2295,1593,22876 20.250000,18.250000,16.375000,15.937500,15.625000,15.500000,15.437500,15.250000 +63 0 13068 4.065094 23981,4958,11733,1032,14696,4514,13068,22876 20.125000,18.000000,17.375000,17.250000,16.875000,16.500000,16.500000,16.125000 +64 0 1294 1.030293 1294,7358,15179,3897,8360,1455,8077,1307 20.625000,19.750000,19.375000,19.000000,18.500000,18.250000,18.250000,17.500000 +65 0 1032 0.023184 1032,1278,3897,5124,7276,6032,1766,5466 25.500000,21.000000,18.625000,18.500000,18.250000,18.125000,18.000000,17.750000 +66 0 1050 0.000177 1050,1049,1048,1052,1051,1053,1032,1056 26.375000,16.375000,15.625000,15.625000,15.062500,14.812500,14.687500,14.625000 +67 0 1048 0.000116 1048,1032,1050,1049,3897,1051,1052,1055 26.500000,16.125000,15.812500,14.937500,14.187500,13.937500,13.687500,13.500000 +68 0 1048 0.113373 1048,1049,1050,1737,2247,1056,1032,1051 27.000000,24.875000,19.500000,14.812500,14.750000,14.437500,14.250000,14.062500 +69 0 1056 0.850134 1056,1055,1057,1054,1053,1052,1051,1050 23.500000,23.250000,22.500000,21.625000,20.000000,18.125000,17.250000,16.875000 +70 0 3897 1.260323 2247,3897,1737,1775,4295,1319,2251,15179 21.875000,21.625000,20.125000,20.125000,20.125000,19.125000,18.750000,18.750000 +71 0 10298 8.887136 15179,2960,1278,28140,1531,3907,1799,1925 15.875000,15.625000,15.562500,15.500000,15.312500,15.312500,14.875000,14.500000 +72 0 7421 1.366634 1307,7421,5661,2247,1032,15247,53728,32022 18.500000,17.875000,17.000000,16.375000,15.250000,15.187500,14.875000,14.687500 +73 0 9225 1.860413 1576,9225,92791,1321,7114,1681,1766,1032 16.875000,16.750000,16.375000,16.250000,15.250000,14.937500,14.875000,14.625000 +74 0 67465 0.000130 67465,14397,6575,1575,2097,2249,3862,1616 33.000000,23.875000,21.625000,21.250000,18.125000,16.125000,15.875000,15.375000 +75 0 1536 6.198950 1321,15179,22594,3897,1319,8360,47715,2247 23.750000,21.500000,21.125000,20.750000,20.625000,20.375000,19.125000,18.875000 +76 0 130537 4.969790 41520,16566,3986,7456,6031,12245,15723,36356 17.125000,16.875000,16.500000,16.500000,16.375000,16.375000,16.375000,16.125000 +77 0 6993 11.104012 6031,7456,3986,15723,16566,1349,8720,8343 17.250000,16.875000,16.750000,16.250000,16.125000,15.937500,15.875000,15.812500 +78 0 2301 6.395235 2163,10307,15450,13586,5303,2301,39316,2247 19.250000,17.000000,15.062500,13.625000,13.500000,13.000000,12.375000,12.187500 +79 0 125371 1.945326 125371,11461,48227,7976,1398,73607,1760,51039 13.500000,13.375000,12.812500,12.500000,11.625000,11.375000,11.125000,11.125000 +80 0 3897 1.769909 1321,4728,3897,75094,1305,4295,3614,2247 20.000000,19.875000,19.125000,16.750000,16.500000,15.875000,15.500000,15.437500 +81 0 1321 0.451180 1321,22594,1799,7125,47715,1294,1261,12440 21.750000,19.750000,19.500000,19.375000,19.000000,17.875000,17.500000,17.500000 +82 0 6797 7.257715 1531,1278,1294,2095,1656,29251,1349,1434 17.375000,15.750000,15.000000,14.562500,14.312500,14.062500,14.000000,14.000000 +83 0 3892 1.292631 3892,1576,5303,14258,60878,20040,2405,6174 17.000000,16.500000,15.875000,15.062500,14.812500,14.500000,14.437500,14.375000 +84 0 1390 7.045465 2247,1536,3897,2662,2314,10689,1319,4295 15.812500,14.375000,13.937500,13.750000,13.375000,12.750000,12.375000,12.312500 +85 0 8919 0.000059 8919,61245,1906,6465,1518,3257,1783,12200 26.250000,15.062500,14.312500,14.250000,13.875000,13.562500,13.375000,13.187500 +86 0 15179 0.966709 15179,2247,1536,47715,3897,1319,8360,4295 21.250000,21.125000,19.375000,19.125000,18.875000,18.625000,18.625000,18.250000 +87 0 1536 0.004276 1536,1321,3897,22594,1294,2790,2126,2247 25.625000,18.500000,17.625000,17.500000,17.375000,17.375000,17.125000,17.125000 +88 0 8595 1.577899 8595,26305,43252,36356,7114,40006,33407,9878 13.750000,12.687500,12.187500,11.625000,11.375000,11.375000,11.187500,11.062500 +89 0 1405 0.079136 1405,1269,5625,2302,6291,11254,19019,3230 20.875000,17.500000,17.375000,15.125000,14.625000,14.500000,14.250000,13.687500 +90 0 88542 0.953245 88542,28958,10307,2247,26515,66533,84292,78734 15.750000,14.562500,14.500000,13.187500,13.187500,12.750000,12.750000,12.562500 +91 0 2247 0.112402 2247,18988,3897,4295,11159,1319,1321,1294 25.125000,22.375000,20.875000,20.875000,20.875000,18.875000,18.875000,18.625000 +92 0 1656 2.934446 2182,1032,12953,6797,1531,1656,10298,6304 20.875000,20.125000,19.875000,19.125000,19.000000,19.000000,18.750000,17.875000 +93 0 4454 6.092888 1032,1278,2960,6346,6866,4454,6633,5821 24.750000,21.625000,19.750000,19.500000,19.125000,18.750000,18.750000,18.625000 +94 0 1032 0.073021 1032,1307,3897,1321,1462,1278,1576,6633 23.875000,20.500000,19.500000,19.250000,17.375000,17.000000,16.875000,16.875000 +95 0 1050 0.001629 1050,1048,1049,1051,1056,1053,1054,1052 26.875000,19.750000,19.000000,17.750000,17.375000,17.250000,16.875000,16.750000 +96 0 1048 0.000851 1048,1055,1032,1050,1051,1053,1056,1057 26.125000,17.125000,16.875000,16.750000,16.625000,16.625000,16.250000,16.250000 +97 0 1048 0.583193 1048,1049,1050,1057,3897,1051,2247,1056 24.875000,24.625000,20.500000,15.687500,14.500000,14.125000,13.750000,13.562500 +98 0 1056 0.862974 1057,1056,1055,1054,1053,1052,1050,1051 24.500000,24.250000,21.250000,20.250000,19.750000,18.375000,17.500000,17.375000 +99 0 3897 0.640310 3897,1514,89656,1278,1794,43378,1321,2165 20.000000,19.250000,18.500000,17.125000,15.750000,14.937500,14.625000,14.500000 +100 0 89656 1.034754 89656,1514,1278,1531,10298,1494,1794,1261 17.250000,17.125000,15.750000,15.062500,14.187500,13.812500,13.312500,13.062500 +101 0 1397 0.007292 1397,1400,3701,5032,21985,2306,31814,1781 20.500000,13.875000,13.812500,13.812500,12.750000,12.375000,12.000000,11.125000 +102 0 3214 4.779621 8308,10354,5537,1486,1880,37101,52158,56671 19.125000,18.625000,18.000000,17.875000,17.750000,17.625000,17.625000,17.375000 +103 0 1261 0.683270 1261,1794,1420,24836,96432,1278,3866,29857 23.625000,23.250000,22.250000,18.500000,18.250000,18.125000,17.500000,17.500000 +104 0 24836 0.717749 24836,96432,4958,13426,6061,1850,8602,4076 20.500000,20.125000,18.125000,17.375000,17.000000,16.500000,16.000000,15.687500 +105 0 13988 0.021297 13988,47715,96432,14696,8281,2126,9655,29857 23.125000,17.375000,17.000000,16.875000,16.375000,16.125000,16.000000,15.750000 +106 0 1408 0.490801 1408,1294,1435,22594,12440,1454,1513,1321 25.125000,24.375000,22.750000,21.875000,20.500000,19.125000,19.000000,18.750000 +107 0 1261 4.933411 1278,22876,1531,19895,14279,1420,1261,14696 20.125000,17.875000,17.625000,17.625000,16.125000,15.875000,15.562500,15.437500 +108 0 2295 6.826847 19895,6061,22876,78172,14696,1875,12321,41500 16.500000,16.375000,16.375000,15.875000,15.625000,15.500000,15.500000,15.375000 +109 0 2126 7.081820 24926,1462,1805,109431,1045,50824,15011,2126 23.375000,20.000000,19.625000,18.625000,17.875000,17.750000,16.500000,16.375000 +110 0 1045 0.049957 1045,1462,93023,1032,1123,18242,1064,1925 21.375000,17.750000,15.312500,14.937500,14.875000,14.312500,13.812500,13.812500 +111 0 1064 0.003304 1064,93023,115587,73856,59962,2126,3332,1050 26.625000,20.375000,18.500000,18.125000,17.625000,17.500000,17.000000,16.375000 +112 0 1805 0.091255 1805,12271,15011,1635,31194,14400,2786,106817 23.500000,20.250000,20.125000,18.250000,18.250000,16.625000,16.500000,16.500000 +113 0 15011 0.116035 15011,41500,8303,6061,19895,5395,13268,2247 20.875000,16.500000,16.375000,16.125000,16.125000,15.625000,15.562500,15.187500 +114 0 19976 13.768942 1307,1408,29818,4418,1294,28140,1319,2126 25.750000,20.500000,20.500000,19.000000,18.875000,18.750000,18.625000,18.375000 +115 0 1307 0.366890 1307,1408,1294,29818,1435,28140,1394,3897 22.750000,21.625000,19.625000,18.375000,18.250000,18.250000,17.250000,16.750000 +116 0 1278 0.094358 1278,1531,3797,15087,22876,43193,2813,81873 20.875000,16.250000,15.500000,15.500000,15.437500,15.437500,15.312500,15.250000 +117 0 14696 0.520361 14696,19895,22876,7520,9371,5150,81873,2730 20.000000,18.250000,18.000000,17.750000,17.250000,15.875000,15.687500,15.625000 +118 0 5395 0.113461 5395,25865,18823,2629,14593,50584,3467,119435 22.500000,19.625000,19.000000,17.875000,17.500000,16.500000,16.125000,15.875000 +119 0 1488 3.882605 3897,15087,1531,2813,51288,43193,1488,45776 17.750000,15.875000,15.437500,15.375000,15.312500,15.187500,14.812500,14.375000 +120 0 6812 0.353123 6812,89220,24591,6970,2939,5635,10178,25632 19.375000,18.375000,14.687500,14.187500,13.562500,13.500000,13.312500,13.250000 +121 0 1278 0.028637 1278,1531,8250,29251,3897,2247,2015,1032 21.250000,16.875000,15.562500,15.437500,14.312500,14.125000,13.937500,13.625000 +122 0 29251 0.030895 29251,19213,1407,8776,39447,49944,3897,49377 21.875000,15.875000,15.812500,15.562500,15.375000,15.375000,15.187500,15.125000 +123 0 3897 3.426485 2247,18988,11159,1435,3897,22594,29818,1454 22.625000,20.750000,19.750000,19.625000,19.625000,19.500000,19.250000,18.875000 +124 0 7125 4.104105 12440,1435,28140,1294,29818,7125,22594,8308 20.750000,17.625000,17.625000,17.500000,17.500000,17.000000,17.000000,16.875000 +125 0 1536 0.069681 1536,1294,29598,2015,4025,1261,1278,17558 25.625000,22.250000,21.000000,20.375000,19.875000,19.750000,19.500000,19.250000 +126 0 1420 2.163845 1261,1420,3866,24836,1278,1794,15720,29857 22.250000,20.500000,19.625000,19.500000,19.125000,19.125000,18.250000,17.375000 +127 0 13988 0.432508 13988,15011,1457,7481,47353,1408,87170,107254 21.250000,19.875000,19.250000,16.875000,16.625000,16.000000,16.000000,16.000000 diff --git a/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_ssmkernel0_w1ctx512.tsv b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_ssmkernel0_w1ctx512.tsv new file mode 100644 index 000000000..336d8c2c4 --- /dev/null +++ b/benchmarks/logit_traces/rocm_gfx1151_96cbce84/rocm_gfx1151_96cbce84_nemotron-3-nano-30b-a3b_ssmkernel0_w1ctx512.tsv @@ -0,0 +1,140 @@ +[NemotronH] Loading config... +[NemotronH] Config loaded: 52 layers (23 mamba, 6 attention, 0 mlp, 23 moe) +[NemotronH] Loading weights from safetensors... +[NemotronH] Building model... +[NemotronH] Model loaded successfully +# model models/mlx/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit +# device GPU +# corpus tests/fixtures/wikitext2_excerpt.txt tokens 301749 +# chunks 128 chunk_tokens 1 topk 8 prefill 512 +# columns chunk pos target nll top_ids top_logits +# start_token 512 +# contiguous false +0 0 1048 0.042322 1048,24926,1052,1054,1805,1462,1053,1056 24.500000,20.000000,19.500000,19.125000,19.125000,18.500000,18.000000,17.750000 +1 0 1048 0.289989 1048,1049,1050,1462,1411,2545,40438,91559 25.375000,24.250000,20.500000,18.125000,17.625000,17.500000,17.250000,16.875000 +2 0 1054 1.248182 1055,1054,1056,1053,1057,1052,1051,1050 24.125000,23.875000,23.250000,22.625000,22.500000,20.750000,20.375000,20.250000 +3 0 15011 0.560434 15011,19895,22876,41500,14696,8281,17058,5395 20.125000,18.875000,17.750000,17.750000,17.125000,16.500000,16.375000,16.000000 +4 0 1307 0.006403 1307,2126,1319,1408,29818,1429,4418,28140 26.000000,19.250000,18.750000,18.750000,18.375000,18.000000,17.875000,17.875000 +5 0 1278 0.346380 1278,1531,51288,19895,109492,22876,2813,1429 20.000000,17.125000,16.500000,16.125000,16.000000,15.937500,15.750000,15.687500 +6 0 14696 0.267420 14696,19895,22876,7520,8281,81873,41500,9371 20.875000,18.625000,18.250000,17.250000,16.750000,16.250000,16.125000,15.687500 +7 0 5395 0.123075 5395,25865,18823,41500,2629,3467,18944,14593 22.625000,19.750000,19.125000,18.125000,17.750000,17.625000,16.875000,16.625000 +8 0 3897 6.093842 51288,1531,2813,43193,15087,109492,1488,1373 17.375000,16.750000,16.625000,16.125000,15.687500,15.562500,15.312500,15.062500 +9 0 43193 4.173510 2813,1531,51288,1799,3140,19148,43193,109492 15.812500,15.375000,14.875000,14.500000,14.500000,13.875000,13.875000,13.812500 +10 0 1839 0.000216 1839,1606,1270,1428,6094,69555,2398,82089 27.750000,18.875000,16.375000,15.875000,15.875000,15.750000,15.687500,15.437500 +11 0 3897 1.045639 2247,3897,1435,18988,1294,12440,4295,11159 22.375000,22.375000,20.750000,20.375000,19.875000,19.750000,19.500000,19.375000 +12 0 7125 9.907101 12440,1435,1294,22594,4418,28140,2478,1321 23.125000,21.375000,20.375000,19.500000,19.250000,19.125000,19.000000,18.625000 +13 0 1536 0.032520 1536,1294,29598,1278,2015,1454,1261,10912 25.875000,21.375000,20.500000,20.125000,20.125000,19.875000,19.125000,18.750000 +14 0 1261 0.401934 1261,1420,1278,3866,24836,1794,15720,12440 22.000000,20.000000,19.250000,19.250000,19.125000,19.000000,18.250000,17.500000 +15 0 4934 1.395081 4934,24836,17317,47715,96432,68393,1032,3709 19.375000,19.250000,18.375000,18.125000,18.000000,17.500000,17.375000,17.375000 +16 0 1294 0.297014 1294,1435,1408,22594,12440,1307,3897,15019 24.750000,23.375000,21.125000,21.125000,20.625000,19.500000,18.875000,18.625000 +17 0 1278 0.080421 1278,1032,1261,1531,22876,1420,19895,81873 21.250000,16.875000,16.750000,16.375000,16.250000,16.000000,15.187500,14.812500 +18 0 1032 1.831333 3354,1032,4076,14696,19895,4958,22876,37848 19.625000,19.000000,18.625000,18.500000,18.500000,17.250000,17.000000,16.500000 +19 0 1050 0.003378 1050,1049,1055,1051,1010,1053,1052,1054 27.625000,20.750000,20.125000,19.500000,19.375000,19.375000,19.250000,19.000000 +20 0 1048 0.000601 1048,1049,1051,1298,1050,1052,1053,1055 26.375000,17.375000,16.500000,16.250000,16.125000,16.125000,16.125000,16.000000 +21 0 1048 0.018791 1048,1049,1050,1411,1055,1462,1051,2247 27.500000,23.500000,19.250000,19.125000,16.750000,16.375000,15.687500,15.062500 +22 0 1055 0.310820 1055,1056,1054,1057,1053,1052,1051,1050 25.875000,24.250000,23.750000,22.625000,21.000000,18.750000,18.125000,17.500000 +23 0 37848 2.809698 4076,14696,3354,19895,37848,22876,7326,15011 19.125000,18.500000,17.750000,17.750000,17.375000,16.625000,16.125000,16.125000 +24 0 5466 0.061943 5466,53458,3354,63239,6267,25349,2629,13079 25.000000,21.750000,20.750000,18.625000,18.250000,18.000000,17.625000,17.375000 +25 0 1307 0.057266 1307,3897,1531,1576,1294,2247,29818,4418 23.125000,18.875000,17.750000,17.625000,16.625000,16.500000,16.375000,16.250000 +26 0 3075 8.127525 1278,1531,15723,1429,22511,1349,6959,19420 19.000000,17.500000,15.250000,15.187500,15.187500,14.812500,14.500000,14.250000 +27 0 1317 2.613325 1278,1362,1531,1317,3870,9595,2247,1576 18.000000,17.125000,16.625000,16.125000,14.750000,14.562500,14.437500,13.687500 +28 0 18275 1.459572 40659,18275,3094,64561,93616,19008,13516,107178 17.875000,17.500000,16.875000,16.125000,15.562500,15.500000,15.062500,14.750000 +29 0 1415 0.128716 1415,3108,1419,1938,39664,86674,2392,3741 20.375000,17.875000,16.875000,15.500000,15.125000,14.750000,14.375000,12.750000 +30 0 15179 16.472836 1294,1278,1454,1261,2247,1513,1321,1420 21.875000,19.375000,19.250000,19.000000,18.750000,18.375000,17.875000,17.625000 +31 0 1536 0.011805 1536,2790,8360,1321,1513,2126,6421,5785 25.750000,20.000000,19.500000,19.125000,18.625000,18.625000,18.500000,17.500000 +32 0 8233 4.734879 15723,26004,26009,16566,6031,3986,7456,16136 16.750000,16.625000,15.937500,15.500000,15.375000,15.250000,15.187500,15.187500 +33 0 1386 0.000402 1386,3390,33151,16926,97044,2755,8101,1704 23.500000,13.875000,13.312500,13.250000,13.250000,13.125000,12.875000,12.562500 +34 0 1434 0.227648 1434,1757,10307,1535,1550,7456,13225,16006 18.250000,15.125000,15.125000,14.375000,14.125000,13.812500,13.187500,13.187500 +35 0 1479 1.293940 10917,1479,1539,2765,25763,47111,4940,15392 15.187500,14.875000,13.937500,12.187500,11.875000,11.312500,11.250000,11.187500 +36 0 1965 0.000193 1965,1107,1121,2247,1985,1046,12623,65526 24.625000,15.250000,13.562500,13.500000,13.000000,11.937500,11.937500,11.937500 +37 0 2247 0.450699 2247,1513,1321,3897,1294,18988,1454,11159 24.500000,23.000000,22.500000,21.625000,21.500000,20.750000,20.500000,20.000000 +38 0 3075 7.015385 2182,1032,1656,1531,2409,8052,4973,3264 20.125000,19.500000,19.375000,18.750000,18.125000,17.625000,17.250000,17.125000 +39 0 1317 0.047411 1317,3870,1278,6793,9327,6421,4058,2247 22.375000,18.875000,16.125000,15.437500,15.375000,15.250000,15.125000,15.062500 +40 0 18275 0.010918 18275,121786,1359,56669,83955,6421,44477,64093 26.000000,20.125000,19.875000,18.625000,18.625000,18.375000,17.875000,17.250000 +41 0 1415 0.000038 1415,1314,1672,2247,1778,1419,32571,3108 27.000000,15.625000,14.625000,14.000000,13.812500,13.750000,13.625000,13.437500 +42 0 1486 0.662741 1486,1395,56671,20209,3897,18166,5537,57062 21.375000,20.375000,18.250000,18.125000,18.000000,18.000000,17.500000,17.375000 +43 0 5537 0.344912 5537,8360,2158,1261,1420,2095,4057,8460 22.000000,20.250000,18.750000,18.250000,18.000000,17.625000,17.625000,17.500000 +44 0 1513 0.196310 1513,1294,1435,22594,1394,1317,1408,1536 25.125000,23.000000,21.375000,20.750000,20.500000,20.000000,19.875000,19.625000 +45 0 27640 9.916943 1278,2960,22410,1531,1335,1261,34614,7893 23.250000,19.000000,18.500000,17.625000,17.500000,17.000000,15.937500,15.750000 +46 0 22410 0.008026 22410,10786,27677,2247,19404,35229,28869,1109 23.375000,17.125000,17.125000,15.937500,15.750000,15.625000,15.437500,15.375000 +47 0 1294 1.286890 1321,1294,2247,3897,1454,2478,4295,1435 22.625000,22.500000,22.500000,20.375000,19.875000,19.125000,19.125000,18.875000 +48 0 1278 10.252029 7893,1032,36541,11250,75728,1278,2247,5663 25.625000,19.000000,16.625000,15.750000,15.562500,15.375000,15.062500,14.937500 +49 0 7893 2.252169 5046,7893,3354,1032,14719,2158,6509,10485 18.000000,16.500000,16.375000,15.750000,14.812500,14.500000,14.375000,14.375000 +50 0 81947 0.209034 81947,13375,2247,70574,3714,37848,5046,4457 20.500000,17.750000,17.500000,17.250000,16.375000,16.000000,15.250000,15.000000 +51 0 1307 0.004282 1307,1281,9214,2247,6421,18988,11159,35549 26.250000,19.625000,19.125000,19.000000,17.875000,17.750000,17.250000,16.500000 +52 0 38199 3.309525 108492,53875,38199,3316,16210,7893,45261,9581 21.375000,21.125000,18.750000,18.250000,17.500000,16.750000,16.750000,16.375000 +53 0 1428 0.000006 1428,1920,52636,8315,49887,1966,2247,95522 30.750000,18.125000,17.125000,15.937500,15.312500,14.562500,14.500000,14.312500 +54 0 52829 0.000187 52829,23825,10786,6989,1425,11564,1109,40158 25.750000,15.875000,15.250000,14.750000,14.375000,14.000000,13.500000,13.375000 +55 0 1321 0.023079 1321,2247,1748,18988,11159,4295,3897,6421 25.625000,21.375000,19.625000,19.500000,18.750000,18.125000,17.875000,17.250000 +56 0 118799 0.010880 118799,38199,75728,10330,44841,21374,39185,1486 22.500000,16.750000,16.625000,15.437500,14.812500,14.687500,14.437500,14.375000 +57 0 2247 0.312292 2247,18988,1321,11159,1294,3897,1454,4295 23.750000,21.625000,21.500000,20.375000,20.125000,20.000000,19.250000,19.000000 +58 0 89656 4.740805 1032,1656,2182,1531,8052,2409,2157,6304 20.000000,19.750000,19.625000,18.375000,17.375000,17.250000,17.000000,16.875000 +59 0 1397 0.001539 1397,5032,1400,3701,1868,2683,1781,21985 23.875000,15.750000,15.687500,15.625000,14.562500,14.062500,13.500000,13.250000 +60 0 56671 3.376667 8308,10354,1576,1880,2095,1486,56671,5537 20.125000,18.875000,18.625000,18.625000,18.500000,18.125000,17.875000,17.625000 +61 0 1294 2.181219 22594,1294,1435,1454,15019,1408,2790,1513 25.375000,23.375000,22.250000,20.125000,19.875000,19.000000,18.875000,18.500000 +62 0 2295 4.997498 1278,1261,1420,1531,1032,2295,1593,22876 20.250000,18.125000,16.250000,15.875000,15.687500,15.500000,15.375000,15.000000 +63 0 13068 4.062016 23981,4958,11733,1032,14696,4514,13068,22876 20.125000,18.000000,17.375000,17.250000,16.875000,16.500000,16.500000,16.125000 +64 0 1294 1.136456 1294,7358,15179,3897,8360,8077,1455,22594 20.625000,20.000000,19.500000,18.875000,18.875000,18.625000,18.125000,17.750000 +65 0 1032 0.020744 1032,1278,3897,5124,6032,7276,1766,7586 25.625000,21.000000,18.625000,18.500000,18.250000,18.125000,18.000000,17.750000 +66 0 1050 0.000162 1050,1049,1048,1052,1051,1053,1056,1057 26.750000,16.625000,15.937500,15.937500,15.312500,15.062500,15.062500,14.937500 +67 0 1048 0.000116 1048,1032,1050,1049,3897,1051,1052,1055 26.500000,16.125000,15.812500,14.875000,14.312500,13.937500,13.625000,13.437500 +68 0 1048 0.113438 1048,1049,1050,1737,2247,1056,1032,1051 27.000000,24.875000,19.625000,14.750000,14.750000,14.562500,14.187500,14.125000 +69 0 1056 0.823215 1056,1055,1057,1054,1053,1052,1051,1050 23.625000,23.375000,22.500000,21.625000,20.125000,18.125000,17.250000,17.000000 +70 0 3897 1.304632 2247,3897,1737,1775,4295,1319,2251,4254 22.125000,21.750000,20.250000,20.250000,20.250000,19.250000,19.000000,18.875000 +71 0 10298 8.887825 15179,2960,1278,1531,28140,3907,1799,1925 15.875000,15.625000,15.562500,15.375000,15.375000,15.312500,14.937500,14.562500 +72 0 7421 1.261198 1307,7421,5661,2247,1032,15247,53728,32022 18.500000,18.000000,16.875000,16.375000,15.312500,15.062500,14.875000,14.687500 +73 0 9225 2.051453 1576,9225,1321,92791,7114,1681,1766,1335 17.125000,16.625000,16.375000,16.375000,15.437500,15.062500,15.062500,14.687500 +74 0 67465 0.000153 67465,14397,1575,6575,2097,3862,2249,1616 33.250000,24.375000,21.375000,21.375000,18.375000,16.125000,16.000000,15.312500 +75 0 1536 6.314381 1321,15179,22594,1319,3897,8360,47715,2247 23.875000,21.625000,21.250000,20.750000,20.750000,20.375000,19.250000,19.000000 +76 0 130537 4.978004 41520,16566,3986,6031,7456,12245,15723,36356 17.125000,16.875000,16.500000,16.500000,16.500000,16.375000,16.375000,16.125000 +77 0 6993 11.047720 6031,7456,3986,15723,16566,1349,8343,8720 17.250000,16.875000,16.750000,16.250000,16.125000,16.000000,15.875000,15.812500 +78 0 2301 6.565687 2163,10307,15450,13586,5303,2301,2247,39316 19.375000,17.000000,15.000000,13.562500,13.437500,12.937500,12.312500,12.312500 +79 0 125371 1.893908 125371,11461,48227,7976,1398,73607,1760,51039 13.562500,13.375000,12.750000,12.562500,11.750000,11.375000,11.125000,11.125000 +80 0 3897 1.923199 1321,4728,3897,1305,75094,4295,3614,2247 19.875000,19.875000,18.875000,16.500000,16.500000,15.687500,15.500000,15.250000 +81 0 1321 0.473541 1321,22594,1799,7125,47715,1294,1261,12440 21.625000,19.750000,19.375000,19.375000,18.875000,17.875000,17.375000,17.375000 +82 0 6797 7.211567 1531,1278,1294,2095,1656,29251,1349,1434 17.375000,15.812500,14.937500,14.500000,14.312500,14.125000,14.000000,14.000000 +83 0 3892 1.324293 3892,1576,5303,14258,60878,20040,2405,6174 17.000000,16.625000,15.937500,14.937500,14.812500,14.500000,14.437500,14.437500 +84 0 1390 7.021889 2247,1536,3897,2662,2314,10689,1319,4295 16.000000,14.500000,14.187500,13.937500,13.625000,12.750000,12.625000,12.625000 +85 0 8919 0.000059 8919,61245,1518,1906,6465,3257,1783,12200 26.250000,15.062500,14.312500,14.000000,13.937500,13.625000,13.375000,13.250000 +86 0 15179 0.990368 15179,2247,1536,47715,3897,1319,8360,4295 21.250000,21.125000,19.500000,19.250000,18.875000,18.750000,18.750000,18.250000 +87 0 1536 0.004879 1536,1321,3897,2790,22594,1294,2126,2247 25.500000,18.500000,17.625000,17.500000,17.500000,17.250000,17.250000,17.125000 +88 0 8595 1.592922 8595,26305,43252,36356,40006,7114,9878,33407 13.750000,12.812500,12.000000,11.687500,11.500000,11.437500,11.125000,11.125000 +89 0 1405 0.070250 1405,1269,5625,2302,6291,11254,19019,3230 21.000000,17.500000,17.375000,15.125000,14.687500,14.500000,14.187500,13.812500 +90 0 88542 0.873457 88542,28958,10307,26515,2247,66533,84292,78734 15.875000,14.562500,14.500000,13.250000,13.062500,12.750000,12.625000,12.562500 +91 0 2247 0.113979 2247,18988,11159,3897,4295,1319,1294,1321 25.125000,22.375000,21.000000,20.875000,20.875000,18.875000,18.750000,18.750000 +92 0 1656 2.974361 2182,1032,12953,6797,1531,1656,10298,6304 20.875000,20.250000,20.000000,19.125000,19.000000,19.000000,18.750000,17.875000 +93 0 4454 6.092216 1032,1278,2960,6346,6866,4454,6633,5821 24.750000,21.625000,19.750000,19.375000,19.125000,18.750000,18.750000,18.625000 +94 0 1032 0.066786 1032,1307,3897,1321,1462,1278,6633,1576 24.000000,20.500000,19.500000,19.375000,17.625000,17.000000,17.000000,16.875000 +95 0 1050 0.001833 1050,1048,1049,1051,1056,1053,1054,1052 27.125000,20.125000,19.375000,18.125000,17.750000,17.625000,17.375000,17.125000 +96 0 1048 0.000856 1048,1055,1032,1050,1053,1051,1056,1057 26.125000,17.125000,17.000000,16.750000,16.625000,16.500000,16.250000,16.250000 +97 0 1048 0.530760 1048,1049,1050,1057,3897,1051,2247,1056 24.875000,24.500000,20.500000,15.625000,14.437500,14.187500,13.750000,13.562500 +98 0 1056 0.932623 1057,1056,1055,1054,1053,1052,1050,1051 24.625000,24.250000,21.250000,20.250000,19.750000,18.375000,17.500000,17.375000 +99 0 3897 0.699791 3897,1514,89656,1278,1794,43378,1321,2165 19.875000,19.250000,18.500000,17.125000,15.750000,14.875000,14.625000,14.375000 +100 0 89656 1.153080 1514,89656,1278,1531,10298,1494,1794,1261 17.250000,17.125000,15.687500,15.062500,14.125000,13.812500,13.250000,13.062500 +101 0 1397 0.007441 1397,1400,5032,3701,21985,2306,31814,1781 20.500000,13.875000,13.875000,13.812500,12.750000,12.375000,12.000000,11.125000 +102 0 3214 4.705956 8308,10354,5537,1486,1880,37101,52158,56671 19.000000,18.500000,18.000000,17.875000,17.750000,17.500000,17.500000,17.250000 +103 0 1261 0.625673 1261,1794,1420,24836,96432,1278,3866,29857 23.625000,23.125000,22.125000,18.500000,18.250000,18.125000,17.500000,17.500000 +104 0 24836 0.666208 24836,96432,4958,13426,6061,1850,8602,4076 20.500000,20.000000,18.125000,17.125000,16.875000,16.375000,15.937500,15.687500 +105 0 13988 0.022766 13988,47715,96432,14696,8281,2126,9655,29857 23.125000,17.500000,17.125000,17.000000,16.375000,16.250000,16.000000,15.812500 +106 0 1408 0.484875 1408,1294,1435,22594,12440,1513,1454,1321 25.125000,24.375000,22.625000,21.875000,20.625000,19.125000,19.000000,18.750000 +107 0 1261 4.873711 1278,22876,1531,19895,14279,1420,1261,14696 20.125000,17.875000,17.625000,17.625000,16.125000,15.937500,15.625000,15.437500 +108 0 2295 6.861080 19895,6061,22876,78172,14696,12321,1875,41291 16.500000,16.375000,16.375000,15.875000,15.687500,15.562500,15.500000,15.437500 +109 0 2126 6.956696 24926,1462,1805,109431,1045,50824,15011,2126 23.250000,19.750000,19.500000,18.750000,17.750000,17.625000,16.500000,16.375000 +110 0 1045 0.049372 1045,1462,93023,1032,1123,18242,1925,39450 21.375000,17.750000,15.125000,15.000000,14.812500,14.312500,13.812500,13.687500 +111 0 1064 0.002832 1064,93023,115587,73856,2126,59962,3332,1050 26.750000,20.375000,18.625000,18.000000,17.625000,17.500000,16.750000,16.000000 +112 0 1805 0.088970 1805,15011,12271,1635,31194,14400,106817,2786 23.500000,20.250000,20.000000,18.375000,18.375000,16.625000,16.625000,16.500000 +113 0 15011 0.112116 15011,41500,6061,8303,19895,5395,13268,2247 20.875000,16.500000,16.250000,16.250000,16.000000,15.625000,15.562500,15.125000 +114 0 19976 13.766954 1307,1408,29818,4418,1294,1319,28140,2126 25.875000,20.500000,20.500000,19.125000,18.875000,18.750000,18.750000,18.250000 +115 0 1307 0.333412 1307,1408,1294,29818,28140,1435,1394,3897 22.750000,21.500000,19.500000,18.375000,18.250000,18.125000,17.125000,16.625000 +116 0 1278 0.100233 1278,1531,22876,2813,3797,15087,43193,81873 20.875000,16.500000,15.750000,15.500000,15.500000,15.375000,15.312500,15.312500 +117 0 14696 0.517912 14696,19895,22876,7520,9371,5150,81873,2730 20.000000,18.250000,18.000000,17.750000,17.250000,15.812500,15.750000,15.625000 +118 0 5395 0.113316 5395,25865,18823,2629,14593,50584,3467,119435 22.500000,19.625000,19.000000,17.875000,17.375000,16.625000,16.125000,15.875000 +119 0 1488 3.878300 3897,15087,2813,1531,51288,43193,1488,10495 17.750000,15.812500,15.437500,15.375000,15.312500,15.125000,14.812500,14.375000 +120 0 6812 0.352888 6812,89220,24591,6970,2939,5635,10178,1493 19.375000,18.375000,14.687500,14.125000,13.562500,13.437500,13.437500,13.312500 +121 0 1278 0.032867 1278,1531,8250,29251,3897,2247,2015,1032 21.125000,16.875000,15.562500,15.500000,14.375000,14.187500,13.937500,13.625000 +122 0 29251 0.030470 29251,19213,1407,8776,49944,39447,3897,49377 21.875000,16.000000,15.750000,15.687500,15.375000,15.312500,15.062500,15.000000 +123 0 3897 3.419935 2247,18988,11159,1435,3897,22594,29818,1454 22.750000,20.875000,19.875000,19.750000,19.750000,19.625000,19.375000,18.875000 +124 0 7125 4.009254 12440,1435,28140,1294,29818,7125,22594,8308 20.625000,17.625000,17.625000,17.500000,17.375000,17.000000,17.000000,16.875000 +125 0 1536 0.068083 1536,1294,29598,2015,4025,1261,1278,17558 25.625000,22.250000,21.000000,20.250000,19.875000,19.625000,19.375000,19.250000 +126 0 1420 2.163057 1261,1420,3866,24836,1278,1794,15720,29857 22.250000,20.500000,19.625000,19.500000,19.125000,19.125000,18.250000,17.375000 +127 0 13988 0.374763 13988,15011,1457,7481,47353,87170,1408,107254 21.375000,19.750000,19.250000,17.000000,16.500000,16.250000,16.000000,15.937500 diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp index 2064b76f4..e9591d08a 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp @@ -1732,18 +1732,19 @@ void ssm_update_kernel( {"H", h}, {"G", g}, }; - // The hipRTC module cache is keyed on the name plus this list, while the - // generated signature takes every input's runtime dtype, so an input whose - // dtype is not named here could reuse a module compiled for another dtype + // The CUDA and HIP JIT caches key a module on the kernel name plus this + // list, while the generated signature takes every input's runtime dtype, + // so an input whose dtype is not named here can reuse a module compiled + // for another dtype and read its buffer through the wrong pointer type // (scripts/ci/check_kernel_dtype_keys.py). T and U cover X, D (cast to T - // below) and state_in, and dt is always float32 from compute_dt; A_log, B - // and C are not tied to T, so the HIP port names them too. Appended on - // ROCm only, which leaves the Metal and CUDA kernel names unchanged. - if (mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Rocm) { - template_args.push_back({"TA", A_log.inner.dtype()}); - template_args.push_back({"TB", B.inner.dtype()}); - template_args.push_back({"TC", C.inner.dtype()}); - } + // below) and state_in, and dt is always float32 from compute_dt. A_log, B + // and C are not tied to T: Nemotron-H stores A_log in f32 next to bf16 + // activations, granite in bf16 (#2067). The kernels never name TA, TB or + // TC, so this changes module names only, not the arithmetic; Metal + // already keys its name on every input dtype. + template_args.push_back({"TA", A_log.inner.dtype()}); + template_args.push_back({"TB", B.inner.dtype()}); + template_args.push_back({"TC", C.inner.dtype()}); std::vector inputs = { hidden_states.inner, A_log.inner, B.inner, C.inner, From 3780b0b20765580c2436e0e1cb6b83cd2032aacf Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 4 Oct 2026 21:19:03 +0900 Subject: [PATCH 3/7] docs(rocm): record the HIP SSM update kernel and MLXCEL_SSM_KERNEL - docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md: decode before and after on gfx1151 (granite-4.0-h-tiny-4bit 60.39 -> 88.43 tok/s, Nemotron-3-Nano-30B-A3B-4bit 51.35 -> 74.37 tok/s, medians of three guarded runs), the parity test, the w1ctx512 kernel-against-kernel traces against Metal and against the ROCm graph, and the greedy comparison. - environment-variables.md: MLXCEL_SSM_KERNEL as the cross-backend switch, MLXCEL_SSM_CUDA_KERNEL as its alias. - installation.md ROCm status table and the open item in the 2026-09-30 correctness report point at the new results. - benchmarks/: the raw bench_decode rows for both arms. Refs #2067 --- ...halo-gfx1151_2026-10-04_ssm-kernel-off.csv | 7 ++ ...xhalo-gfx1151_2026-10-04_ssm-kernel-on.csv | 7 ++ .../rocm-correctness-gfx1151-2026-09-30.md | 2 +- ...cm-ssm-update-kernel-gfx1151-2026-10-04.md | 65 +++++++++++++++++++ docs/environment-variables.md | 3 +- docs/installation.md | 1 + 6 files changed, 83 insertions(+), 2 deletions(-) create mode 100644 benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-off.csv create mode 100644 benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-on.csv create mode 100644 docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md diff --git a/benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-off.csv b/benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-off.csv new file mode 100644 index 000000000..2a696dd8c --- /dev/null +++ b/benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-off.csv @@ -0,0 +1,7 @@ +model,model_path,prompt_tokens,generated_tokens,prefill_ms,prefill_tok_s,decode_ms,decode_tok_s,date,hardware,mlxcel_version,build_type,max_tokens,prompt,prompt_target_len,mlxcel_commit,mlx_commit,mlx_rocm_overlay_commit,hip_version +granite-4.0-h-tiny-4bit,models/mlx/granite-4.0-h-tiny-4bit,512,128,1014.85,504.51,2113.20,60.57,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 +granite-4.0-h-tiny-4bit,models/mlx/granite-4.0-h-tiny-4bit,512,128,851.06,601.60,2119.61,60.39,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 +granite-4.0-h-tiny-4bit,models/mlx/granite-4.0-h-tiny-4bit,512,128,1265.81,404.49,2131.73,60.05,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 +NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,models/mlx/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,512,128,1712.54,298.97,2478.21,51.65,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 +NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,models/mlx/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,512,128,1837.73,278.61,2516.24,50.87,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 +NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,models/mlx/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,512,128,1709.97,299.42,2492.62,51.35,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 diff --git a/benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-on.csv b/benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-on.csv new file mode 100644 index 000000000..bc8054ff4 --- /dev/null +++ b/benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-on.csv @@ -0,0 +1,7 @@ +model,model_path,prompt_tokens,generated_tokens,prefill_ms,prefill_tok_s,decode_ms,decode_tok_s,date,hardware,mlxcel_version,build_type,max_tokens,prompt,prompt_target_len,mlxcel_commit,mlx_commit,mlx_rocm_overlay_commit,hip_version +granite-4.0-h-tiny-4bit,models/mlx/granite-4.0-h-tiny-4bit,512,128,868.77,589.34,1448.68,88.36,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 +granite-4.0-h-tiny-4bit,models/mlx/granite-4.0-h-tiny-4bit,512,128,1349.37,379.44,1443.15,88.69,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 +granite-4.0-h-tiny-4bit,models/mlx/granite-4.0-h-tiny-4bit,512,128,842.09,608.01,1447.52,88.43,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 +NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,models/mlx/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,512,128,1728.62,296.19,1709.73,74.87,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 +NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,models/mlx/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,512,128,1772.23,288.90,1756.62,72.87,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 +NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,models/mlx/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit,512,128,1794.99,285.24,1721.18,74.37,2026-10-04,AMD_Radeon_8060S_Graphics_gfx1151_ROCm10.0.0_96GB,0.7.0,release,128,"Hello, how are you today?",512,844bd94c,81ba1c6a,75915908,7.15.26333 diff --git a/docs/benchmark_results/rocm-correctness-gfx1151-2026-09-30.md b/docs/benchmark_results/rocm-correctness-gfx1151-2026-09-30.md index 3f580d7ad..8948637c5 100644 --- a/docs/benchmark_results/rocm-correctness-gfx1151-2026-09-30.md +++ b/docs/benchmark_results/rocm-correctness-gfx1151-2026-09-30.md @@ -275,6 +275,6 @@ The Metal traces were produced with the loop in `benchmarks/logit_traces/rocm_gf ## Open - The second acceptance criterion of #1809 (decided-position mismatch reported per model and width against Metal) is now met for all nine models: fifteen of these sixteen pairs pass and `nemotron-3-nano-30b-a3b` `w1` is inconclusive, with the same checkpoint's `w8` and `w256` pairs passing. The two `w1ctx512` rows, which cover the fused SSM update kernel on Metal, also pass. -- Kernel against kernel for the SSM update still needs the ROCm port in #1814; the `w1ctx512` rows compare Metal's kernel with ROCm's graph. +- Kernel against kernel for the SSM update: the `w1ctx512` rows here compare Metal's kernel with ROCm's graph. The HIP port (#2067) added kernel-against-kernel rows, with zero decided-position mismatches for both hybrids, in [rocm-ssm-update-kernel-gfx1151-2026-10-04.md](rocm-ssm-update-kernel-gfx1151-2026-10-04.md#model-logits-kernel-against-kernel). - `src/models/nemotron_h.rs` prints its loading messages to stdout, which puts five non-trace lines into every Nemotron-H trace and stops `compare_logit_traces.py`. Either the loader should log to stderr or the comparison script should skip them. - No noise floor was measured; the first run's reasoning for not needing one still applies. diff --git a/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md b/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md new file mode 100644 index 000000000..7b3c126a2 --- /dev/null +++ b/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md @@ -0,0 +1,65 @@ +# HIP port of the fused SSM update kernel on gfx1151 (2026-10-04) + +lablup/mlxcel#2067, part of #1814. Before this change every single-token Mamba2 decode step on ROCm ran the ~55-op SSD graph (`ssm_step`), because `ssm_ports()` had no `.rocm` entry and `ssm_kernel_available()` answered with `cu::is_available()` off Apple. The decode profile ([rocm-decode-profile-gfx1151-2026-09-30.md](rocm-decode-profile-gfx1151-2026-09-30.md)) put that graph at 29.8% of granite-4.0-h-tiny's and 20.0% of Nemotron-H's decode GPU time, the largest share of any #1814 port. + +The port (`SSM_HIP_SOURCE` in `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`) is the CUDA kernel with `__shfl_down(acc, o, 32)` for the lane fold and a preprocessor guard that refuses to compile for a wave64 target. `ssm_kernel_available()` now returns `has_kernel_port(ssm_ports())` on every platform, so the model gates (`seq_len == 1 && ssm_kernel_available()` in granitemoehybrid, falcon_h1, plamo2 and nemotron_h) reach it with no model-side change. `MLXCEL_SSM_KERNEL=0` forces the graph on every backend; `MLXCEL_SSM_CUDA_KERNEL=0` is kept as an alias. + +## Environment + +AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5), 96 GiB VRAM carve-out, Debian 13, ROCm 10.0.0 (HIP 7.15.26333). mlxcel `844bd94c` (the PR branch on `origin/main` `c0b71344`), MLX pin `81ba1c6a`, ROCm overlay `75915908`, `cargo build --release --features rocm`. Checkpoints: `mlx-community/granite-4.0-h-tiny-4bit` and `mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit`, the same local directories as the earlier correctness rows. + +## Decode throughput + +`scripts/bench_decode.sh` at its default pp512/tg128 shape, one model per run, `MLXCEL_SSM_KERNEL=0` for the graph and unset for the kernel, on the same binary. Three runs per arm, the arm order alternated between runs. Every run went through `scripts/rocm_gpu_guard.sh` (90 s with `/sys/class/kfd/kfd/proc` empty and no compiler, then 1 Hz monitoring); a parallel development unit shared the GPU, and every attempt that overlapped its work was rejected and rerun. The first run used the guard on `main`, the rest the copy fixed in lablup/mlxcel#2098 (commit `c3eab7b2`), which stops rejecting the command's own exited children; the old guard's rejections were false, never its acceptances. Raw rows: `benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-off.csv` and `..._ssm-kernel-on.csv`. + +| Model | Graph (`MLXCEL_SSM_KERNEL=0`) tok/s | Kernel tok/s | Speedup (medians) | +|---|---|---|---| +| granite-4.0-h-tiny-4bit | 60.57 / 60.39 / 60.05, median 60.39 | 88.36 / 88.69 / 88.43, median 88.43 | 1.46x | +| NVIDIA-Nemotron-3-Nano-30B-A3B-4bit | 51.65 / 50.87 / 51.35, median 51.35 | 74.87 / 72.87 / 74.37, median 74.37 | 1.45x | + +Run-to-run spread is under 1% for granite and under 3% for Nemotron-H in either arm, far below the change. `compare_bench_csv.py --before --after ` (which keeps the last row per model) reports 1.47x and 1.45x. Prefill is not on this path (a prompt is more than one token) and is not compared; it varied between 379 and 608 tok/s for granite across both arms on this shared host. + +## Correctness + +### Kernel against the graph, one step + +`ssm_update_parity_tests` (`src/lib/mlxcel-core/src/ssm_update_parity_tests.rs`) compares the kernel's output and new state with a float32 MLX-op reference of the single-token SSD step at the granite-4.0-h-tiny shape (48 heads of 64, one group, state 128) and the Nemotron-H shape (64 heads of 64, eight groups, state 128, batch 2, clipped dt), f32 and bf16. Tolerances, normalized RMS / max: f32 1e-5 / 1e-4, bf16 1.6e-2 / 7e-2; the f32 state is held to the f32 budget in both. It passes on gfx1151. It fails with normalized RMS 0.65 to 0.89 when the lane fold starts at 8 instead of 16. + +A third test runs the same shape with `A_log` first in bf16 and then in f32 in one process. Nemotron-H stores `A_log` in f32 next to bf16 activations and granite in bf16, and the generated kernel signature takes each input's runtime dtype, so the launch now names the `A_log`, `B` and `C` dtypes in its template arguments (the cache key on CUDA and ROCm). Without them the second launch reuses the bf16 module and the test fails at normalized RMS 0.80. + +### Model logits: kernel against kernel + +The `w1ctx512` traces (`1 128 8 512` with `MLXCEL_TRACE_START_TOKEN=512`) prefill 512 corpus tokens into a fresh cache and then trace one token, so every traced step meets an SSM state and takes the fused step. Until now ROCm ran the graph there and the rows compared Metal's kernel with ROCm's graph ([rocm-correctness-gfx1151-2026-09-30.md](rocm-correctness-gfx1151-2026-09-30.md#the-fused-ssm-kernel-against-the-graph-with-state)). With the port they compare kernel with kernel. New traces: `benchmarks/logit_traces/rocm_gfx1151_96cbce84/`, built at the PR's first commit (the second changes only the kernel's cache name, the same on ROCm in both). `ssmkernel0` is the same binary with `MLXCEL_SSM_KERNEL=0`. `python3 scripts/compare_logit_traces.py --decided 2.0`, Nemotron-H files filtered with `grep -v '^\[NemotronH\] '`. + +| Model | Reference | Candidate | Top-1 disagreement | Decided mismatches | Largest gap | Perplexity ref / cand | +|---|---|---|---|---|---|---| +| granite-4.0-h-tiny | Metal kernel (`metal_m5_3c9edea0`) | ROCm kernel | 3 / 128 | 0 / 60 | 0.125 | 7.090 / 7.216 | +| granite-4.0-h-tiny | Metal kernel | ROCm graph (`ssmkernel0`) | 2 / 128 | 0 / 60 | 0.125 | 7.090 / 7.290 | +| granite-4.0-h-tiny | ROCm graph (`ssmkernel0`) | ROCm kernel | 1 / 128 | 0 / 59 | 0.000 | 7.290 / 7.216 | +| granite-4.0-h-tiny | ROCm graph at `3c9edea0` | ROCm kernel | 1 / 128 | 0 / 60 | 0.125 | 7.234 / 7.216 | +| nemotron-3-nano-30b-a3b | Metal kernel (`metal_m5_3c9edea0`) | ROCm kernel | 3 / 128 | 0 / 71 | 0.250 | 7.489 / 7.442 | +| nemotron-3-nano-30b-a3b | Metal kernel | ROCm graph (`ssmkernel0`) | 4 / 128 | 0 / 71 | 0.250 | 7.489 / 7.460 | +| nemotron-3-nano-30b-a3b | ROCm graph (`ssmkernel0`) | ROCm kernel | 1 / 128 | 0 / 68 | 0.125 | 7.460 / 7.442 | +| nemotron-3-nano-30b-a3b | ROCm graph at `3c9edea0` | ROCm kernel | 4 / 128 | 0 / 68 | 0.125 | 7.489 / 7.442 | + +Every row has zero disagreements on a decided position, and every disagreement puts the reference's token at rank 2 or 3 of the candidate, at a reference gap of 0.25 or less. The kernel and graph traces from one ROCm binary differ (different perplexities, one flipped token each), which shows the `default` traces took the kernel. Against Metal's kernel the ROCm kernel is as close as the ROCm graph was (3 and 3 top-1 disagreements, against 2 and 4 for the graph here and 4 and 7 in the earlier `3c9edea0` comparison). + +The `w1` rows (`1 128 8 0`) have no prefill, so their single token never meets a state and runs the graph with or without the kernel. Against the graph-path traces in `benchmarks/logit_traces/rocm_gfx1151_c5fe9a16/` both are identical at every position (granite 0 / 128, 0 of its decided positions; Nemotron-H 0 / 128, no decided position). They show the change leaves the stateless path alone, not anything about the kernel. + +### Greedy output + +`mlxcel generate -p "The history of the Roman Empire" -n 128 -t 0 --no-chat-template`, with and without `MLXCEL_SSM_KERNEL=0`. Both arms of both models produce coherent text, but they are not byte-identical (on CUDA they were, per [hybrid-ssm-decode-cuda-kernel-gb10-2026-07-10.md](hybrid-ssm-decode-cuda-kernel-gb10-2026-07-10.md)): granite parts at the third generated token ("a rich and complex tapestry" against "a fascinating and complex subject"), Nemotron-H after about a dozen ("a vast and complex tapestry spanning centuries," is common to both). A free-running generation conditions everything after the first flip on different text, which is why the teacher-forced traces above are the comparison; they show the flips land only where the model was undecided. + +## Reproducing + +```bash +cargo build --release --features rocm --bin mlxcel --bin mlxcel-bench-decode --example logit_trace +cargo test --release --features rocm -p mlxcel-core --lib ssm_update_parity_tests -- --test-threads=1 +for m in granite-4.0-h-tiny-4bit NVIDIA-Nemotron-3-Nano-30B-A3B-4bit; do + MLXCEL_SSM_KERNEL=0 scripts/rocm_gpu_guard.sh -- env MODELS_DIR=models/mlx ./scripts/bench_decode.sh models/mlx/$m --output off.csv + scripts/rocm_gpu_guard.sh -- env MODELS_DIR=models/mlx ./scripts/bench_decode.sh models/mlx/$m --output on.csv +done +python3 scripts/compare_bench_csv.py --before off.csv --after on.csv +``` + +The trace loop is in `benchmarks/logit_traces/rocm_gfx1151_96cbce84/README.md`. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 2e03ac38e..977087194 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -550,7 +550,8 @@ recommended as normal deployment settings. | `MLXCEL_QMV_MULTIROW` | `0` disables; any other value or unset enables | on | **CUDA only.** Gates the weight-amortizing multirow qmv path (issue #725). With broadcast weights and `2 <= M*B <= 8` input rows (batched decode with B in [2,8), speculative-verify `[1, K]` forwards), one warp applies each dequantized weight tile to every input row instead of launching one weight-rereading block column per row, so weight DRAM traffic is O(1) in the row count. Per-row outputs are bit-identical to the stock per-row launches (pinned by `qmv_multirow_matches_per_row_qmv_bitwise`); classic `M*B == 1` decode is untouched. On GB10 this flips B=1 MTP speculative decoding from 0.52-0.77x to 1.31-1.46x and lifts 4-client aggregate serving decode from ~50 to ~74 tok/s on llama-3.1-8b-4bit. Set to `0` to force the stock per-row dispatch (A/B, rollback). See `docs/benchmark_results/qmv-multirow-gb10-2026-07-11.md`. | | `MLXCEL_QMV_ROWS_PER_WARP` | `1`, `2` or `4`; any other value (or unset) uses the per-architecture default | `2` on sm_70, `1` elsewhere | **CUDA only.** Output rows each warp covers in `qmv`, for classic `M*B == 1` decode. The stock kernel gives one warp one output row and has that warp read the whole activation vector, so a launch moves `n * k * 2` bytes of activations against `n * k / 2` bytes of 4-bit weights. Measured on a V100 over a controlled 4-bit against 8-bit pair of gemma-4-12B-it with identical launch counts, that activation term is 84% of a 4-bit launch, which is why decode below Ampere is not weight-bandwidth bound. `R` rows per warp load the activation tile once and apply it to all `R` weight rows, dividing the activation traffic by `R` while leaving weight traffic alone; the per-row arithmetic is unchanged, so output is byte-identical to `R = 1` at temperature 0. The cost is that the same output rows are covered by `R` times fewer warps and the accumulator array grows to `R * elems_per_thread`. On a V100 at 4 bits, R = 2 moves decode from 44.7 to 60.1 tok/s on gemma-4-12B-it-4bit under the `MLXCEL_QMV_MIN_BLOCKS` default of 3, and from 20.6 to 24.9 on qwen3.8-27B-4bit; R = 4 gives it all back on the 12B (96 to 142 registers, 1 to 2 resident blocks per SM against the stock kernel's 4 to 5) while costing nothing on the wider 27B. Turing defaults to 1 because no sm_75 device was available to measure and the R = 4 reversal shows the tradeoff does not extrapolate. `1` routes to the stock kernel with the stock grid, so it is the rollback. Read once per process. | | `MLXCEL_QMV_MIN_BLOCKS` | `1`, `3` or `4`; any other value (or unset) uses the per-architecture default | `3` on sm_70, `1` elsewhere | **CUDA only.** Resident-block floor handed to `__launch_bounds__` on the row-blocked `qmv` kernel, which is how ptxas is told to cap the register count: for a 256-thread block, 3 caps at 85 registers per thread and 4 caps at 64. This is the knob, not the row count. Measured on a V100 at 4 bits on gemma-4-12B-it, decode as the slope between `-n 60` and `-n 200`: R = 2 unconstrained sits at 92 to 95 registers and 2 blocks per SM for 56.0 tok/s, the same R = 2 under a floor of 3 sits at 80 registers and 3 blocks for 60.1, and under a floor of 4 it sits at 64 registers and falls back to 54.9, so the budget can be set too tight as well as too loose. Nothing spills at any floor (`LOCAL` is 0 in every instantiation). R = 3 under the same floor of 3 has identical registers and residency and moves 17% fewer memory transactions per output row, and it still lands at 58.8, which is why the floor rather than the row count carries the default. `1` is the unconstrained case and is the rollback. Read once per process. | -| `MLXCEL_SSM_CUDA_KERNEL` | `0` disables; any other value or unset enables | on | **CUDA only.** Gates the fused single-token SSM decode kernel port (issue #631). The Metal-only `ssm_update_kernel` (one launch replacing the ~55-op SSD scan graph per SSM layer) now has a `mx.fast.cuda_kernel` port, selected automatically on the CUDA backend for the hybrid SSM models (granite-4.0-h, falcon-h1, plamo-2, nemotron-h). Fixes the 0.29-0.36x hybrid-SSM decode ratio vs Metal M1 Ultra; measured 2.6-4.5x decode speedup on the granite/falcon family on GB10, greedy parity byte-identical, pure-mamba2 control unchanged. Set to `0` to force the graph path (A/B, rollback). See `docs/benchmark_results/hybrid-ssm-decode-cuda-kernel-gb10-2026-07-10.md`. | +| `MLXCEL_SSM_KERNEL` | `0` disables; any other value or unset enables | on | Gates the fused single-token SSM decode kernel (`ssm_update_kernel`) on every backend that has a port: Metal, CUDA (issue #631) and ROCm (issue #2067). It replaces the ~55-op SSD scan graph (`ssm_step`) with one launch per SSM layer per decode step for the hybrid SSM models (granite-4.0-h, falcon-h1, plamo-2, nemotron-h). `ssm_kernel_available()` answers from the kernel's port table (`ssm_ports()`), so a backend with no port takes the graph path on its own. Set to `0` to force the graph path (A/B, rollback). Measured on GB10 (CUDA): 2.6-4.5x decode on the granite/falcon family, see `docs/benchmark_results/hybrid-ssm-decode-cuda-kernel-gb10-2026-07-10.md`; on gfx1151 (ROCm): see `docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md`. | +| `MLXCEL_SSM_CUDA_KERNEL` | `0` disables; any other value or unset enables | on | Older name of `MLXCEL_SSM_KERNEL`, from when the CUDA port was the only one off Apple, kept as an alias so existing scripts still work. Either variable set to `0` forces the graph path, now on every backend (Metal and ROCm included). | | `MLXCEL_FUSED_QK_NORM` | `1`/`true`/`on`/`yes` enable; any other value or unset disables | off (opt-in) | Fused single-token QKV projection + Q/K RMSNorm + RoPE kernel (#326) for Qwen3 and Qwen3-MoE decode. Opt-in: set to `1` to enable. Matches the graph path within RMS < 5e-3 (the reduction is over the transpose-invariant head_dim axis), but greedy temp-0 is not byte-identical over long generation; on CUDA the graph path is itself non-deterministic run-to-run from GPU FP-reduction order, while the fused path is deterministic, so its output stays inside the graph baseline's own envelope. The kernel cuts Rust/C++ FFI crossings rather than MLX op count, so it does not speed up the GPU/bandwidth-bound decode loop: on M1 Ultra it measured 1 to 3.4% slower (qwen3-0.6b 275 vs 284, qwen3-8b 82.3 vs 83.2 tok/s); on GB10/CUDA (SM 12.1) it is also slower (qwen3-0.6b 0.96x, qwen3-8b ~1.0x, qwen3-30b-a3b 0.92x fused/graph; see `docs/benchmark_results/fused-qk-norm-decode-gb10.md`), so there is no per-backend win. Ships as a reusable shared primitive for the deferred QK-norm families and stays opt-in (default off) on every measured backend (M1 Ultra, M5 Max, GB10/CUDA), mirroring the opt-in `MLXCEL_FUSED_MOE_RELU2`. Active only when `l == 1` (decode) and weights are quantized. | | `MLXCEL_FUSED_ADD_RMSNORM` | `0`/`false`/`off`/`no` disable; `1`/`true`/`on`/`yes` enable; unset or unrecognised keeps the default | off (opt-in) | Fused residual-add + RMSNorm decode kernel (#905). Computes `new_residual = x + residual` and `normed = rms_norm(new_residual) * (weight_bias + weight)` in one dispatch instead of an elementwise `Add` followed by `fast::rms_norm`, removing one kernel launch and one full-width intermediate at every pre-norm residual join. Adopted by the Llama3-family block (which also serves Qwen2/Qwen2.5 and every VLM whose text backbone is `Llama3Model`/`Qwen2Model`) and by Gemma, where the `(1 + w)` convention rides in as a scalar `weight_bias = 1.0` folded in the weight's own dtype rather than as a second weight tensor. Numerics are pinned to MLX's own `rms_norm` kernel (fp32 accumulation over the dtype-rounded sum, `precise::rsqrt`, the `x * inv_mean` product rounded to the activation dtype before the weight multiply), so the fused and unfused paths agree to within a small multiple of one ulp; greedy temp-0 is not guaranteed byte-identical over long generation, because a near-tie argmax can land on either side of that difference. Set to `0` to force the `add` + `fast_rms_norm` pair (A/B, rollback). The compiled-in default lives in `FUSED_ADD_RMSNORM_DEFAULT` in `src/lib/mlxcel-core/src/layers.rs`. Read once at first use and cached for the process lifetime. | Ships **opt-in (default off)**: the op-level microbench on M1 Ultra (Metal, f16, hidden {2048, 4096, 8192} x batch {1, 4, 8}, three repetitions) measured the fused path at parity, speedups scattering 0.77x to 1.14x with no structure and the extremes not reproducing. Per-iteration time stayed near 280-410us across a 4x change in hidden size, so the harness is dominated by dispatch cost rather than the op. Issue #905's measure-then-keep policy lands a fusion unwired absent a demonstrated win, matching how `MLXCEL_FUSED_QK_NORM` (#326) shipped. The production decode path agrees: on M1 Ultra at pp512/tg128, Llama 3.1 8B decodes at 113.18 tok/s off and 113.01 on, Qwen2.5 7B at 115.38 and 115.34, with prefill unchanged. Set to `1` to opt in. See `docs/benchmark_results/fused-norm-rope-m1ultra-2026-07-31.md` and `docs/benchmark_results/fused-add-rmsnorm-decode-m1ultra-2026-09-27.md`. | `MLXCEL_FUSED_ROPE_APPEND` | `0`/`false`/`off`/`no` disable; `1`/`true`/`on`/`yes` enable; unset or unrecognised keeps the default | off (opt-in) | Fused q/k RoPE + KV-append-layout decode kernel (#905). Reads the row-contiguous fused-QKV projection output whole, applies rotary embedding to the q and k blocks, and emits q, k and v already in their consumers' layouts, replacing three trailing-axis slices, three reshape/transpose pairs and two `fast::rope` calls with one dispatch. The rotation is transcribed from MLX's own `rope.metal`, so it agrees with `fast::rope` including the `traditional` convention and partial `rope_dims`. The trig call is the one thing chosen per backend rather than transcribed: each body calls whatever MLX's own RoPE kernel for that backend calls, so the Metal body uses `metal::fast::cos`/`sin` after `rope.metal` and the CUDA body uses libdevice `cosf`/`sinf` after `rope.cu`. Using the CUDA fast intrinsics `__cosf`/`__sinf` there instead cost about 1.2e-2 rad of applied phase at absolute position 131071 and failed the parity suite at the two large-offset cases (issue #1049). Adopted by the Llama3-family dense decode path, except for a checkpoint whose `rope_scaling` block selects a frequency table: the kernel derives its frequencies from `rope_base` and has no parameter for a table, so such a model takes the graph path (#1355), and a one-time notice on stderr says so when this variable asked for the kernel. `traditional` and the position scale are real parameters and are honored, so neither triggers that notice for this kernel. The paged block-pool destination layout the kernel also supports is implemented and tested but not wired, because the batched paged decode path belongs to issue #899. The kernel emits the append payload rather than writing the cache slab, because MLX custom-kernel outputs are always freshly allocated while the existing `slice_update` append donates the slab buffer, so an in-kernel write would trade an O(1) donated update for an O(capacity) copy. Set to `0` to force the reshape/transpose/`fast_rope` graph (A/B, rollback). The compiled-in default lives in `FUSED_ROPE_APPEND_DEFAULT` in `src/lib/mlxcel-core/src/layers.rs`. Read once at first use and cached for the process lifetime. | Ships **opt-in (default off)** for the same measured reason as `MLXCEL_FUSED_ADD_RMSNORM`, with one cell consistently below parity rather than merely noisy: hidden 8192 at batch 1 measured 0.94x, 0.90x and 0.89x across three repetitions. Set to `1` to opt in. On GB10/CUDA the same harness measures 1.02x to 2.06x with every cell at or above parity, a different picture from Apple Silicon that has not been acted on; flipping the default is its own decision with its own evidence bar. See `docs/benchmark_results/fused-norm-rope-m1ultra-2026-07-31.md` and `docs/benchmark_results/fused-rope-cuda-trig-gb10-2026-08-06.md`. diff --git a/docs/installation.md b/docs/installation.md index b1c5848b2..c7b518ca0 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -553,6 +553,7 @@ so `mlxcel inspect` and `--estimate-memory` give the same answers as before. | NVFP4 checkpoints | No native kernel. ModelOpt NVFP4 checkpoints (for example the Gemma 4 NVFP4 exports) are converted to affine 4-bit at load with no environment variables, and the load log names the route and the reason. A layer that cannot be converted fails the load with its name and the reason. MLX-native NVFP4 exports (`"mode": "nvfp4"` in `config.json`, such as `mlx-community/*-nvfp4`) have no load-time conversion and are refused at load; use an affine export instead (lablup/mlxcel#1806). | | Affine MoE models (for example Qwen3-30B-A3B) | Run natively. The fused MoE path used to abort on ROCm; lablup/mlxcel#1803 routes it to the MLX graph fallback, so `MLXCEL_FUSED_MOE=0` is no longer needed. | | mlxcel's fused kernels (sampling, fused norm, RoPE + KV append, paged attention) | Run as MLX graph fallbacks on the paths that have one (lablup/mlxcel#1803); ROCm ports are lablup/mlxcel#1814. | +| Hybrid SSM decode (granite-4.0-h, falcon-h1, plamo-2, Nemotron-H) | The single-token Mamba2 SSM update runs the fused HIP kernel instead of the ~55-op SSD graph (lablup/mlxcel#2067): decode on `gfx1151` went from 60.4 to 88.4 tok/s on granite-4.0-h-tiny-4bit and from 51.4 to 74.4 tok/s on Nemotron-3-Nano-30B-A3B-4bit, with model logits matching the graph and Metal's kernel at every decided position ([results](benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md)). `MLXCEL_SSM_KERNEL=0` forces the graph. | | GPU faults | Reported as errors (lablup/mlxcel#1804). A launch HIP rejects (an oversized block, no code object for the device) fails the evaluation that issued it and the device stays usable. An asynchronous fault (an out-of-bounds access) fails the evaluation waiting on it within about a second instead of hanging or returning NaN; the server fails that request and answers later ones with the same error, because after a queue fault the HIP runtime rejects every call for the rest of the process, so the process has to restart, and its shutdown may need a SIGKILL (HIP's teardown waits on callbacks the faulted queue never runs). A kernel that never finishes is not detected; `MLX_ROCM_GPU_WATCHDOG_SECS` (default off) fails any single host wait that outlives it. | | Memory estimation on UMA hosts | Correct. Measured on the tested configuration: the ROCm allocator reports a nonzero cap (76.80 GiB of the 96 GiB carve-out), which the estimator reads before it would ever reach host RAM, so nothing that fits the carve-out is refused for that reason (lablup/mlxcel#1805). | | Diagnostics | Report the AMD vendor, device name, `gfx` target and device memory, and no longer print a CUDA compute capability for it (lablup/mlxcel#1805). A binary whose compiled `gfx` list does not cover the device refuses to start rather than failing at the first kernel launch. | From bd2e8f313be93d5e60bc0db8a5bdb48802945387 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 4 Oct 2026 21:31:22 +0900 Subject: [PATCH 4/7] fix(kernels): keep the SSM kernel off the CPU device and fix review notes - ssm_kernel_available() answers false when the default device is not the GPU, as mamba1_scan_kernel_accepts already does. With MLXCEL_DEVICE=cpu on a ROCm build the gate used to fall to the graph only because the predicate was false there; with the port it would have launched a custom kernel on the CPU stream and thrown. - The HIP wave32 guard comments now say what holds: ROCm 10's AMD clang 23 defines neither __AMDGCN_WAVEFRONT_SIZE spelling for gfx1151, gfx942 or gfx90a, so the guard is inert there, and the width-32 shuffle keeps the fold inside one row on wave64 anyway. - ssm_update_parity_tests: granite runs with its real time_step_limit (0.001, 100.0), the dt-window and reference comments are corrected, a head_dim-60 shape covers the padded rows of the last threadgroup, and the tests skip when the default device is the CPU. - Results doc: a reproduce loop that keeps every run's rows, a note that Nemotron-H's gate selects the whole fused mixer, and the CPU-device rule; installation.md says falcon-h1 and plamo-2 were not run on ROCm. Validated with ssm_update_parity_tests on gfx1151 (4 passed). Refs #2067 --- ...cm-ssm-update-kernel-gfx1151-2026-10-04.md | 19 ++++++-- docs/installation.md | 2 +- src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h | 5 ++- src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp | 37 +++++++++------- src/lib/mlxcel-core/src/lib.rs | 5 ++- .../src/ssm_update_parity_tests.rs | 44 ++++++++++++++----- 6 files changed, 78 insertions(+), 34 deletions(-) diff --git a/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md b/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md index 7b3c126a2..1381a3e39 100644 --- a/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md +++ b/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md @@ -2,7 +2,7 @@ lablup/mlxcel#2067, part of #1814. Before this change every single-token Mamba2 decode step on ROCm ran the ~55-op SSD graph (`ssm_step`), because `ssm_ports()` had no `.rocm` entry and `ssm_kernel_available()` answered with `cu::is_available()` off Apple. The decode profile ([rocm-decode-profile-gfx1151-2026-09-30.md](rocm-decode-profile-gfx1151-2026-09-30.md)) put that graph at 29.8% of granite-4.0-h-tiny's and 20.0% of Nemotron-H's decode GPU time, the largest share of any #1814 port. -The port (`SSM_HIP_SOURCE` in `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`) is the CUDA kernel with `__shfl_down(acc, o, 32)` for the lane fold and a preprocessor guard that refuses to compile for a wave64 target. `ssm_kernel_available()` now returns `has_kernel_port(ssm_ports())` on every platform, so the model gates (`seq_len == 1 && ssm_kernel_available()` in granitemoehybrid, falcon_h1, plamo2 and nemotron_h) reach it with no model-side change. `MLXCEL_SSM_KERNEL=0` forces the graph on every backend; `MLXCEL_SSM_CUDA_KERNEL=0` is kept as an alias. +The port (`SSM_HIP_SOURCE` in `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`) is the CUDA kernel with `__shfl_down(acc, o, 32)` for the lane fold. It carries the `__AMDGCN_WAVEFRONT_SIZE` `#error` guard the #1814 ports share, which is inert with ROCm 10's AMD clang 23 (it defines neither macro for gfx1151, gfx942 or gfx90a); the fold does not need it, because the shuffle width of 32 keeps each reduction inside the 32 lanes of one row on a wave64 target too. `ssm_kernel_available()` now returns `has_kernel_port(ssm_ports())` on every platform, so the model gates (`seq_len == 1 && ssm_kernel_available()` in granitemoehybrid, falcon_h1, plamo2 and nemotron_h) reach it with no model-side change. On Nemotron-H that gate (`nemotron_h.rs`, `NemotronHMamba2Mixer::forward`) selects `fused_mamba2_forward`, which runs the whole single-token mixer (input projection, convolution, this kernel, gated norm, output projection) as one C++ call, so its decode figures below measure that path against the Rust graph mixer, not the SSM step alone. The predicate also answers false when the default device is the CPU (`MLXCEL_DEVICE=cpu`), where a custom kernel cannot run. `MLXCEL_SSM_KERNEL=0` forces the graph on every backend; `MLXCEL_SSM_CUDA_KERNEL=0` is kept as an alias. ## Environment @@ -55,9 +55,20 @@ The `w1` rows (`1 128 8 0`) have no prefill, so their single token never meets a ```bash cargo build --release --features rocm --bin mlxcel --bin mlxcel-bench-decode --example logit_trace cargo test --release --features rocm -p mlxcel-core --lib ssm_update_parity_tests -- --test-threads=1 -for m in granite-4.0-h-tiny-4bit NVIDIA-Nemotron-3-Nano-30B-A3B-4bit; do - MLXCEL_SSM_KERNEL=0 scripts/rocm_gpu_guard.sh -- env MODELS_DIR=models/mlx ./scripts/bench_decode.sh models/mlx/$m --output off.csv - scripts/rocm_gpu_guard.sh -- env MODELS_DIR=models/mlx ./scripts/bench_decode.sh models/mlx/$m --output on.csv +# bench_decode.sh truncates --output, so every run writes its own file. +for i in 1 2 3; do + if [ $((i % 2)) = 1 ]; then arms="off on"; else arms="on off"; fi + for m in granite-4.0-h-tiny-4bit NVIDIA-Nemotron-3-Nano-30B-A3B-4bit; do + for arm in $arms; do + if [ $arm = off ]; then e=MLXCEL_SSM_KERNEL=0; else e=MLXCEL_SSM_KERNEL=1; fi + env $e scripts/rocm_gpu_guard.sh -- env MODELS_DIR=models/mlx \ + ./scripts/bench_decode.sh models/mlx/$m --no-cooldown --output $arm-${m%%-*}-r$i.csv + done + done +done +for arm in off on; do + head -1 $arm-granite-r1.csv > $arm.csv + for f in $arm-*-r*.csv; do tail -n +2 "$f" >> $arm.csv; done done python3 scripts/compare_bench_csv.py --before off.csv --after on.csv ``` diff --git a/docs/installation.md b/docs/installation.md index c7b518ca0..a6dac83c8 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -553,7 +553,7 @@ so `mlxcel inspect` and `--estimate-memory` give the same answers as before. | NVFP4 checkpoints | No native kernel. ModelOpt NVFP4 checkpoints (for example the Gemma 4 NVFP4 exports) are converted to affine 4-bit at load with no environment variables, and the load log names the route and the reason. A layer that cannot be converted fails the load with its name and the reason. MLX-native NVFP4 exports (`"mode": "nvfp4"` in `config.json`, such as `mlx-community/*-nvfp4`) have no load-time conversion and are refused at load; use an affine export instead (lablup/mlxcel#1806). | | Affine MoE models (for example Qwen3-30B-A3B) | Run natively. The fused MoE path used to abort on ROCm; lablup/mlxcel#1803 routes it to the MLX graph fallback, so `MLXCEL_FUSED_MOE=0` is no longer needed. | | mlxcel's fused kernels (sampling, fused norm, RoPE + KV append, paged attention) | Run as MLX graph fallbacks on the paths that have one (lablup/mlxcel#1803); ROCm ports are lablup/mlxcel#1814. | -| Hybrid SSM decode (granite-4.0-h, falcon-h1, plamo-2, Nemotron-H) | The single-token Mamba2 SSM update runs the fused HIP kernel instead of the ~55-op SSD graph (lablup/mlxcel#2067): decode on `gfx1151` went from 60.4 to 88.4 tok/s on granite-4.0-h-tiny-4bit and from 51.4 to 74.4 tok/s on Nemotron-3-Nano-30B-A3B-4bit, with model logits matching the graph and Metal's kernel at every decided position ([results](benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md)). `MLXCEL_SSM_KERNEL=0` forces the graph. | +| Hybrid SSM decode (granite-4.0-h, falcon-h1, plamo-2, Nemotron-H) | The single-token Mamba2 SSM update runs the fused HIP kernel instead of the ~55-op SSD graph (lablup/mlxcel#2067): decode on `gfx1151` went from 60.4 to 88.4 tok/s on granite-4.0-h-tiny-4bit and from 51.4 to 74.4 tok/s on Nemotron-3-Nano-30B-A3B-4bit, with model logits matching the graph and Metal's kernel at every decided position ([results](benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md)). Measured on granite and Nemotron-H; falcon-h1 and plamo-2 take the same gate but have not been run on ROCm. `MLXCEL_SSM_KERNEL=0` forces the graph. | | GPU faults | Reported as errors (lablup/mlxcel#1804). A launch HIP rejects (an oversized block, no code object for the device) fails the evaluation that issued it and the device stays usable. An asynchronous fault (an out-of-bounds access) fails the evaluation waiting on it within about a second instead of hanging or returning NaN; the server fails that request and answers later ones with the same error, because after a queue fault the HIP runtime rejects every call for the rest of the process, so the process has to restart, and its shutdown may need a SIGKILL (HIP's teardown waits on callbacks the faulted queue never runs). A kernel that never finishes is not detected; `MLX_ROCM_GPU_WATCHDOG_SECS` (default off) fails any single host wait that outlives it. | | Memory estimation on UMA hosts | Correct. Measured on the tested configuration: the ROCm allocator reports a nonzero cap (76.80 GiB of the 96 GiB carve-out), which the estimator reads before it would ever reach host RAM, so nothing that fits the carve-out is refused for that reason (lablup/mlxcel#1805). | | Diagnostics | Report the AMD vendor, device name, `gfx` target and device memory, and no longer print a CUDA compute capability for it (lablup/mlxcel#1805). A binary whose compiled `gfx` list does not cover the device refuses to start rather than failing at the first kernel launch. | diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h index f6770673f..90a3ced83 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h @@ -1958,8 +1958,9 @@ std::unique_ptr fused_moe_forward( ); // Whether the fused single-token SSM update kernel can run: the GPU backend -// has a port in ssm_ports() (Metal, CUDA, ROCm) and MLXCEL_SSM_KERNEL=0 (or -// its older alias MLXCEL_SSM_CUDA_KERNEL=0) is not set. +// has a port in ssm_ports() (Metal, CUDA, ROCm), the default device is the +// GPU, and MLXCEL_SSM_KERNEL=0 (or its older alias MLXCEL_SSM_CUDA_KERNEL=0) +// is not set. bool ssm_kernel_available(); // Mamba1 selective scan fused over the sequence (Jamba, issue #2005). diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp index e9591d08a..caeaa207d 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp @@ -637,18 +637,19 @@ namespace { // outputs, grid and template arguments as the CUDA port above; the body // differs in the lane reduction only. `__shfl_down_sync` is a HIP // compatibility shim that ignores its mask, so the native - // `__shfl_down(var, delta, width)` is used with the width stated. The fold - // starts at 16 because one 32-lane wavefront owns one (head, row) pair: - // the launch's threadgroup is (32, 8, 1), so each `threadIdx.y` row is one - // wave on a wave32 target. On a wave64 target (CDNA: gfx90a, gfx942) a - // 64-lane wave would span two rows and this fold would mix them, which is - // why the source refuses to compile there. + // `__shfl_down(var, delta, width)` is used with the width stated. One + // (head, row) pair is reduced by the 32 lanes that share a `threadIdx.y` + // (the threadgroup is (32, 8, 1)), and the explicit width of 32 keeps the + // fold inside those 32 lanes: on a wave32 target (gfx11, gfx12) they are + // the whole wave, and on a wave64 target (CDNA) the width splits the wave + // into two 32-lane segments that are again exactly two rows. static const char* SSM_HIP_SOURCE = R"( - // A fold that starts at 16 is correct only for a 32-lane wavefront; - // on wave64 it would return a finite, plausible, wrong sum. Checked - // with the preprocessor because `static_assert(warpSize == 32)` does - // not compile in HIP (`warpSize` is not a constant expression). Both - // macro spellings are checked; the trailing-underscore one is older. + // The wave32 guard every #1814 port carries (#2067 port + // requirements), spelled as the bitlinear port spells it because + // `static_assert(warpSize == 32)` does not compile in HIP. It is a + // no-op with ROCm 10's AMD clang 23, which defines neither macro for + // gfx1151, gfx942 or gfx90a. This fold does not rely on it: the + // shuffle width of 32 keeps it inside one row on wave64 as well. #if defined(__AMDGCN_WAVEFRONT_SIZE__) && __AMDGCN_WAVEFRONT_SIZE__ != 32 #error "ssm_kernel_hip assumes a 32-lane wavefront" #endif @@ -1678,10 +1679,16 @@ bool ssm_kernel_available() { return false; } } - // Read from the same table the dispatch in ssm_update_kernel selects - // from, so the Rust gates (`seq_len == 1 && ssm_kernel_available()`) can - // `expect` the launch: a true here means select_kernel_port will not - // refuse. + // Custom kernels run only on the GPU stream; with MLXCEL_DEVICE=cpu the + // model gates then take the ssm_step graph on the CPU instead of the + // launch throwing (the same check as mamba1_scan_kernel_accepts). + if (mlx::core::default_device() != mlx::core::Device::gpu) { + return false; + } + // Otherwise read from the same table the dispatch in ssm_update_kernel + // selects from, so the Rust gates (`seq_len == 1 && + // ssm_kernel_available()`) can `expect` the launch: a true here means + // select_kernel_port will not refuse. return mlxcel::has_kernel_port(ssm_ports()); } diff --git a/src/lib/mlxcel-core/src/lib.rs b/src/lib/mlxcel-core/src/lib.rs index bb17baf15..edd3a8b1f 100644 --- a/src/lib/mlxcel-core/src/lib.rs +++ b/src/lib/mlxcel-core/src/lib.rs @@ -1839,10 +1839,11 @@ mod ffi { bits: i32, ) -> Result>; - // SSM (Mamba2) fused Metal kernel. + // SSM (Mamba2) fused kernel. /// Whether the fused single-token SSM update kernel can run: the /// resolved GPU backend has a port in `ssm_ports()` (Metal, CUDA, - /// ROCm) and neither `MLXCEL_SSM_KERNEL=0` nor its older alias + /// ROCm), the default device is the GPU, and neither + /// `MLXCEL_SSM_KERNEL=0` nor its older alias /// `MLXCEL_SSM_CUDA_KERNEL=0` is set. fn ssm_kernel_available() -> bool; diff --git a/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs b/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs index 1136d4c8d..7bf5dc318 100644 --- a/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs +++ b/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs @@ -26,9 +26,12 @@ //! ``` //! //! with B and C shared by the heads of a group. The reference below composes -//! exactly that from MLX ops in float32 (the graph path also promotes to -//! float32), and the kernel's output and new state are compared with it at the -//! granite-4.0-h-tiny and Nemotron-H mixer shapes, in f32 and bf16. +//! exactly that from MLX ops in float32, as the kernel computes it (the graph +//! path promotes x, B, C and dt to float32 too, but forms `-exp(A_log)` in +//! `A_log`'s own dtype, so for a bf16 `A_log` it rounds A where the kernel does +//! not). The kernel's output and new state are compared with it at the +//! granite-4.0-h-tiny and Nemotron-H mixer shapes, in f32 and bf16, and at a +//! head size that leaves padded rows in the last threadgroup. //! //! Tolerances are normalized RMS and normalized max deviation (both divided by //! the reference's RMS): f32 1e-5 / 1e-4 and bf16 1.6e-2 / 7e-2, the bf16 @@ -36,8 +39,9 @@ //! every case, as the models carry it, so it is held to the f32 budget even //! when the activations are bf16. //! -//! These tests return early only on a build with no GPU backend, or when the -//! `MLXCEL_SSM_KERNEL=0` / `MLXCEL_SSM_CUDA_KERNEL=0` kill switch is set. On +//! These tests return early only on a build with no GPU backend, when the +//! default device is the CPU, or when the `MLXCEL_SSM_KERNEL=0` / +//! `MLXCEL_SSM_CUDA_KERNEL=0` kill switch is set. On //! Metal, CUDA and ROCm a false `ssm_kernel_available()` is itself a defect //! (each has a port), so the test fails there instead of skipping. //! @@ -75,6 +79,10 @@ fn skip() -> bool { eprintln!("skipping ssm_update_parity_tests: no GPU backend in this build"); return true; } + if !default_device_is_gpu() { + eprintln!("skipping ssm_update_parity_tests: the default device is the CPU"); + return true; + } if kill_switch_set() { eprintln!("skipping ssm_update_parity_tests: SSM kernel kill switch is set"); return true; @@ -132,11 +140,12 @@ struct Shape { head_dim: i32, groups: i32, state_dim: i32, - /// `time_step_limit` of the model. + /// The `(min, max)` clip applied to `softplus(dt + dt_bias)`. dt_limits: (f32, f32), } -/// granite-4.0-h-tiny: 48 heads of 64, one group, state 128, no dt clip. +/// granite-4.0-h-tiny: 48 heads of 64, one group, state 128, with granite's +/// default `time_step_limit` of `(0.001, 100.0)`. const GRANITE_TINY: Shape = Shape { name: "granite-4.0-h-tiny", batch: 1, @@ -144,12 +153,14 @@ const GRANITE_TINY: Shape = Shape { head_dim: 64, groups: 1, state_dim: 128, - dt_limits: (0.0, f32::INFINITY), + dt_limits: (0.001, 100.0), }; /// Nemotron-3-Nano-30B-A3B: 64 heads of 64 in 8 groups, state 128. Run at -/// batch 2 with a clipping dt window so the group/batch indexing and the clip -/// are both exercised. +/// batch 2 with a dt window of `(1e-3, 0.1)` (the checkpoint's +/// `time_step_min` / `time_step_max`; its `time_step_limit` is `(0, inf)`), +/// narrow enough to clip, so the group/batch indexing and the clip are both +/// exercised. const NEMOTRON_H: Shape = Shape { name: "nemotron-h", batch: 2, @@ -160,6 +171,18 @@ const NEMOTRON_H: Shape = Shape { dt_limits: (1e-3, 0.1), }; +/// A head size that is not a multiple of the threadgroup's 8 rows, so the last +/// threadgroup has rows past `Dh` that must return before touching memory. +const PADDED_ROWS: Shape = Shape { + name: "head_dim 60", + batch: 1, + heads: 4, + head_dim: 60, + groups: 2, + state_dim: 64, + dt_limits: (0.0, f32::INFINITY), +}; + struct Case { x: UniquePtr, a_log: UniquePtr, @@ -307,6 +330,7 @@ fn ssm_update_kernel_matches_graph_step_f32() { } check(&GRANITE_TINY, 2067, dtype::FLOAT32, dtype::FLOAT32); check(&NEMOTRON_H, 2068, dtype::FLOAT32, dtype::FLOAT32); + check(&PADDED_ROWS, 2073, dtype::FLOAT32, dtype::FLOAT32); } #[test] From 3807103807fdd7f47e51180ca94f792ede40f9c6 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 4 Oct 2026 21:47:58 +0900 Subject: [PATCH 5/7] fix(kernels): refuse SSM update shapes the kernel cannot index ssm_update_kernel passes its shapes to the Metal, CUDA and HIP kernels as template constants and indexes raw buffers with them. A state width that is not a multiple of 32 left state columns unwritten (and with one under 32, the whole state came from uninitialized memory), heads that do not divide into groups read past B and C, and a mismatched state, A_log, D or dt read out of bounds. The graph path rejects such configs inside reshape or repeat; the kernel path did not. The host now checks rank, a single token, matching batch, the group divisibility, the state width and every input's element count, and throws std::invalid_argument, which the bridge returns to Rust as an error. This was shared by every port; the HIP port made it reachable on ROCm. ssm_update_kernel_refuses_unsupported_shapes covers a 48-wide state and 4 heads in 3 groups. The env mutations in the kill-switch test now carry a SAFETY note. Validated on gfx1151: ssm_update_parity_tests 5 passed; greedy 128-token output for granite-4.0-h-tiny-4bit and Nemotron-3-Nano-30B-A3B-4bit identical to the run before this change. Refs #2067 --- src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp | 24 +++++++++ .../src/ssm_update_parity_tests.rs | 53 +++++++++++++++++++ 2 files changed, 77 insertions(+) diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp index caeaa207d..506d9feb5 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp @@ -1708,6 +1708,11 @@ void ssm_update_kernel( ) { using namespace mlx::core; + if (hidden_states.inner.ndim() != 4 || B.inner.ndim() != 4) { + throw std::invalid_argument( + "[ssm_update_kernel] hidden_states and B must be rank 4 " + "([batch, 1, heads, head_dim] and [batch, 1, groups, state])"); + } auto shape = hidden_states.inner.shape(); int n = shape[0]; // batch int h = shape[2]; // num_heads @@ -1715,6 +1720,25 @@ void ssm_update_kernel( auto b_shape = B.inner.shape(); int hb = b_shape[2]; // n_groups int ds = b_shape[3]; // state_dim + // The kernel trusts these shapes as template constants and indexes raw + // buffers with them, so a checkpoint config the graph path would reject + // inside reshape/repeat (a state width that is not a multiple of 32, heads + // that do not divide into groups, a mismatched state) would read or leave + // unwritten GPU memory instead. Refuse it here; the bridge returns the + // throw to Rust as an error. + if (C.inner.shape() != B.inner.shape() || shape[1] != 1 || + b_shape[0] != n || b_shape[1] != 1 || hb <= 0 || h % hb != 0 || + ds < 32 || ds % 32 != 0 || + A_log.inner.size() != static_cast(h) || + D.inner.size() != static_cast(h) || + dt.inner.size() != static_cast(n) * h || + state_in.inner.size() != static_cast(n) * h * dh * ds) { + throw std::invalid_argument( + "[ssm_update_kernel] unsupported shapes for the fused SSM step " + "(single token, state width a multiple of 32, heads divisible by " + "groups, state [batch, heads, head_dim, state]); use the ssm_step " + "graph path"); + } int g = h / hb; // heads per group auto input_type = hidden_states.inner.dtype(); diff --git a/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs b/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs index 7bf5dc318..8ccdde235 100644 --- a/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs +++ b/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs @@ -358,6 +358,56 @@ fn ssm_update_kernel_keys_on_a_log_dtype() { check(&GRANITE_TINY, 2072, dtype::BFLOAT16, dtype::FLOAT32); } +/// Shapes the kernel cannot index safely are refused on the host, as an error +/// the caller sees, rather than launched: a state width that is not a multiple +/// of 32 would leave state columns unwritten, and heads that do not divide +/// into groups would read past B and C. +#[test] +fn ssm_update_kernel_refuses_unsupported_shapes() { + let _guard = crate::test_support::env_lock::env_lock(); + if skip() { + return; + } + let bad = [ + Shape { + name: "state 48", + state_dim: 48, + ..PADDED_ROWS + }, + Shape { + name: "4 heads in 3 groups", + groups: 3, + ..PADDED_ROWS + }, + ]; + for s in &bad { + let case = make_case(s, 2074, dtype::FLOAT32, dtype::FLOAT32); + let mut out = UniquePtr::null(); + let mut state = UniquePtr::null(); + let result = ssm_update_kernel( + &case.x, + &case.a_log, + &case.b, + &case.c, + &case.d, + &case.dt, + &case.dt_bias, + &case.state, + s.dt_limits.0, + s.dt_limits.1, + &mut out, + &mut state, + ); + let err = result.expect_err(s.name); + assert!( + err.what().contains("unsupported shapes"), + "{}: unexpected error {}", + s.name, + err.what() + ); + } +} + /// `MLXCEL_SSM_KERNEL=0` and its older alias `MLXCEL_SSM_CUDA_KERNEL=0` each /// turn the predicate off, on every backend with a port. #[test] @@ -367,6 +417,9 @@ fn ssm_kernel_kill_switches_turn_predicate_off() { .iter() .map(|name| (*name, std::env::var(name).ok())) .collect(); + // SAFETY (every env mutation in this test): the crate-wide env_lock guard + // is held for the whole window, and every other test in this crate that + // mutates the environment takes the same lock. for name in KILL_SWITCHES { unsafe { std::env::remove_var(name) }; } From 2f17f286e7b73c2ffe777f2aaeb807f6ebd7f781 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 4 Oct 2026 22:15:23 +0900 Subject: [PATCH 6/7] test(kernels): cover more SSM shape refusals and document the tests ssm_update_kernel_refuses_unsupported_shapes now also checks a 16-wide state (under 32, where the kernel wrote no state at all) and a D one element longer than the head count, beside the 48-wide state and the 4-heads-in-3-groups cases. The results doc names the commit its rows were measured at and what changed after it, and lists the padded-rows shape, the refusal test, the kill-switch test and the untested CPU-device rule. environment-variables.md says the predicate is also false when the default device is the CPU. Validated on gfx1151: ssm_update_parity_tests 5 passed; clippy -p mlxcel-core --lib --tests --features rocm clean. Refs #2067 --- ...cm-ssm-update-kernel-gfx1151-2026-10-04.md | 4 +++- docs/environment-variables.md | 2 +- .../src/ssm_update_parity_tests.rs | 20 ++++++++++++++++--- 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md b/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md index 1381a3e39..5e7e5b5ef 100644 --- a/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md +++ b/docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md @@ -6,7 +6,7 @@ The port (`SSM_HIP_SOURCE` in `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`) is ## Environment -AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5), 96 GiB VRAM carve-out, Debian 13, ROCm 10.0.0 (HIP 7.15.26333). mlxcel `844bd94c` (the PR branch on `origin/main` `c0b71344`), MLX pin `81ba1c6a`, ROCm overlay `75915908`, `cargo build --release --features rocm`. Checkpoints: `mlx-community/granite-4.0-h-tiny-4bit` and `mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit`, the same local directories as the earlier correctness rows. +AMD Ryzen AI MAX+ 395 with Radeon 8060S (`gfx1151`, RDNA 3.5), 96 GiB VRAM carve-out, Debian 13, ROCm 10.0.0 (HIP 7.15.26333). mlxcel `844bd94c` (the PR branch on `origin/main` `c0b71344`) for the decode and logit rows; the later commits of the PR add only the CPU-device check in `ssm_kernel_available()` and host shape checks in front of the same launch, and the greedy 128-token output of both models was re-checked identical after them, MLX pin `81ba1c6a`, ROCm overlay `75915908`, `cargo build --release --features rocm`. Checkpoints: `mlx-community/granite-4.0-h-tiny-4bit` and `mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-4bit`, the same local directories as the earlier correctness rows. ## Decode throughput @@ -27,6 +27,8 @@ Run-to-run spread is under 1% for granite and under 3% for Nemotron-H in either A third test runs the same shape with `A_log` first in bf16 and then in f32 in one process. Nemotron-H stores `A_log` in f32 next to bf16 activations and granite in bf16, and the generated kernel signature takes each input's runtime dtype, so the launch now names the `A_log`, `B` and `C` dtypes in its template arguments (the cache key on CUDA and ROCm). Without them the second launch reuses the bf16 module and the test fails at normalized RMS 0.80. +The f32 test also runs a head size of 60, which leaves rows past `Dh` in the last threadgroup of 8. `ssm_update_kernel_refuses_unsupported_shapes` pins that shapes the kernel cannot index (a state width that is not a multiple of 32, heads that do not divide into groups; the host also checks rank, a single token, the batch and every input's element count) throw `std::invalid_argument` before any launch, and a fifth test that `MLXCEL_SSM_KERNEL=0` and `MLXCEL_SSM_CUDA_KERNEL=0` each turn the predicate off. The CPU-device rule is not covered by a test: switching the process's default device inside the shared test binary would move other tests to the CPU. + ### Model logits: kernel against kernel The `w1ctx512` traces (`1 128 8 512` with `MLXCEL_TRACE_START_TOKEN=512`) prefill 512 corpus tokens into a fresh cache and then trace one token, so every traced step meets an SSM state and takes the fused step. Until now ROCm ran the graph there and the rows compared Metal's kernel with ROCm's graph ([rocm-correctness-gfx1151-2026-09-30.md](rocm-correctness-gfx1151-2026-09-30.md#the-fused-ssm-kernel-against-the-graph-with-state)). With the port they compare kernel with kernel. New traces: `benchmarks/logit_traces/rocm_gfx1151_96cbce84/`, built at the PR's first commit (the second changes only the kernel's cache name, the same on ROCm in both). `ssmkernel0` is the same binary with `MLXCEL_SSM_KERNEL=0`. `python3 scripts/compare_logit_traces.py --decided 2.0`, Nemotron-H files filtered with `grep -v '^\[NemotronH\] '`. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 977087194..368039ac9 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -550,7 +550,7 @@ recommended as normal deployment settings. | `MLXCEL_QMV_MULTIROW` | `0` disables; any other value or unset enables | on | **CUDA only.** Gates the weight-amortizing multirow qmv path (issue #725). With broadcast weights and `2 <= M*B <= 8` input rows (batched decode with B in [2,8), speculative-verify `[1, K]` forwards), one warp applies each dequantized weight tile to every input row instead of launching one weight-rereading block column per row, so weight DRAM traffic is O(1) in the row count. Per-row outputs are bit-identical to the stock per-row launches (pinned by `qmv_multirow_matches_per_row_qmv_bitwise`); classic `M*B == 1` decode is untouched. On GB10 this flips B=1 MTP speculative decoding from 0.52-0.77x to 1.31-1.46x and lifts 4-client aggregate serving decode from ~50 to ~74 tok/s on llama-3.1-8b-4bit. Set to `0` to force the stock per-row dispatch (A/B, rollback). See `docs/benchmark_results/qmv-multirow-gb10-2026-07-11.md`. | | `MLXCEL_QMV_ROWS_PER_WARP` | `1`, `2` or `4`; any other value (or unset) uses the per-architecture default | `2` on sm_70, `1` elsewhere | **CUDA only.** Output rows each warp covers in `qmv`, for classic `M*B == 1` decode. The stock kernel gives one warp one output row and has that warp read the whole activation vector, so a launch moves `n * k * 2` bytes of activations against `n * k / 2` bytes of 4-bit weights. Measured on a V100 over a controlled 4-bit against 8-bit pair of gemma-4-12B-it with identical launch counts, that activation term is 84% of a 4-bit launch, which is why decode below Ampere is not weight-bandwidth bound. `R` rows per warp load the activation tile once and apply it to all `R` weight rows, dividing the activation traffic by `R` while leaving weight traffic alone; the per-row arithmetic is unchanged, so output is byte-identical to `R = 1` at temperature 0. The cost is that the same output rows are covered by `R` times fewer warps and the accumulator array grows to `R * elems_per_thread`. On a V100 at 4 bits, R = 2 moves decode from 44.7 to 60.1 tok/s on gemma-4-12B-it-4bit under the `MLXCEL_QMV_MIN_BLOCKS` default of 3, and from 20.6 to 24.9 on qwen3.8-27B-4bit; R = 4 gives it all back on the 12B (96 to 142 registers, 1 to 2 resident blocks per SM against the stock kernel's 4 to 5) while costing nothing on the wider 27B. Turing defaults to 1 because no sm_75 device was available to measure and the R = 4 reversal shows the tradeoff does not extrapolate. `1` routes to the stock kernel with the stock grid, so it is the rollback. Read once per process. | | `MLXCEL_QMV_MIN_BLOCKS` | `1`, `3` or `4`; any other value (or unset) uses the per-architecture default | `3` on sm_70, `1` elsewhere | **CUDA only.** Resident-block floor handed to `__launch_bounds__` on the row-blocked `qmv` kernel, which is how ptxas is told to cap the register count: for a 256-thread block, 3 caps at 85 registers per thread and 4 caps at 64. This is the knob, not the row count. Measured on a V100 at 4 bits on gemma-4-12B-it, decode as the slope between `-n 60` and `-n 200`: R = 2 unconstrained sits at 92 to 95 registers and 2 blocks per SM for 56.0 tok/s, the same R = 2 under a floor of 3 sits at 80 registers and 3 blocks for 60.1, and under a floor of 4 it sits at 64 registers and falls back to 54.9, so the budget can be set too tight as well as too loose. Nothing spills at any floor (`LOCAL` is 0 in every instantiation). R = 3 under the same floor of 3 has identical registers and residency and moves 17% fewer memory transactions per output row, and it still lands at 58.8, which is why the floor rather than the row count carries the default. `1` is the unconstrained case and is the rollback. Read once per process. | -| `MLXCEL_SSM_KERNEL` | `0` disables; any other value or unset enables | on | Gates the fused single-token SSM decode kernel (`ssm_update_kernel`) on every backend that has a port: Metal, CUDA (issue #631) and ROCm (issue #2067). It replaces the ~55-op SSD scan graph (`ssm_step`) with one launch per SSM layer per decode step for the hybrid SSM models (granite-4.0-h, falcon-h1, plamo-2, nemotron-h). `ssm_kernel_available()` answers from the kernel's port table (`ssm_ports()`), so a backend with no port takes the graph path on its own. Set to `0` to force the graph path (A/B, rollback). Measured on GB10 (CUDA): 2.6-4.5x decode on the granite/falcon family, see `docs/benchmark_results/hybrid-ssm-decode-cuda-kernel-gb10-2026-07-10.md`; on gfx1151 (ROCm): see `docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md`. | +| `MLXCEL_SSM_KERNEL` | `0` disables; any other value or unset enables | on | Gates the fused single-token SSM decode kernel (`ssm_update_kernel`) on every backend that has a port: Metal, CUDA (issue #631) and ROCm (issue #2067). It replaces the ~55-op SSD scan graph (`ssm_step`) with one launch per SSM layer per decode step for the hybrid SSM models (granite-4.0-h, falcon-h1, plamo-2, nemotron-h). `ssm_kernel_available()` answers from the kernel's port table (`ssm_ports()`), so a backend with no port takes the graph path on its own. It is also false when the default device is the CPU (`MLXCEL_DEVICE=cpu`). Set to `0` to force the graph path (A/B, rollback). Measured on GB10 (CUDA): 2.6-4.5x decode on the granite/falcon family, see `docs/benchmark_results/hybrid-ssm-decode-cuda-kernel-gb10-2026-07-10.md`; on gfx1151 (ROCm): see `docs/benchmark_results/rocm-ssm-update-kernel-gfx1151-2026-10-04.md`. | | `MLXCEL_SSM_CUDA_KERNEL` | `0` disables; any other value or unset enables | on | Older name of `MLXCEL_SSM_KERNEL`, from when the CUDA port was the only one off Apple, kept as an alias so existing scripts still work. Either variable set to `0` forces the graph path, now on every backend (Metal and ROCm included). | | `MLXCEL_FUSED_QK_NORM` | `1`/`true`/`on`/`yes` enable; any other value or unset disables | off (opt-in) | Fused single-token QKV projection + Q/K RMSNorm + RoPE kernel (#326) for Qwen3 and Qwen3-MoE decode. Opt-in: set to `1` to enable. Matches the graph path within RMS < 5e-3 (the reduction is over the transpose-invariant head_dim axis), but greedy temp-0 is not byte-identical over long generation; on CUDA the graph path is itself non-deterministic run-to-run from GPU FP-reduction order, while the fused path is deterministic, so its output stays inside the graph baseline's own envelope. The kernel cuts Rust/C++ FFI crossings rather than MLX op count, so it does not speed up the GPU/bandwidth-bound decode loop: on M1 Ultra it measured 1 to 3.4% slower (qwen3-0.6b 275 vs 284, qwen3-8b 82.3 vs 83.2 tok/s); on GB10/CUDA (SM 12.1) it is also slower (qwen3-0.6b 0.96x, qwen3-8b ~1.0x, qwen3-30b-a3b 0.92x fused/graph; see `docs/benchmark_results/fused-qk-norm-decode-gb10.md`), so there is no per-backend win. Ships as a reusable shared primitive for the deferred QK-norm families and stays opt-in (default off) on every measured backend (M1 Ultra, M5 Max, GB10/CUDA), mirroring the opt-in `MLXCEL_FUSED_MOE_RELU2`. Active only when `l == 1` (decode) and weights are quantized. | | `MLXCEL_FUSED_ADD_RMSNORM` | `0`/`false`/`off`/`no` disable; `1`/`true`/`on`/`yes` enable; unset or unrecognised keeps the default | off (opt-in) | Fused residual-add + RMSNorm decode kernel (#905). Computes `new_residual = x + residual` and `normed = rms_norm(new_residual) * (weight_bias + weight)` in one dispatch instead of an elementwise `Add` followed by `fast::rms_norm`, removing one kernel launch and one full-width intermediate at every pre-norm residual join. Adopted by the Llama3-family block (which also serves Qwen2/Qwen2.5 and every VLM whose text backbone is `Llama3Model`/`Qwen2Model`) and by Gemma, where the `(1 + w)` convention rides in as a scalar `weight_bias = 1.0` folded in the weight's own dtype rather than as a second weight tensor. Numerics are pinned to MLX's own `rms_norm` kernel (fp32 accumulation over the dtype-rounded sum, `precise::rsqrt`, the `x * inv_mean` product rounded to the activation dtype before the weight multiply), so the fused and unfused paths agree to within a small multiple of one ulp; greedy temp-0 is not guaranteed byte-identical over long generation, because a near-tie argmax can land on either side of that difference. Set to `0` to force the `add` + `fast_rms_norm` pair (A/B, rollback). The compiled-in default lives in `FUSED_ADD_RMSNORM_DEFAULT` in `src/lib/mlxcel-core/src/layers.rs`. Read once at first use and cached for the process lifetime. | Ships **opt-in (default off)**: the op-level microbench on M1 Ultra (Metal, f16, hidden {2048, 4096, 8192} x batch {1, 4, 8}, three repetitions) measured the fused path at parity, speedups scattering 0.77x to 1.14x with no structure and the extremes not reproducing. Per-iteration time stayed near 280-410us across a 4x change in hidden size, so the harness is dominated by dispatch cost rather than the op. Issue #905's measure-then-keep policy lands a fusion unwired absent a demonstrated win, matching how `MLXCEL_FUSED_QK_NORM` (#326) shipped. The production decode path agrees: on M1 Ultra at pp512/tg128, Llama 3.1 8B decodes at 113.18 tok/s off and 113.01 on, Qwen2.5 7B at 115.38 and 115.34, with prefill unchanged. Set to `1` to opt in. See `docs/benchmark_results/fused-norm-rope-m1ultra-2026-07-31.md` and `docs/benchmark_results/fused-add-rmsnorm-decode-m1ultra-2026-09-27.md`. diff --git a/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs b/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs index 8ccdde235..1bcf25ca6 100644 --- a/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs +++ b/src/lib/mlxcel-core/src/ssm_update_parity_tests.rs @@ -360,8 +360,9 @@ fn ssm_update_kernel_keys_on_a_log_dtype() { /// Shapes the kernel cannot index safely are refused on the host, as an error /// the caller sees, rather than launched: a state width that is not a multiple -/// of 32 would leave state columns unwritten, and heads that do not divide -/// into groups would read past B and C. +/// of 32 (or under 32) would leave state columns unwritten, heads that do not +/// divide into groups would read past B and C, and a per-head parameter of the +/// wrong length would be read past its end. #[test] fn ssm_update_kernel_refuses_unsupported_shapes() { let _guard = crate::test_support::env_lock::env_lock(); @@ -374,14 +375,27 @@ fn ssm_update_kernel_refuses_unsupported_shapes() { state_dim: 48, ..PADDED_ROWS }, + Shape { + name: "state 16", + state_dim: 16, + ..PADDED_ROWS + }, Shape { name: "4 heads in 3 groups", groups: 3, ..PADDED_ROWS }, + // Valid shape; `D` is given one element too many below. + Shape { + name: "D of length heads + 1", + ..PADDED_ROWS + }, ]; for s in &bad { - let case = make_case(s, 2074, dtype::FLOAT32, dtype::FLOAT32); + let mut case = make_case(s, 2074, dtype::FLOAT32, dtype::FLOAT32); + if s.name.starts_with("D of length") { + case.d = uniform(0.5, 1.5, &[s.heads + 1]); + } let mut out = UniquePtr::null(); let mut state = UniquePtr::null(); let result = ssm_update_kernel( From c70609499e8bc6819c601a5cabf28a36d736a02b Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Sun, 4 Oct 2026 22:32:06 +0900 Subject: [PATCH 7/7] docs: add technical report for PR #2099 --- ...99-rocm-ssm-update-hip-port-20261004.en.md | 213 ++++++++++++++++++ ...99-rocm-ssm-update-hip-port-20261004.ko.md | 213 ++++++++++++++++++ 2 files changed, 426 insertions(+) create mode 100644 TECHNICAL_REPORTS/2099-rocm-ssm-update-hip-port-20261004.en.md create mode 100644 TECHNICAL_REPORTS/2099-rocm-ssm-update-hip-port-20261004.ko.md diff --git a/TECHNICAL_REPORTS/2099-rocm-ssm-update-hip-port-20261004.en.md b/TECHNICAL_REPORTS/2099-rocm-ssm-update-hip-port-20261004.en.md new file mode 100644 index 000000000..bcc6bba31 --- /dev/null +++ b/TECHNICAL_REPORTS/2099-rocm-ssm-update-hip-port-20261004.en.md @@ -0,0 +1,213 @@ +# Technical Report: PR #2099 - Port ssm_update_kernel to HIP and read the port table + +**Date**: 2026-10-04 + +**Status**: Implemented and measured on the gfx1151 host; head `2f17f286` on origin/main `c0b71344`, pending merge. + +**Languages**: C++ (HIP kernel source, kernel holder, support predicate, host shape checks), Rust (FFI doc comments, parity tests), Markdown (environment variables, installation, benchmark results), CSV/TSV (bench rows, logit traces) + +**Risk Level**: Medium (it changes the default single-token decode path of the hybrid SSM models on ROCm, and touches the predicate, the template arguments and the host checks that Metal and CUDA share, neither of which could be run on the development host) + +## Executive Summary + +Issue #2067 (part of #1814, epic #1801) is the first port in the order that PR #2086 measured: the single-token Mamba2 SSM update, which the decode profile put at 29.8% of granite-4.0-h-tiny's and 20.0% of Nemotron-H's decode GPU time on gfx1151. Before this PR every such step on ROCm ran the ~55-op SSD graph (`ssm_step`). `ssm_ports()` had no `.rocm` entry, and `ssm_kernel_available()` did not read the table at all, so a filled slot would have stayed unreachable. + +The PR does four things in `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`. It adds `SSM_HIP_SOURCE`, the CUDA kernel body with the lane fold rewritten as `__shfl_down(acc, o, 32)`, and fills `.rocm`. It makes `ssm_kernel_available()` return `has_kernel_port(ssm_ports())` on every platform, with a new `MLXCEL_SSM_KERNEL=0` switch and the old `MLXCEL_SSM_CUDA_KERNEL=0` kept as an alias. It adds the `A_log`, `B` and `C` dtypes to the launch's template arguments on every backend, because Nemotron-H stores `A_log` in f32 beside bf16 activations and the JIT caches key on template arguments. It refuses on the host, as an error, the shapes the kernel cannot index. + +On gfx1151, with the same binary and `MLXCEL_SSM_KERNEL=0` as the "before" arm, decode went from 60.39 to 88.43 tok/s on granite-4.0-h-tiny-4bit (1.46x) and from 51.35 to 74.37 tok/s on NVIDIA-Nemotron-3-Nano-30B-A3B-4bit (1.45x), medians of three runs. The `w1ctx512` logit traces, where every traced step meets a 512-token SSM state, now compare ROCm's kernel with Metal's kernel on M5 and find 0 of 60 (granite) and 0 of 71 (Nemotron-H) decided-position mismatches. Greedy 128-token output with and without the kernel is coherent on both models but not byte-identical; the traces put every flip at an undecided position. + +One finding goes beyond this issue: the wave32 `#error` guard that every #1814 port is required to carry is inert with ROCm 10's AMD clang 23, which defines neither `__AMDGCN_WAVEFRONT_SIZE__` nor `__AMDGCN_WAVEFRONT_SIZE` for gfx1151, gfx942 or gfx90a. This port does not depend on it, but the idiom guards nothing for the ports that follow. + +## 1. Problem Statement + +### 1.1 An empty slot behind a predicate that never asked + +Hybrid SSM models (granite-4.0-h, falcon-h1, plamo-2, Nemotron-H) gate their fused single-token path with `seq_len == 1 && ssm_kernel_available()` in `granitemoehybrid.rs`, `falcon_h1.rs`, `plamo2.rs` and `nemotron_h.rs`, and then `expect` the launch. Before this PR the predicate was: + +```cpp +#ifdef __APPLE__ + return mlx::core::metal::is_available(); +#else + // MLXCEL_SSM_CUDA_KERNEL=0 forces the graph path + ... + return mlx::core::cu::is_available(); +#endif +``` + +On ROCm that answered false, so the gates took the graph. Filling `.rocm` alone would have changed nothing, and filling it while the predicate stayed backend-enumerated is the shape `verify-kernel-port-dispatch` exists to prevent: a predicate and a dispatch that can disagree. + +### 1.2 What the profile said it was worth + +PR #2086 measured the SSM update as the largest recoverable share of any #1814 port on gfx1151: 29.8% of granite-4.0-h-tiny's and 20.0% of Nemotron-H's decode GPU time, and more than half of their dispatches. The CUDA port (#631) had measured 2.6-4.5x decode on the granite/falcon family on GB10. + +### 1.3 No kernel-level parity test existed + +Parity for the CUDA port was end-to-end greedy output on CUDA. Nothing compared the kernel with the graph step on the same inputs, so a reduction bug that returns finite, plausible values (the failure the wave32 rule in #1814 describes) would only show up as drifting model output. + +## 2. Change Summary + +| Area | Change | +|---|---| +| `mlx_cxx_kernels.cpp`, kernel | `SSM_HIP_SOURCE`, `SsmKernelHolderHip` / `get_ssm_kernel_hip()` calling `fast::hip_kernel` under `MLXCEL_BRIDGE_ROCM_BACKEND`; `.rocm` getter in `ssm_ports()` | +| `mlx_cxx_kernels.cpp`, predicate | `ssm_kernel_available()`: kill switches, then false off the GPU device, then `has_kernel_port(ssm_ports())` | +| `mlx_cxx_kernels.cpp`, launch | `TA`, `TB`, `TC` template arguments (dtypes of `A_log`, `B`, `C`) on every backend | +| `mlx_cxx_kernels.cpp`, host | Rank and shape checks in `ssm_update_kernel` that throw `std::invalid_argument` before any launch | +| `ssm_update_parity_tests.rs` (new) | Five tests: f32 and bf16 parity at three shapes, the `A_log` dtype key, shape refusals, both kill switches | +| `mlx_cxx_bridge.h`, `lib.rs` | Doc comments for the predicate and the kernel; test module registration | +| Docs | `MLXCEL_SSM_KERNEL` row and alias row in `docs/environment-variables.md`; ROCm row in `docs/installation.md`; the "Open" item in `rocm-correctness-gfx1151-2026-09-30.md` now points to the kernel-against-kernel rows; new `rocm-ssm-update-kernel-gfx1151-2026-10-04.md` | +| Data | `benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-{off,on}.csv`; six traces plus `METADATA.txt`, `RUNS.txt`, `SHA256SUMS` and `README.md` in `benchmarks/logit_traces/rocm_gfx1151_96cbce84/` | + +The branch has six commits: the port and predicate (`96cbce84`), the dtype keys on every backend (`844bd94c`), the docs (`3780b0b2`), the CPU-device check and review fixes (`bd2e8f31`), the host shape refusal (`38071038`), and more refusal tests plus test docs (`2f17f286`). + +## 3. The Port + +### 3.1 CUDA body, HIP fold + +The HIP source keeps the CUDA port's kernel name pattern, inputs (`X`, `A_log`, `B`, `C`, `D`, `dt`, `state_in`), outputs (`out`, `state_out`), grid and template arguments. Each (head, row) pair is reduced by the 32 lanes that share a `threadIdx.y` in a (32, 8, 1) threadgroup; each lane walks `Ds / 32` state elements, updates the state, and accumulates `state * C`. The only body change is the fold: + +```cpp +for (int o = 16; o > 0; o >>= 1) { + acc += __shfl_down(acc, o, 32); +} +``` + +`__shfl_down_sync` exists in HIP only as a compatibility shim that ignores its mask, so the native `__shfl_down` is used and the width is stated. The width is what makes the fold correct regardless of wave size: on a wave32 target (gfx11, gfx12) the 32 lanes are the whole wave, and on a wave64 target (CDNA) a width of 32 splits the wave into two 32-lane segments, which are again exactly two rows. The early return for rows past `Dh` is wave-uniform because all 32 lanes of a wave share `threadIdx.y`. + +### 3.2 The wave32 guard, and why it does nothing here + +The #1814 port requirements tell every port that reduces across lanes to carry two `#error` checks, on `__AMDGCN_WAVEFRONT_SIZE__` and on `__AMDGCN_WAVEFRONT_SIZE`, as the bitlinear port does. The port carries them. During the work it turned out that ROCm 10's AMD clang 23 defines neither macro for gfx1151, gfx942 or gfx90a, so the `#if defined(...)` conditions are false on every target and the guard cannot fire. The comment in `SSM_HIP_SOURCE` and the benchmark doc both say so. The fold above is correct on wave64 by its width argument, not because of the guard; see section 9 for what that means for the remaining ports. + +## 4. The Predicate and the Switch + +`ssm_kernel_available()` now reads, in order: + +1. `MLXCEL_SSM_KERNEL=0` or `MLXCEL_SSM_CUDA_KERNEL=0` returns false. The new name says what the switch does on every backend; the old one is kept so existing A/B scripts still work. +2. A default device other than the GPU returns false, the same check `mamba1_scan_kernel_accepts` makes. Custom kernels run only on the GPU stream, so with `MLXCEL_DEVICE=cpu` the gates take `ssm_step` on the CPU instead of the launch throwing. +3. Otherwise `mlxcel::has_kernel_port(ssm_ports())`, which reads the same table `select_kernel_port` reads in the dispatch. A true answer therefore means the dispatch will not refuse, which is what lets the Rust gates `expect` the launch. + +Two behavior changes follow for the other backends, and the docs state both. The alias now also affects Metal and ROCm (it used to be read on the non-Apple path only), and `MLXCEL_SSM_KERNEL=0` is honored on Metal, where previously nothing turned the kernel off. With the environment unset, `has_kernel_port` is true on Metal and CUDA as before. + +No model file changed. On Nemotron-H the gate in `NemotronHMamba2Mixer::forward` selects `fused_mamba2_forward`, which runs the whole single-token mixer (input projection, convolution, this kernel, gated norm, output projection) as one C++ call. Nemotron-H's decode figures therefore measure that path against the Rust graph mixer, not the SSM step alone. + +## 5. Dtype Keys on Every Backend + +The CUDA and HIP JIT caches key a module on the kernel name plus `template_args`, while the generated kernel signature takes each input's runtime dtype. The launch named `T` (activations, also covering `X` and `D`) and `U` (state) only. Nemotron-H stores `A_log` in f32 next to bf16 activations, granite in bf16, so wherever both reach the same cached module (one process, as in the parity test) whichever compiled it first fixes the pointer type the other reads `A_log` through. `B` and `C` are not tied to `T` either. + +The fix adds `TA`, `TB`, `TC` for the `A_log`, `B` and `C` dtypes. The kernel bodies never name them, so only module names change, not arithmetic. They are added on every backend, Metal included (where the name already keys on every input dtype, so the addition is redundant there). The first version (`96cbce84`) appended them inside `if (mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Rocm)`, to leave the Metal and CUDA module names unchanged. That is what rule 1 of `scripts/ci/check_kernel_port_dispatch.py` forbids: a file that launches a custom kernel must not branch on the backend kind outside `kernel_port.cpp` and `gpu_backend.cpp`, because hand-written per-backend branches are how the launchers drifted before `select_kernel_port` existed. The second commit (`844bd94c`) moved them to every backend. The traces were produced at `96cbce84`; on ROCm the launch, its arguments and its cache key are the same in both commits, which `METADATA.txt` records. + +`ssm_update_kernel_keys_on_a_log_dtype` runs the same shape with `A_log` first in bf16 and then in f32 in one process. Without the keys the second launch reuses the bf16 module and the test fails at normalized RMS 0.80. + +## 6. Host-Side Shape Refusal + +The kernel takes `Dh`, `Ds`, `H` and `G` as template constants and indexes raw buffers with them. A checkpoint config the graph path would reject inside `reshape` or `repeat` would instead read or leave unwritten GPU memory. `ssm_update_kernel` now throws `std::invalid_argument` before any launch when: + +- `hidden_states` or `B` is not rank 4; +- `C`'s shape differs from `B`'s; +- there is more than one token, or `B`'s batch differs from the activations'; +- groups are not positive or heads are not divisible by groups; +- the state width is below 32 or not a multiple of 32; +- the element count of `A_log`, `D`, `dt` or `state_in` does not match the shape. + +The bridge returns the throw to Rust as an error. The checks sit in front of the shared launch, so they apply to every port; the HIP port is what made them reachable on ROCm. `ssm_update_kernel_refuses_unsupported_shapes` pins them. + +## 7. Measured Results + +### 7.1 Decode throughput + +`scripts/bench_decode.sh` at pp512/tg128, one model per run, three runs per arm with the arm order alternated, at `844bd94c`, same binary. + +| Model | Graph (`MLXCEL_SSM_KERNEL=0`) tok/s | Kernel tok/s | Speedup (medians) | +|---|---|---|---| +| granite-4.0-h-tiny-4bit | 60.57 / 60.39 / 60.05, median 60.39 | 88.36 / 88.69 / 88.43, median 88.43 | 1.46x | +| NVIDIA-Nemotron-3-Nano-30B-A3B-4bit | 51.65 / 50.87 / 51.35, median 51.35 | 74.87 / 72.87 / 74.37, median 74.37 | 1.45x | + +Run-to-run spread is under 1% for granite and under 3% for Nemotron-H in either arm. `compare_bench_csv.py --before --after `, which keeps the last row per model, reports 1.47x and 1.45x. Prefill is not on this path and is not compared; it varied between 379 and 608 tok/s for granite across both arms on the shared host. + +The gain is smaller than the CUDA port's 2.6-4.5x on GB10. It is also larger than the decode profile's GPU-time shares alone would suggest: removing a 29.8% or 20.0% share entirely is about 1.42x and 1.25x of GPU time. The profile was taken at a different commit under a profiler that itself costs 20 to 23% on these models, and Nemotron-H's fused path replaces the whole mixer rather than the SSM step, so this is a consistency check on the published shares, not a decomposition of the speedup; no measurement in this PR attributes the excess (host-side dispatch of the ~55-op graph is the obvious candidate). + +### 7.2 Idle-GPU guard + +Every run went through `scripts/rocm_gpu_guard.sh` (90 s with `/sys/class/kfd/kfd/proc` empty and no compiler running, then 1 Hz monitoring). A parallel development unit shared the GPU; every attempt that overlapped its work was rejected and rerun. All runs but one used the copy of the guard fixed in #2098 (commit `c3eab7b2`), run from the #2065 unit's worktree and not committed in this PR; that fix stops the guard from rejecting the command's own exited children. The exception is the first run (granite, graph arm, r1), which used the copy on `main`. The old guard's defect produced false rejections, never false acceptances, so that run's acceptance holds. + +### 7.3 Kernel against the graph, one step + +`ssm_update_parity_tests` compares output and new state with a float32 MLX-op reference of the one-token SSD step: + +| Shape | Batch | Heads x head dim | Groups | State | +|---|---|---|---|---| +| granite-4.0-h-tiny | 1 | 48 x 64 | 1 | 128 | +| Nemotron-H (dt clipped to `(1e-3, 0.1)`) | 2 | 64 x 64 | 8 | 128 | +| padded rows | 1 | 4 x 60 | 2 | 64 | + +Tolerances (normalized RMS / max): f32 1e-5 / 1e-4, bf16 1.6e-2 / 7e-2, from the issue. The state is held to the f32 budget in every case because the models carry it in f32. The tests pass on gfx1151. Two mutations show the tests can tell: a lane fold starting at 8 instead of 16 fails at normalized RMS 0.65 to 0.89, and dropping the dtype keys fails at 0.80. The tests return early only when there is no GPU backend, the default device is the CPU, or a kill switch is set; on Metal, CUDA and ROCm a false predicate fails the test instead of skipping it. + +### 7.4 Model logits: kernel against kernel (`w1ctx512`) + +The `w1ctx512` traces prefill 512 corpus tokens and then trace single-token steps, so every traced step meets a state and takes the fused step. Until this PR the corresponding rows in `rocm-correctness-gfx1151-2026-09-30.md` compared Metal's kernel with ROCm's graph. Now the comparison is kernel with kernel. `compare_logit_traces.py --decided 2.0`, Nemotron-H files filtered of their `[NemotronH]` loader lines. + +| Model | Reference | Candidate | Top-1 disagreement | Decided mismatches | Largest gap | Perplexity ref / cand | +|---|---|---|---|---|---|---| +| granite-4.0-h-tiny | Metal kernel (`metal_m5_3c9edea0`) | ROCm kernel | 3 / 128 | 0 / 60 | 0.125 | 7.090 / 7.216 | +| granite-4.0-h-tiny | Metal kernel | ROCm graph (`ssmkernel0`) | 2 / 128 | 0 / 60 | 0.125 | 7.090 / 7.290 | +| granite-4.0-h-tiny | ROCm graph (`ssmkernel0`) | ROCm kernel | 1 / 128 | 0 / 59 | 0.000 | 7.290 / 7.216 | +| granite-4.0-h-tiny | ROCm graph at `3c9edea0` | ROCm kernel | 1 / 128 | 0 / 60 | 0.125 | 7.234 / 7.216 | +| nemotron-3-nano-30b-a3b | Metal kernel (`metal_m5_3c9edea0`) | ROCm kernel | 3 / 128 | 0 / 71 | 0.250 | 7.489 / 7.442 | +| nemotron-3-nano-30b-a3b | Metal kernel | ROCm graph (`ssmkernel0`) | 4 / 128 | 0 / 71 | 0.250 | 7.489 / 7.460 | +| nemotron-3-nano-30b-a3b | ROCm graph (`ssmkernel0`) | ROCm kernel | 1 / 128 | 0 / 68 | 0.125 | 7.460 / 7.442 | +| nemotron-3-nano-30b-a3b | ROCm graph at `3c9edea0` | ROCm kernel | 4 / 128 | 0 / 68 | 0.125 | 7.489 / 7.442 | + +Every pairing has zero decided-position mismatches, and every top-1 disagreement puts the reference's token at rank 2 or 3 of the candidate at a reference gap of 0.25 or less. The `default` and `ssmkernel0` traces from the same binary differ (different perplexities, one flipped token each), which shows the `default` traces actually took the kernel. Against Metal's kernel the ROCm kernel is as close as the ROCm graph: 3 and 3 top-1 disagreements, against 2 and 4 for the graph here and 4 and 7 in the earlier `3c9edea0` comparison. + +The `w1` traces (no prefill) never meet a state and run the graph with or without the kernel. Against `rocm_gfx1151_c5fe9a16` both models agree at every position (granite 0 / 128 top-1, 0 / 116 decided; Nemotron-H 0 / 128 top-1, no decided position). These satisfy the issue's `w1` criterion and show the stateless path is unchanged; they say nothing about the kernel itself, which is why the `w1ctx512` rows carry the correctness argument. + +### 7.5 Greedy divergence + +`mlxcel generate -p "The history of the Roman Empire" -n 128 -t 0 --no-chat-template`, with and without `MLXCEL_SSM_KERNEL=0`. Both arms of both models produce coherent text, but not byte-identical text, unlike the CUDA port on GB10. Granite parts at the third generated token ("a rich and complex tapestry" against "a fascinating and complex subject"); Nemotron-H after about a dozen tokens. A free-running generation conditions everything after the first flip on different text, so it cannot say whether the kernel is wrong or the model was undecided. The teacher-forced traces answer that: the kernel and the graph disagree on one top-1 token per model, and never at a decided position. The divergence is a change of summation order meeting a near-tie, not an error. + +## 8. Technical Decisions + +- **Port the CUDA body, not the Metal one.** CUDA and HIP share the launch shape, grid and `template_args`, and `fast::hip_kernel` compiles through hipRTC like `cuda_kernel`. Keeping the bodies line-for-line except the fold leaves one place to look when they disagree. +- **State the shuffle width instead of relying on the guard.** The width-32 fold is correct on wave32 and wave64 by construction. That turned out to matter, because the guard the requirements prescribe does not fire on any ROCm 10 target. +- **Make the predicate read the table, not the backend.** `has_kernel_port(ssm_ports())` cannot disagree with `select_kernel_port`, and a future port that fills a slot becomes reachable with no predicate edit. The `#ifdef __APPLE__` split and the `cu::is_available()` call are gone, which the issue's first acceptance criterion checks. +- **Rename the switch, keep the alias.** `MLXCEL_SSM_KERNEL` describes a switch that now acts on every backend; keeping `MLXCEL_SSM_CUDA_KERNEL` avoids breaking existing A/B scripts. The cost is that the CUDA-named variable now also turns the kernel off on Metal and ROCm, which the docs state. +- **Dtype keys on every backend, not a ROCm branch.** A backend branch at the launch was the first attempt and was rejected by `verify-kernel-port-dispatch`. Adding the keys everywhere is harmless on Metal and closes the same keying gap on CUDA in the same change. +- **Refuse bad shapes on the host, for all ports.** The graph path fails loudly on these shapes; the kernel would fail silently. Throwing `std::invalid_argument` restores the loud failure without a per-backend condition. +- **Keep the CPU-device check in the predicate.** It is the same rule the Mamba1 scan uses, and it keeps `MLXCEL_DEVICE=cpu` on the graph path rather than turning a working configuration into an error. + +## 9. Follow-up: the Wave64 Guard Idiom Is Inert on ROCm 10 Clang + +The #1814 port requirements prescribe `#if defined(__AMDGCN_WAVEFRONT_SIZE__) && __AMDGCN_WAVEFRONT_SIZE__ != 32` plus `#error`, and the same for `__AMDGCN_WAVEFRONT_SIZE`, so that a port that assumes 32 lanes fails to compile on a wave64 CDNA part instead of returning a finite, wrong result. With ROCm 10's AMD clang 23 neither macro is defined for gfx1151, gfx942 or gfx90a, so both conditions are false and the guard never fires on any of them. + +Today that has no effect on correctness: both HIP ports in `mlx_cxx_kernels.cpp` that reduce across lanes (bitlinear and this one) pass an explicit width of 32 to `__shfl_down`, which is correct on wave64. It matters for the ports still to come (#2065, #2064, #2068, #2063, #2069). A port whose correctness does depend on a 32-lane wave would compile cleanly on CDNA, and the requirement text would make a reviewer believe it was protected. The requirement needs either a mechanism that actually fires on ROCm 10 or a rule that every cross-lane operation states its width and that the requirement stop claiming a compile-time guard. This PR does not choose the replacement, and no issue for it was found when this report was written. + +## 10. Validation + +From the PR body, on gfx1151 (Radeon 8060S, ROCm 10.0.0): + +- `make verify-rocm` at `38071038`: OK; 11,856 passed, 0 failed, 378 ignored across 146 test binaries; ROCm smoke OK. +- `cargo test --release --features rocm -p mlxcel-core --lib ssm_update_parity_tests -- --test-threads=1` at `2f17f286` (test and doc changes only since the gate): 5 passed. +- `cargo clippy -p mlxcel-core --lib --tests --features rocm -- -D warnings`: clean. `make verify-fmt verify-kernel-dtype-keys verify-kernel-port-dispatch`: pass, with `EXPECTED_IN_SCOPE` and its count unchanged. +- Greedy output for both models identical before and after the CPU-device and shape checks, so the decode and trace numbers measured at `844bd94c` and `96cbce84` stand for the head. + +Orchestrator verification on head `2f17f286`, up to date with origin/main `c0b71344`: + +- `make verify-rocm` passed every step: 11,856 tests passed, 0 failed, 378 ignored; smoke OK. + +## 11. Residual Risks and What Was Not Verified + +- **Metal and CUDA were not run.** Both are touched by the predicate (still `has_kernel_port`, true on both, plus the CPU-device and kill-switch rules), the extra `TA`/`TB`/`TC` template arguments (module names only), and the host shape checks. The parity test is written to run on both and will fail, not skip, if the predicate is false there. +- **The alias widened.** Anyone who set `MLXCEL_SSM_CUDA_KERNEL=0` on a Metal or ROCm machine and forgot about it now gets the graph path there. +- **`MLXCEL_DEVICE=cpu` with a hybrid SSM model was not run.** No test covers the CPU-device rule, because switching the default device inside the shared test binary would move other tests to the CPU. +- **falcon-h1 and plamo-2 were not run on ROCm.** They take the same gate; only granite and Nemotron-H were measured. +- **One device, one session, shared GPU.** All numbers are from gfx1151; one development unit shared the GPU and overlapping attempts were rerun. Prefill varied widely on the shared host and was not compared. No wave64 hardware was available, so the width argument for CDNA is reasoned, not run. +- **The measurement guard is not in this PR.** All runs but one used `c3eab7b2` from #2098, outside this branch; the published method references a fix that lands separately. + +## 12. Learning Points + +- **A port is not reachable until its predicate reads the table.** Filling `.rocm` without changing `ssm_kernel_available()` would have produced a kernel nothing calls and a green build. +- **A JIT cache key must name every input dtype the signature takes.** A mixed-dtype checkpoint (f32 `A_log`, bf16 activations) is exactly the case where an incomplete key silently reuses the wrong module, and only a same-process test with both dtypes catches it. +- **Prove a guard can fire before relying on it.** The wave32 `#error` looked like protection and compiled everywhere, which is the same symptom as a guard that checks nothing. +- **Mutation-check the parity test.** A fold starting at 8 and a missing dtype key both fail with large deviations, which shows the tolerance is meaningful rather than loose. +- **Greedy text is the wrong instrument for a reordering change.** Free-running output diverges at the first near-tie and stays diverged; teacher-forced traces with a decided-position threshold separate a precision change from a bug. + +Refs: #2067, #1814, #1801, #2061, #2086, #631, #1862, #1870, #2059, #2065, #2098, #2064, #2068, #2063, #2069. diff --git a/TECHNICAL_REPORTS/2099-rocm-ssm-update-hip-port-20261004.ko.md b/TECHNICAL_REPORTS/2099-rocm-ssm-update-hip-port-20261004.ko.md new file mode 100644 index 000000000..f43b3fb3c --- /dev/null +++ b/TECHNICAL_REPORTS/2099-rocm-ssm-update-hip-port-20261004.ko.md @@ -0,0 +1,213 @@ +# 기술 보고서: PR #2099 - ssm_update_kernel을 HIP으로 port하고 port table을 읽도록 함 + +**날짜**: 2026-10-04 + +**상태**: gfx1151 호스트에서 구현 및 측정 완료. head `2f17f286`, origin/main `c0b71344` 위, 머지 대기 중. + +**언어**: C++ (HIP kernel source, kernel holder, support predicate, host 측 shape 검사), Rust (FFI doc comment, parity test), Markdown (환경 변수, 설치, benchmark 결과), CSV/TSV (bench 행, logit trace) + +**위험도**: 중간 (ROCm에서 hybrid SSM 모델의 기본 single-token decode 경로가 바뀝니다. Metal과 CUDA가 공유하는 predicate, template argument, host 검사도 건드리며, 두 backend는 개발 호스트에서 실행할 수 없었습니다) + +## 요약 + +Issue #2067 (#1814의 일부, epic #1801)은 PR #2086이 측정한 순서의 첫 번째 port입니다. 대상은 single-token Mamba2 SSM update이며, decode profile은 gfx1151에서 이것이 granite-4.0-h-tiny decode GPU 시간의 29.8%, Nemotron-H의 20.0%를 차지한다고 보고했습니다. 이 PR 이전에는 ROCm에서 그런 step마다 약 55개 op로 된 SSD graph (`ssm_step`)가 실행되었습니다. `ssm_ports()`에는 `.rocm` 항목이 없었고, `ssm_kernel_available()`은 table을 전혀 읽지 않았으므로 slot을 채워도 도달할 수 없었습니다. + +이 PR은 `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`에서 네 가지를 합니다. CUDA kernel 본문에서 lane fold만 `__shfl_down(acc, o, 32)`로 바꾼 `SSM_HIP_SOURCE`를 추가하고 `.rocm`을 채웁니다. `ssm_kernel_available()`이 모든 플랫폼에서 `has_kernel_port(ssm_ports())`를 반환하게 하고, 새 `MLXCEL_SSM_KERNEL=0` 스위치를 두며 예전 `MLXCEL_SSM_CUDA_KERNEL=0`은 alias로 남깁니다. Nemotron-H가 `A_log`를 bf16 activation 옆에 f32로 저장하고 JIT cache가 template argument로 key를 만들기 때문에, launch의 template argument에 `A_log`, `B`, `C`의 dtype을 모든 backend에서 추가합니다. Kernel이 index할 수 없는 shape은 host에서 error로 거부합니다. + +gfx1151에서 같은 binary로 `MLXCEL_SSM_KERNEL=0`을 "이전" arm으로 두고 측정한 결과, decode가 granite-4.0-h-tiny-4bit에서 60.39에서 88.43 tok/s (1.46x), NVIDIA-Nemotron-3-Nano-30B-A3B-4bit에서 51.35에서 74.37 tok/s (1.45x)로 올랐습니다 (세 번 실행의 중앙값). 모든 traced step이 512 token SSM state를 만나는 `w1ctx512` logit trace는 이제 ROCm kernel과 M5의 Metal kernel을 비교하며, decided position 불일치는 0 / 60 (granite), 0 / 71 (Nemotron-H)입니다. Kernel 유무에 따른 greedy 128 token 출력은 두 모델 모두 자연스럽지만 byte 단위로 같지는 않으며, trace는 모든 갈림이 undecided position에서 일어났음을 보여 줍니다. + +이 issue 범위를 넘는 발견이 하나 있습니다. 모든 #1814 port가 갖추도록 요구되는 wave32 `#error` guard는 ROCm 10의 AMD clang 23에서 아무 효과가 없습니다. 이 컴파일러는 gfx1151, gfx942, gfx90a 어느 것에도 `__AMDGCN_WAVEFRONT_SIZE__`나 `__AMDGCN_WAVEFRONT_SIZE`를 정의하지 않습니다. 이 port는 guard에 의존하지 않지만, 이후 port에 대해서는 이 관용구가 아무것도 막아 주지 않습니다. + +## 1. 문제 정의 + +### 1.1 묻지 않는 predicate 뒤의 빈 slot + +Hybrid SSM 모델 (granite-4.0-h, falcon-h1, plamo-2, Nemotron-H)은 `granitemoehybrid.rs`, `falcon_h1.rs`, `plamo2.rs`, `nemotron_h.rs`에서 fused single-token 경로를 `seq_len == 1 && ssm_kernel_available()`로 gate하고, 그 뒤 launch를 `expect`합니다. 이 PR 이전의 predicate는 다음과 같았습니다. + +```cpp +#ifdef __APPLE__ + return mlx::core::metal::is_available(); +#else + // MLXCEL_SSM_CUDA_KERNEL=0 forces the graph path + ... + return mlx::core::cu::is_available(); +#endif +``` + +ROCm에서는 false가 나와 gate가 graph로 갔습니다. `.rocm`만 채웠다면 아무것도 바뀌지 않았을 것이고, predicate가 backend를 나열하는 형태로 남은 채 slot을 채우는 것은 `verify-kernel-port-dispatch`가 막으려는 바로 그 모양, 즉 predicate와 dispatch가 서로 다르게 답할 수 있는 상태입니다. + +### 1.2 Profile이 말한 가치 + +PR #2086은 gfx1151에서 SSM update를 #1814 port 중 회수 가능한 비율이 가장 큰 항목으로 측정했습니다. granite-4.0-h-tiny decode GPU 시간의 29.8%, Nemotron-H의 20.0%, 그리고 두 모델 dispatch의 절반 이상입니다. CUDA port (#631)는 GB10에서 granite/falcon 계열 decode를 2.6~4.5x 올렸습니다. + +### 1.3 Kernel 단위 parity test가 없었음 + +CUDA port의 parity는 CUDA에서의 end-to-end greedy 출력이었습니다. 같은 입력으로 kernel과 graph step을 비교하는 것이 없었으므로, 유한하고 그럴듯한 값을 내는 reduction 버그 (#1814의 wave32 규칙이 설명하는 실패)는 모델 출력이 서서히 어긋나는 형태로만 드러났을 것입니다. + +## 2. 변경 요약 + +| 영역 | 변경 | +|---|---| +| `mlx_cxx_kernels.cpp`, kernel | `SSM_HIP_SOURCE`, `MLXCEL_BRIDGE_ROCM_BACKEND`에서 `fast::hip_kernel`을 부르는 `SsmKernelHolderHip` / `get_ssm_kernel_hip()`, `ssm_ports()`의 `.rocm` getter | +| `mlx_cxx_kernels.cpp`, predicate | `ssm_kernel_available()`: kill switch, 그다음 GPU device가 아니면 false, 그다음 `has_kernel_port(ssm_ports())` | +| `mlx_cxx_kernels.cpp`, launch | 모든 backend에서 `TA`, `TB`, `TC` template argument (`A_log`, `B`, `C`의 dtype) | +| `mlx_cxx_kernels.cpp`, host | `ssm_update_kernel`에서 launch 전에 `std::invalid_argument`를 던지는 rank와 shape 검사 | +| `ssm_update_parity_tests.rs` (신규) | Test 다섯 개: 세 shape에서 f32와 bf16 parity, `A_log` dtype key, shape 거부, 두 kill switch | +| `mlx_cxx_bridge.h`, `lib.rs` | predicate와 kernel의 doc comment, test module 등록 | +| 문서 | `docs/environment-variables.md`의 `MLXCEL_SSM_KERNEL` 행과 alias 행, `docs/installation.md`의 ROCm 행, `rocm-correctness-gfx1151-2026-09-30.md`의 "Open" 항목이 kernel 대 kernel 행을 가리키도록 수정, 새 `rocm-ssm-update-kernel-gfx1151-2026-10-04.md` | +| 데이터 | `benchmarks/rocm_strixhalo-gfx1151_2026-10-04_ssm-kernel-{off,on}.csv`, `benchmarks/logit_traces/rocm_gfx1151_96cbce84/`의 trace 여섯 개와 `METADATA.txt`, `RUNS.txt`, `SHA256SUMS`, `README.md` | + +Branch의 commit은 여섯 개입니다. port와 predicate (`96cbce84`), 모든 backend의 dtype key (`844bd94c`), 문서 (`3780b0b2`), CPU device 검사와 review 수정 (`bd2e8f31`), host 측 shape 거부 (`38071038`), 추가 거부 test와 test 문서 (`2f17f286`). + +## 3. Port + +### 3.1 CUDA 본문, HIP fold + +HIP source는 CUDA port의 kernel 이름 형식, 입력 (`X`, `A_log`, `B`, `C`, `D`, `dt`, `state_in`), 출력 (`out`, `state_out`), grid, template argument를 그대로 유지합니다. (32, 8, 1) threadgroup에서 `threadIdx.y`가 같은 32개 lane이 (head, row) 한 쌍을 reduce합니다. 각 lane은 `Ds / 32`개의 state 원소를 돌며 state를 갱신하고 `state * C`를 누적합니다. 본문에서 바뀐 것은 fold뿐입니다. + +```cpp +for (int o = 16; o > 0; o >>= 1) { + acc += __shfl_down(acc, o, 32); +} +``` + +HIP에서 `__shfl_down_sync`는 mask를 무시하는 호환용 shim이므로 native `__shfl_down`을 쓰고 width를 명시합니다. Wave 크기와 무관하게 fold가 맞는 이유가 이 width입니다. Wave32 target (gfx11, gfx12)에서는 32개 lane이 wave 전체이고, wave64 target (CDNA)에서는 width 32가 wave를 32 lane 두 구간으로 나누며 그것이 정확히 row 두 개입니다. `Dh`를 넘는 row의 early return은 wave의 32 lane이 모두 같은 `threadIdx.y`를 가지므로 wave-uniform입니다. + +### 3.2 Wave32 guard와, 여기서 그것이 아무 일도 하지 않는 이유 + +#1814 port 요구사항은 lane 간 reduce를 하는 모든 port에 bitlinear port처럼 `__AMDGCN_WAVEFRONT_SIZE__`와 `__AMDGCN_WAVEFRONT_SIZE`에 대한 `#error` 검사 두 개를 넣도록 합니다. 이 port도 넣었습니다. 작업 중에 ROCm 10의 AMD clang 23이 gfx1151, gfx942, gfx90a 어디에도 두 macro를 정의하지 않는다는 것이 드러났습니다. 따라서 `#if defined(...)` 조건은 모든 target에서 false이고 guard는 발동할 수 없습니다. `SSM_HIP_SOURCE`의 주석과 benchmark 문서 모두 이 사실을 적고 있습니다. 위 fold가 wave64에서 맞는 것은 guard 때문이 아니라 width 논리 때문입니다. 남은 port에 대한 의미는 9절에 있습니다. + +## 4. Predicate와 스위치 + +`ssm_kernel_available()`은 이제 다음 순서로 판단합니다. + +1. `MLXCEL_SSM_KERNEL=0` 또는 `MLXCEL_SSM_CUDA_KERNEL=0`이면 false. 새 이름은 모든 backend에서 스위치가 하는 일을 그대로 말하고, 예전 이름은 기존 A/B script가 계속 동작하도록 남깁니다. +2. 기본 device가 GPU가 아니면 false. `mamba1_scan_kernel_accepts`와 같은 검사입니다. Custom kernel은 GPU stream에서만 실행되므로, `MLXCEL_DEVICE=cpu`에서는 launch가 throw하는 대신 gate가 CPU에서 `ssm_step`을 탑니다. +3. 그 외에는 `mlxcel::has_kernel_port(ssm_ports())`. Dispatch에서 `select_kernel_port`가 읽는 것과 같은 table입니다. 그래서 true라는 답은 dispatch가 거부하지 않는다는 뜻이고, Rust gate가 launch를 `expect`할 수 있는 근거가 됩니다. + +다른 backend에는 두 가지 동작 변화가 생기며 문서에 둘 다 적혀 있습니다. Alias가 이제 Metal과 ROCm에도 적용되고 (예전에는 non-Apple 경로에서만 읽었습니다), Metal에서도 `MLXCEL_SSM_KERNEL=0`이 반영됩니다. 이전에는 Metal에서 kernel을 끌 방법이 없었습니다. 환경 변수가 없으면 Metal과 CUDA에서 `has_kernel_port`는 이전처럼 true입니다. + +모델 파일은 바뀌지 않았습니다. Nemotron-H에서 `NemotronHMamba2Mixer::forward`의 gate는 `fused_mamba2_forward`를 선택하는데, 이것은 single-token mixer 전체 (input projection, convolution, 이 kernel, gated norm, output projection)를 C++ 호출 하나로 실행합니다. 따라서 Nemotron-H의 decode 수치는 SSM step만이 아니라 이 경로 전체를 Rust graph mixer와 비교한 것입니다. + +## 5. 모든 backend의 dtype key + +CUDA와 HIP의 JIT cache는 kernel 이름과 `template_args`로 module의 key를 만들지만, 생성되는 kernel signature는 각 입력의 runtime dtype을 받습니다. Launch는 `T` (activation, `X`와 `D`도 포함)와 `U` (state)만 이름에 넣었습니다. Nemotron-H는 `A_log`를 bf16 activation 옆에 f32로, granite는 bf16으로 저장하므로, 두 모델이 같은 cached module에 닿는 곳 (parity test처럼 한 process 안)에서는 먼저 compile한 쪽이 다른 쪽이 `A_log`를 읽을 pointer type을 정해 버립니다. `B`와 `C`도 `T`에 묶여 있지 않습니다. + +수정은 `A_log`, `B`, `C`의 dtype으로 `TA`, `TB`, `TC`를 추가합니다. Kernel 본문은 이 이름을 쓰지 않으므로 바뀌는 것은 module 이름뿐이고 연산은 같습니다. 이것은 Metal을 포함한 모든 backend에 추가됩니다 (Metal은 이미 이름에 모든 입력 dtype을 넣으므로 중복입니다). 첫 버전 (`96cbce84`)은 Metal과 CUDA module 이름을 바꾸지 않으려고 `if (mlxcel::gpu_kernel_backend() == mlxcel::GpuKernelBackend::Rocm)` 안에서 추가했습니다. 이것은 `scripts/ci/check_kernel_port_dispatch.py`의 규칙 1이 금지하는 형태입니다. Custom kernel을 launch하는 파일은 `kernel_port.cpp`와 `gpu_backend.cpp` 밖에서 backend 종류로 분기하면 안 됩니다. `select_kernel_port`가 생기기 전 launcher들이 서로 어긋난 원인이 손으로 쓴 backend별 분기였기 때문입니다. 두 번째 commit (`844bd94c`)이 key를 모든 backend로 옮겼습니다. Trace는 `96cbce84`에서 만들었으며, ROCm에서는 두 commit의 launch, argument, cache key가 같다는 것을 `METADATA.txt`가 기록합니다. + +`ssm_update_kernel_keys_on_a_log_dtype`은 같은 shape을 한 process에서 `A_log` bf16으로 먼저, f32로 다음에 실행합니다. Key가 없으면 두 번째 launch가 bf16 module을 재사용하고 test는 normalized RMS 0.80으로 실패합니다. + +## 6. Host 측 shape 거부 + +Kernel은 `Dh`, `Ds`, `H`, `G`를 template 상수로 받아 raw buffer를 index합니다. Graph 경로라면 `reshape`이나 `repeat`에서 거부했을 checkpoint 설정이, kernel에서는 GPU memory를 잘못 읽거나 쓰지 않은 채 남깁니다. `ssm_update_kernel`은 이제 다음 경우 launch 전에 `std::invalid_argument`를 던집니다. + +- `hidden_states`나 `B`가 rank 4가 아님 +- `C`의 shape이 `B`와 다름 +- token이 하나보다 많거나, `B`의 batch가 activation과 다름 +- group이 양수가 아니거나 head가 group으로 나누어지지 않음 +- state 폭이 32 미만이거나 32의 배수가 아님 +- `A_log`, `D`, `dt`, `state_in`의 원소 수가 shape과 맞지 않음 + +Bridge는 이 throw를 Rust에 error로 돌려줍니다. 검사는 공유 launch 앞에 있으므로 모든 port에 적용되며, HIP port가 ROCm에서 이 경로를 도달 가능하게 만든 것입니다. `ssm_update_kernel_refuses_unsupported_shapes`가 이를 고정합니다. + +## 7. 측정 결과 + +### 7.1 Decode 처리량 + +`scripts/bench_decode.sh`, pp512/tg128, 실행당 모델 하나, arm마다 세 번, arm 순서를 번갈아, `844bd94c`, 같은 binary. + +| 모델 | Graph (`MLXCEL_SSM_KERNEL=0`) tok/s | Kernel tok/s | 속도 향상 (중앙값) | +|---|---|---|---| +| granite-4.0-h-tiny-4bit | 60.57 / 60.39 / 60.05, 중앙값 60.39 | 88.36 / 88.69 / 88.43, 중앙값 88.43 | 1.46x | +| NVIDIA-Nemotron-3-Nano-30B-A3B-4bit | 51.65 / 50.87 / 51.35, 중앙값 51.35 | 74.87 / 72.87 / 74.37, 중앙값 74.37 | 1.45x | + +실행 간 편차는 두 arm 모두 granite 1% 미만, Nemotron-H 3% 미만입니다. 모델별 마지막 행을 쓰는 `compare_bench_csv.py --before --after `은 1.47x와 1.45x를 보고합니다. Prefill은 이 경로에 있지 않아 비교하지 않았고, 공유 호스트에서 granite prefill은 두 arm에 걸쳐 379~608 tok/s 사이로 변했습니다. + +향상 폭은 GB10의 CUDA port 2.6~4.5x보다 작습니다. 동시에 decode profile의 GPU 시간 비율만으로 예상할 수 있는 것보다는 큽니다. 29.8%나 20.0% 비율을 통째로 없애면 GPU 시간 기준 약 1.42x와 1.25x입니다. 다만 profile은 다른 commit에서, 이 모델들에서 그 자체로 20~23%의 비용이 드는 profiler 아래에서 측정되었고, Nemotron-H의 fused 경로는 SSM step이 아니라 mixer 전체를 대체합니다. 그러므로 이것은 공개된 비율과의 일관성 점검이지 속도 향상의 분해가 아닙니다. 초과분을 설명하는 측정은 이 PR에 없습니다 (약 55개 op graph의 host 측 dispatch가 가장 유력한 후보입니다). + +### 7.2 Idle-GPU guard + +모든 실행은 `scripts/rocm_gpu_guard.sh`를 거쳤습니다 (`/sys/class/kfd/kfd/proc`가 비어 있고 compiler가 없는 상태로 90초, 그다음 1 Hz 감시). 병렬 개발 unit이 GPU를 공유했으며, 그 작업과 겹친 시도는 모두 거부되어 다시 실행했습니다. 한 번을 제외한 모든 실행은 #2098에서 고친 guard (commit `c3eab7b2`)를 #2065 unit의 worktree에서 실행했으며, 이 PR에는 커밋되지 않았습니다. 그 수정은 guard가 명령 자신의 종료된 자식 process를 이유로 거부하던 문제를 없앱니다. 예외는 첫 실행 (granite, graph arm, r1)으로, `main`의 guard를 사용했습니다. 예전 guard의 결함은 잘못된 거부만 만들고 잘못된 수용은 만들지 않았으므로, 그 실행의 수용은 유효합니다. + +### 7.3 Kernel 대 graph, 한 step + +`ssm_update_parity_tests`는 출력과 새 state를 one-token SSD step의 float32 MLX-op reference와 비교합니다. + +| Shape | Batch | Head 수 x head 차원 | Group | State | +|---|---|---|---|---| +| granite-4.0-h-tiny | 1 | 48 x 64 | 1 | 128 | +| Nemotron-H (dt를 `(1e-3, 0.1)`로 clip) | 2 | 64 x 64 | 8 | 128 | +| padded row | 1 | 4 x 60 | 2 | 64 | + +허용 오차 (normalized RMS / max)는 issue에서 정한 f32 1e-5 / 1e-4, bf16 1.6e-2 / 7e-2입니다. 모델이 state를 f32로 들고 다니므로 state는 모든 경우에 f32 기준을 적용합니다. Gfx1151에서 통과합니다. 두 가지 변형으로 test가 구별력을 가진다는 것을 보였습니다. Lane fold를 16 대신 8에서 시작하면 normalized RMS 0.65~0.89로 실패하고, dtype key를 빼면 0.80으로 실패합니다. Test가 일찍 반환하는 것은 GPU backend가 없을 때, 기본 device가 CPU일 때, kill switch가 설정되었을 때뿐입니다. Metal, CUDA, ROCm에서 predicate가 false이면 skip이 아니라 실패입니다. + +### 7.4 모델 logit: kernel 대 kernel (`w1ctx512`) + +`w1ctx512` trace는 corpus token 512개를 prefill한 뒤 single-token step을 trace하므로, 모든 traced step이 state를 만나 fused step을 탑니다. 이 PR 전까지 `rocm-correctness-gfx1151-2026-09-30.md`의 해당 행은 Metal kernel과 ROCm graph를 비교했습니다. 이제는 kernel과 kernel을 비교합니다. `compare_logit_traces.py --decided 2.0`, Nemotron-H 파일은 `[NemotronH]` loader 행을 걸러 냈습니다. + +| 모델 | Reference | Candidate | Top-1 불일치 | Decided 불일치 | 최대 gap | Perplexity ref / cand | +|---|---|---|---|---|---|---| +| granite-4.0-h-tiny | Metal kernel (`metal_m5_3c9edea0`) | ROCm kernel | 3 / 128 | 0 / 60 | 0.125 | 7.090 / 7.216 | +| granite-4.0-h-tiny | Metal kernel | ROCm graph (`ssmkernel0`) | 2 / 128 | 0 / 60 | 0.125 | 7.090 / 7.290 | +| granite-4.0-h-tiny | ROCm graph (`ssmkernel0`) | ROCm kernel | 1 / 128 | 0 / 59 | 0.000 | 7.290 / 7.216 | +| granite-4.0-h-tiny | `3c9edea0`의 ROCm graph | ROCm kernel | 1 / 128 | 0 / 60 | 0.125 | 7.234 / 7.216 | +| nemotron-3-nano-30b-a3b | Metal kernel (`metal_m5_3c9edea0`) | ROCm kernel | 3 / 128 | 0 / 71 | 0.250 | 7.489 / 7.442 | +| nemotron-3-nano-30b-a3b | Metal kernel | ROCm graph (`ssmkernel0`) | 4 / 128 | 0 / 71 | 0.250 | 7.489 / 7.460 | +| nemotron-3-nano-30b-a3b | ROCm graph (`ssmkernel0`) | ROCm kernel | 1 / 128 | 0 / 68 | 0.125 | 7.460 / 7.442 | +| nemotron-3-nano-30b-a3b | `3c9edea0`의 ROCm graph | ROCm kernel | 4 / 128 | 0 / 68 | 0.125 | 7.489 / 7.442 | + +모든 조합에서 decided position 불일치는 0이고, 모든 top-1 불일치에서 reference의 token은 candidate의 2위나 3위이며 reference gap은 0.25 이하입니다. 같은 binary의 `default`와 `ssmkernel0` trace가 서로 다르다는 점 (perplexity가 다르고 각각 token 하나가 뒤집힘)이 `default` trace가 실제로 kernel을 탔다는 증거입니다. Metal kernel 기준으로 ROCm kernel은 ROCm graph만큼 가깝습니다. Top-1 불일치가 3과 3인데, 여기서 graph는 2와 4, 예전 `3c9edea0` 비교에서는 4와 7이었습니다. + +`w1` trace (prefill 없음)는 state를 만나지 않으므로 kernel 유무와 관계없이 graph를 탑니다. `rocm_gfx1151_c5fe9a16`과 비교하면 두 모델 모두 모든 position에서 일치합니다 (granite top-1 0 / 128, decided 0 / 116, Nemotron-H top-1 0 / 128, decided position 없음). 이것은 issue의 `w1` 기준을 충족하고 stateless 경로가 바뀌지 않았음을 보이지만 kernel 자체에 대해서는 아무것도 말하지 않습니다. 그래서 정확성 논거는 `w1ctx512` 행이 담당합니다. + +### 7.5 Greedy 출력의 갈림 + +`mlxcel generate -p "The history of the Roman Empire" -n 128 -t 0 --no-chat-template`을 `MLXCEL_SSM_KERNEL=0` 유무로 실행했습니다. 두 모델의 두 arm 모두 자연스러운 글을 내지만 byte 단위로 같지는 않으며, GB10의 CUDA port와는 다릅니다. Granite는 세 번째 생성 token에서 갈리고 ("a rich and complex tapestry" 대 "a fascinating and complex subject"), Nemotron-H는 열두 token 정도 뒤에 갈립니다. 자유 생성은 첫 갈림 이후 모든 것이 다른 글에 조건화되므로, kernel이 틀렸는지 모델이 망설였는지를 구별하지 못합니다. Teacher-forced trace가 그 답입니다. Kernel과 graph는 모델마다 top-1 token 하나에서만 다르고, decided position에서는 한 번도 다르지 않습니다. 갈림은 합산 순서 변화가 near-tie를 만난 결과이지 오류가 아닙니다. + +## 8. 기술적 선택과 그 이유 + +- **Metal이 아니라 CUDA 본문을 port.** CUDA와 HIP은 launch 형태, grid, `template_args`를 공유하고, `fast::hip_kernel`은 `cuda_kernel`처럼 hipRTC로 compile합니다. Fold를 빼고 본문을 한 줄씩 같게 두면 둘이 다를 때 볼 곳이 하나뿐입니다. +- **Guard에 기대지 않고 shuffle width를 명시.** Width 32 fold는 구조상 wave32와 wave64 모두에서 맞습니다. 요구사항이 정한 guard가 ROCm 10의 어느 target에서도 발동하지 않으므로 이것이 실제로 중요했습니다. +- **Predicate가 backend가 아니라 table을 읽게 함.** `has_kernel_port(ssm_ports())`는 `select_kernel_port`와 다르게 답할 수 없고, 앞으로 slot을 채우는 port는 predicate 수정 없이 도달 가능해집니다. `#ifdef __APPLE__` 분기와 `cu::is_available()` 호출이 없어졌고, issue의 첫 acceptance criterion이 이를 확인합니다. +- **스위치 이름을 바꾸고 alias는 유지.** `MLXCEL_SSM_KERNEL`은 이제 모든 backend에 작용하는 스위치를 그대로 설명하고, `MLXCEL_SSM_CUDA_KERNEL`을 남겨 기존 A/B script가 깨지지 않게 합니다. 대가는 CUDA라는 이름의 변수가 이제 Metal과 ROCm에서도 kernel을 끈다는 것이며, 문서에 적혀 있습니다. +- **ROCm 분기가 아니라 모든 backend에 dtype key.** Launch에서의 backend 분기가 첫 시도였고 `verify-kernel-port-dispatch`에 거부되었습니다. 모든 곳에 key를 추가하는 것은 Metal에서 무해하고, 같은 변경으로 CUDA의 같은 key 공백도 메웁니다. +- **잘못된 shape은 모든 port에 대해 host에서 거부.** Graph 경로는 이런 shape에서 시끄럽게 실패하지만 kernel은 조용히 실패합니다. `std::invalid_argument`를 던지면 backend별 조건 없이 시끄러운 실패가 돌아옵니다. +- **CPU device 검사를 predicate에 둠.** Mamba1 scan과 같은 규칙이고, `MLXCEL_DEVICE=cpu`를 동작하는 설정에서 error로 바꾸지 않고 graph 경로에 둡니다. + +## 9. 후속 작업: wave64 guard 관용구는 ROCm 10 clang에서 무력함 + +#1814 port 요구사항은 `#if defined(__AMDGCN_WAVEFRONT_SIZE__) && __AMDGCN_WAVEFRONT_SIZE__ != 32`와 `#error`, 그리고 `__AMDGCN_WAVEFRONT_SIZE`에 대한 같은 검사를 요구합니다. 32 lane을 가정하는 port가 wave64 CDNA에서 유한하지만 틀린 결과를 내는 대신 compile에 실패하게 하려는 것입니다. ROCm 10의 AMD clang 23은 gfx1151, gfx942, gfx90a 어디에도 두 macro를 정의하지 않으므로, 두 조건이 모두 false이고 guard는 이 중 어디에서도 발동하지 않습니다. + +현재 정확성에는 영향이 없습니다. `mlx_cxx_kernels.cpp`에서 lane 간 reduce를 하는 HIP port 두 개 (bitlinear와 이 port)는 모두 `__shfl_down`에 width 32를 명시하고, 이는 wave64에서도 맞습니다. 문제는 앞으로 올 port (#2065, #2064, #2068, #2063, #2069)입니다. 32 lane wave에 정확성이 실제로 의존하는 port가 CDNA에서 아무 문제 없이 compile되고, 요구사항 문구는 reviewer가 보호받고 있다고 믿게 만듭니다. 요구사항에는 ROCm 10에서 실제로 발동하는 장치가 필요하거나, 모든 lane 간 연산이 width를 명시하도록 하는 규칙으로 바꾸고 compile 시점 guard를 주장하지 않아야 합니다. 이 PR은 대체 방법을 정하지 않으며, 이 보고서를 쓰는 시점에 이에 대한 issue는 찾지 못했습니다. + +## 10. 검증 + +PR 본문 기준, gfx1151 (Radeon 8060S, ROCm 10.0.0)에서: + +- `38071038`에서 `make verify-rocm`: OK. Test binary 146개에 걸쳐 11,856 passed, 0 failed, 378 ignored. ROCm smoke OK. +- `2f17f286`에서 `cargo test --release --features rocm -p mlxcel-core --lib ssm_update_parity_tests -- --test-threads=1` (gate 이후 test와 문서만 변경): 5 passed. +- `cargo clippy -p mlxcel-core --lib --tests --features rocm -- -D warnings`: clean. `make verify-fmt verify-kernel-dtype-keys verify-kernel-port-dispatch`: pass, `EXPECTED_IN_SCOPE`와 그 개수는 그대로. +- CPU device와 shape 검사 전후로 두 모델의 greedy 출력이 동일하므로, `844bd94c`와 `96cbce84`에서 측정한 decode와 trace 수치가 head에도 적용됩니다. + +Orchestrator 검증, head `2f17f286` (origin/main `c0b71344`과 최신 상태): + +- `make verify-rocm`이 모든 단계를 통과: 11,856 tests passed, 0 failed, 378 ignored, smoke OK. + +## 11. 남은 위험과 검증하지 않은 부분 + +- **Metal과 CUDA는 실행하지 않음.** 둘 다 predicate (여전히 `has_kernel_port`이고 둘 다 true, 여기에 CPU device와 kill switch 규칙), 추가된 `TA`/`TB`/`TC` template argument (module 이름만), host shape 검사의 영향을 받습니다. Parity test는 둘 다에서 실행되도록 작성되었고, 거기서 predicate가 false이면 skip이 아니라 실패합니다. +- **Alias의 범위가 넓어짐.** Metal이나 ROCm 머신에서 `MLXCEL_SSM_CUDA_KERNEL=0`을 설정해 두고 잊은 사용자는 이제 그곳에서 graph 경로를 탑니다. +- **Hybrid SSM 모델로 `MLXCEL_DEVICE=cpu`는 실행하지 않음.** 공유 test binary 안에서 기본 device를 바꾸면 다른 test가 CPU로 옮겨지므로 CPU device 규칙을 다루는 test는 없습니다. +- **falcon-h1과 plamo-2는 ROCm에서 실행하지 않음.** 같은 gate를 타지만 측정한 것은 granite와 Nemotron-H뿐입니다. +- **장치 하나, 세션 하나, 공유 GPU.** 모든 수치는 gfx1151에서 나왔고, 개발 unit 하나가 GPU를 공유했으며 겹친 시도는 다시 실행했습니다. 공유 호스트에서 prefill은 크게 변해 비교하지 않았습니다. Wave64 하드웨어가 없었으므로 CDNA에 대한 width 논리는 실행이 아니라 추론입니다. +- **측정용 guard는 이 PR에 없음.** 한 번을 제외한 모든 실행은 이 branch 밖의 #2098 `c3eab7b2`를 사용했고, 공개된 방법은 별도로 들어가는 수정을 참조합니다. + +## 12. 학습 포인트 + +- **Predicate가 table을 읽기 전까지 port는 도달 불가능.** `ssm_kernel_available()`을 바꾸지 않고 `.rocm`만 채웠다면 아무도 부르지 않는 kernel과 녹색 build가 남았을 것입니다. +- **JIT cache key는 signature가 받는 모든 입력 dtype을 포함해야 함.** dtype이 섞인 checkpoint (f32 `A_log`, bf16 activation)가 바로 불완전한 key가 엉뚱한 module을 조용히 재사용하는 경우이고, 두 dtype을 한 process에서 돌리는 test만 이를 잡습니다. +- **Guard에 기대기 전에 발동할 수 있음을 증명.** Wave32 `#error`는 보호처럼 보였고 어디서나 compile되었는데, 이는 아무것도 검사하지 않는 guard와 같은 증상입니다. +- **Parity test를 변형으로 점검.** 8에서 시작하는 fold와 빠진 dtype key가 둘 다 큰 편차로 실패하므로, 허용 오차가 느슨하지 않고 의미가 있다는 것이 드러납니다. +- **순서를 바꾸는 변경에 greedy 글은 맞지 않는 도구.** 자유 출력은 첫 near-tie에서 갈라져 계속 갈라진 채로 남습니다. Decided position 기준을 둔 teacher-forced trace가 정밀도 변화와 버그를 구별합니다. + +Refs: #2067, #1814, #1801, #2061, #2086, #631, #1862, #1870, #2059, #2065, #2098, #2064, #2068, #2063, #2069.