[Rui Ren (rui-ren)]

Description

Adds com.microsoft::Mamba2Verify version 1 (CUDA), a snapshot-aware Mamba2 operator
for Nemotron Bolt. Fuses causal convolution, selective recurrence, and gated grouped
RMS normalization. A device INT32 selector chooses the starting snapshot; each
processed token writes its post-token convolution and SSM snapshots so the caller
can select the accepted/final state on the next invocation.

The retained contract uses BF16 activations/convolution, FP32 recurrent state,
readout and normalization coefficients, head_dim=64, d_state=128, and 1-16 token
rows. It preserves state aliasing, trailing snapshots, eight-row specialization,
and invalid-selector non-finite output behavior. Longer prefills require chunking.
This layer adds no projection, routing, or expert operator and requires no custom-op plugin.

Dependency: #33963 in
native stack #33967;
the eventual landing target is main. The slot-aware single-row replay readout contract is included in this introducing layer.

Summary of Changes

Schema and Kernel Registration

File	Change
onnxruntime/core/graph/contrib_ops/fused_state_space_defs.h	Define Mamba2Verify inputs/outputs, type constraints, and shape inference.
onnxruntime/core/graph/contrib_ops/contrib_defs.cc, onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Register the contrib schema and CUDA kernel.

Implementation, Tests, and Documentation

File	Change
onnxruntime/contrib_ops/cuda/sequence/fused_state_space.cc	Validate geometry/snapshots and dispatch the CUDA implementation.
onnxruntime/contrib_ops/cuda/sequence/fused_state_space_impl.{cu,h}	Implement convolution, recurrence, readout, normalization, and snapshot writes.
onnxruntime/test/contrib_ops/fused_state_space_op_test.cc, onnxruntime/test/contrib_ops/cuda_kernels/fused_state_space_test.cc	Exercise provider-level state/selector behavior and internal CUDA kernels.
docs/contrib_ops/cuda/mamba2_verify.md, docs/ContribOperators.md, docs/OperatorKernels.md	Describe the operator ABI, snapshot behavior, and CUDA support.

Testing

Review-fix snapshot: 4f2bd30fe1f3017213ce9eb620b34ab507f5e392.

Fixed static rank/dimension validation with CPU-only schema tests, dt clamp/large-dt and NaN behavior, strict BF16 ULP/fraction checking and nonuniform 1024-wide normalization. The slot-aware replay correction and MambaSingleRowReplayPreservesMultiRowReadout regression are now owned by this introducing PR (moved from #34003), with both host references updated and the single-slot path preserved. The permanent contract is now docs/contrib_ops/cuda/mamba2_verify.md.

Owning-layer validation: 8 provider + 3 CPU schema + 4 actual internal CUDA gtests; zero failures/skips.

Integrated tip regression: 51 C++ provider/schema tests, 5 actual internal CUDA tests and 2 Python IOBinding tests passed, zero failures/skips. All owning-layer provider cases are retained. The final overlay recompiles the changed CUDA/host/schema/test sources over ABI-matched preserved ORT support; not a clean full-source build, full-CI approval, independent HF parity or new end-to-end qualification. Targeted configured clang-format/Ruff and whitespace checks passed.

Historical source provenance moved out of permanent documentation

Origin: #33940, commit 491dc043952066267b2212ba979f0325da318573, frozen external source export-nemotron-pr-20261030/src/bolt_ops.cu, SHA256 1f587d9eacc356805c37d19e94e7aa4465dd50302576b2ca4b9579d7c5ae8ba5. This identifies the original retained implementation, not a dependency in the merged tree or a promise that later dispatch fixes are absent. The slot-aware replay correction and current contract are documented in this PR and the permanent operator documentation. This is the Bolt precision contract, not independent HuggingFace model parity.

Historical measurements and qualification evidence (not current-head acceptance)

These previously reported results are retained for provenance. They have not been rerun on the rebased head above; earlier CI/build status and benchmark results do not establish current-head CI, accuracy, or merge readiness.

Measured CUDA kernel latency - 2026-11-07

BF16 projection/convolution, FP32 SSM state; 128 heads, eight groups, head dimension 64, state size 128, channels 10240, projection stride 18560, four-tap convolution. One fused R-row production launch is compared with R serial single-row production calls with exact state handoff; every timed invocation starts from the same immutable initial state. This compares three fused kernels with 3*R serial kernels.

Frozen Bolt precision contract: FP32 gate/normalization/gamma followed by BF16 output; this is not an acceptance test for the earlier Lightning cast-before-gamma contract. Conv snapshots, SSM snapshots and xBC are bitwise equal for all three cases in both runs. Y/readout pass normalized RMSE <=0.03 and cosine >=0.999; exact raw metrics are below. The one-row control is approximately unchanged; eight/sixteen-row fusion reduces latency.

Rows	Serial single-row calls us	Fused multi-row call us	Ratio
1	8.91	8.91	1.00x
8	82.98	22.17	3.74x
16	167.01	41.41	4.03x

H100 80GB HBM3, CUDA 12.8.93, driver 580.159.04, GCC 11.4.0. CUDA-event median per dispatch; 64 dispatches per captured CUDA graph, 50 warmups, three graph warmups, and 15 paired timed batches alternating AB/BA. Four rotating allocation banks contain deterministic synthetic weights (seed 617); this is not a cold-cache guarantee or a checkpoint-weight benchmark. Raw p95 values describe batch-average latency, not individual-invocation tails. No profiler runs in the timing window. The assigned GPU had empty admission, one stable compute client during execution, and clean release afterward; clocks were not locked. Two independent 39-case runs passed, including the published replay runner.

Scope: same-build production CUDA dispatch alternatives, not historical parent-commit binaries. The measured runtime is an isolated CUDA provider relink with preserved core/support, not a clean full-source ORT build or end-to-end model qualification. Each measured operator's CUDA source hash matches its owning PR head. A clean source build is a separate replication variant and must be measured again. Kernel-only timings do not establish model TTFT/tok/s, MMLU/LCB, speculative integration, or full-model readiness.

Complete primary/repeat results, numerical metrics, source hashes and benchmark code. All 39 cases passed in both supervised runs; full-model work remains paused. No source push, stack change, merge, auto-merge or ready-for-review promotion is included.

Checklist

• Review feedback addressed with focused regressions
• Documentation and supported contracts updated
• Native stack order and immediate-parent dependencies preserved
• Clean full-source/full-CI qualification (not claimed by focused overlays)

[Rui Ren (rui-ren)]

Source-pinned H100 kernel evidence - 2026-11-07

BF16 projection/convolution, FP32 SSM state; 128 heads, eight groups, head dimension 64, state size 128, channels 10240, projection stride 18560, four-tap convolution. One fused R-row production launch is compared with R serial single-row production calls with exact state handoff; every timed invocation starts from the same immutable initial state. This compares three fused kernels with 3*R serial kernels.

Frozen Bolt precision contract: FP32 gate/normalization/gamma followed by BF16 output; this is not an acceptance test for the earlier Lightning cast-before-gamma contract. Conv snapshots, SSM snapshots and xBC are bitwise equal for all three cases in both runs. Y/readout pass normalized RMSE <=0.03 and cosine >=0.999; exact raw metrics are below. The one-row control is approximately unchanged; eight/sixteen-row fusion reduces latency.

Rows	Serial single-row calls us	Fused multi-row call us	Ratio
1	8.91	8.91	1.00x
8	82.98	22.17	3.74x
16	167.01	41.41	4.03x

H100 80GB HBM3, CUDA 12.8.93, driver 580.159.04, GCC 11.4.0. CUDA-event median per dispatch; 64 dispatches per captured CUDA graph, 50 warmups, three graph warmups, and 15 paired timed batches alternating AB/BA. Four rotating allocation banks contain deterministic synthetic weights (seed 617); this is not a cold-cache guarantee or a checkpoint-weight benchmark. Raw p95 values describe batch-average latency, not individual-invocation tails. No profiler runs in the timing window. The assigned GPU had empty admission, one stable compute client during execution, and clean release afterward; clocks were not locked. Two independent 39-case runs passed, including the published replay runner.

Scope: same-build production CUDA dispatch alternatives, not historical parent-commit binaries. The measured runtime is an isolated CUDA provider relink with preserved core/support, not a clean full-source ORT build or end-to-end model qualification. Each measured operator's CUDA source hash matches its owning PR head. A clean source build is a separate replication variant and must be measured again. Kernel-only timings do not establish model TTFT/tok/s, MMLU/LCB, speculative integration, or full-model readiness.

Source/runtime hashes and supervision receipts

{
  "operator_sources": {
    "projection": {
      "head": "ca5bf6a4c1308b4374238b37884d07edf4c3bae5",
      "cuda_sha256": "9df42138bbe8cbd40114307053398db59f64a5ab25908361a8a655230ce90e71"
    },
    "moe": {
      "head": "a9bdbce374a3b2757e1fd11ef83ce6f83502b0d9",
      "cuda_sha256": "7d9a1454e416753b61342cd78c3ba11f90308320d58f96cc5361c6bfa6ffa486"
    },
    "mamba": {
      "head": "374c7d51d93bfdd29cb0f2efbc76b88e1e812aa7",
      "cuda_sha256": "5e2568fe8aa5794456bf84f122b2e50020f050572043d578fb736ea03796923a"
    }
  },
  "loaded_library_sha256": {
    "libonnxruntime.so.1.32.0": "8e08a5f9b0ae8e67f31bbeaea1942a14f97d2a568d62f1e6b13ad72fcc187d4c",
    "libonnxruntime_providers_cuda.so": "c1b303583039d412416a5aabc04c1748a2d5433ce691c4aa10b9d58ab3d76e0a",
    "libonnxruntime_providers_shared.so": "98cc35d20a8701b425d3b51ebbbd9186744c9abcec5d2270879f6c165a735228"
  },
  "cuda": "12.8.93",
  "driver": "580.159.04",
  "host_compiler": "GCC 11.4.0",
  "device": "NVIDIA H100 80GB HBM3",
  "seed": 617,
  "allocation_banks": 4,
  "primary_supervisor": {
    "status": "PASS",
    "worker_exit": 0,
    "peak_process_bytes": 3147825152
  },
  "repeat_supervisor": {
    "status": "PASS",
    "worker_exit": 0,
    "peak_process_bytes": 3147825152
  },
  "scope": "Isolated production CUDA dispatch microbenchmark; recorded production objects/flags linked with preserved ORT support, NOT a clean full runtime or end-to-end model benchmark.",
  "cuda_and_host_dependency_sha256": {
    "libcudnn_graph.so.9": "2e8a0e4e66c0c8f9dc273c486698614c0561832626fadf76a8a37b94d0e2f895",
    "libcudnn.so.9": "0f43113ec6b171f93236560541f3b95104fed2516b19000a0e7047a5cc12f8ee",
    "libcuda.so.580.159.04": "821992937649fe6b01b0633fa4f47665ff5c02f273ea4696763ad19154cb7c55",
    "libcudart.so.12.8.90": "c3a75b33af334a3486d197dbd1584a2985183ba4688d237a2be5f2f679329920",
    "libcublas.so.12.8.4.1": "031ce6c2cbfbb9468f040527cab5c599069ce5609e73e28f87503881063eac21",
    "libcublasLt.so.12.8.4.1": "10b5e6631cf8115c661eb895ed1533826308b58f7956466f53d236a40c9b622c",
    "libstdc++.so.6": "1fd75fe70354a416d75aef22bcae68c47bd25d20e2d0568c30b1a9838cf62f11",
    "libc.so.6": "3a15d66867d83762c7f2f1e37359cb8f6c5743edb369c65285cb0b1c4f7498bf"
  },
  "production_cuda_codegen_options": [
    "-O3",
    "-std=c++20",
    "--generate-code=arch=compute_90a,code=[sm_90a]"
  ]
}

Primary run: all cases, numerical metrics and all paired latency samples

{"pr":33969,"op":"Mamba2Verify","rows":1,"heads":128,"groups":8,"head_dim":64,"state_size":128,"correctness":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"conv_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"ssm_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"xbc":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"readout":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"baseline":{"median_us":8.909000084,"batch_p95_us":8.926300146,"trials_us":[8.948000148,8.917000145,8.917000145,8.903999813,8.906000294,8.905000053,8.912499994,8.901000023,8.905500174,8.909500204,8.910000324,8.909000084,8.910999633,8.906999603,8.907499723]},"current":{"median_us":8.910999633,"batch_p95_us":8.914550208,"trials_us":[8.910000324,8.910999633,8.902999572,8.911499754,8.906500414,8.917000145,8.913500234,8.906999603,8.910999633,8.913500234,8.909500204,8.912499994,8.907999843,8.908499964,8.913500234]}}
{"pr":33969,"op":"Mamba2Verify","rows":8,"heads":128,"groups":8,"head_dim":64,"state_size":128,"correctness":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"conv_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"ssm_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"xbc":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"readout":{"nrmse":6.08055523951e-08,"cosine":1,"max_abs":7.45058059692e-09,"bitwise":false},"baseline":{"median_us":82.982502878,"batch_p95_us":83.074299246,"trials_us":[83.089001477,82.929499447,82.973003387,82.982502878,83.06799829,82.997001708,83.065502346,82.953996956,82.974500954,82.959502935,83.065502346,82.901999354,83.008997142,82.908503711,82.9859972]},"current":{"median_us":22.168500349,"batch_p95_us":22.269850411,"trials_us":[22.153500468,22.141499445,22.220000625,22.210000083,22.216500714,22.208500654,22.168500349,22.215500474,22.286999971,22.140000015,22.262500599,22.110000253,22.151000798,22.144500166,22.139500827]}}
{"pr":33969,"op":"Mamba2Verify","rows":16,"heads":128,"groups":8,"head_dim":64,"state_size":128,"correctness":{"nrmse":4.23097936912e-09,"cosine":1,"max_abs":2.38418579102e-07,"bitwise":false},"conv_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"ssm_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"xbc":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"readout":{"nrmse":6.15534900543e-08,"cosine":1,"max_abs":1.49011611938e-08,"bitwise":false},"baseline":{"median_us":167.011499405,"batch_p95_us":167.137302458,"trials_us":[167.110994458,167.027503252,167.034998536,167.187005281,167.011499405,166.930004954,166.954502463,166.852995753,166.971504688,166.990503669,167.058005929,167.006492615,167.054504156,166.977003217,167.116001248]},"current":{"median_us":41.409499943,"batch_p95_us":41.497398913,"trials_us":[41.386000812,41.437998414,41.49049893,41.478998959,41.390001774,41.390001774,41.429001838,41.513498873,41.450001299,41.409499943,41.407998651,41.381001472,41.288498789,41.375499219,41.450999677]}}

Independent replay: all cases, numerical metrics and all paired latency samples

{"pr":33969,"op":"Mamba2Verify","rows":1,"heads":128,"groups":8,"head_dim":64,"state_size":128,"correctness":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"conv_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"ssm_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"xbc":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"readout":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"baseline":{"median_us":8.898000233,"batch_p95_us":8.945599757,"trials_us":[8.988999762,8.895000443,8.899499662,8.926999755,8.898000233,8.897500113,8.901500143,8.902500384,8.890500292,8.895499632,8.893500082,8.906999603,8.895999752,8.892999962,8.918999694]},"current":{"median_us":8.907999843,"batch_p95_us":8.976550307,"trials_us":[8.985999972,8.911999874,8.898000233,8.905500174,8.907999843,8.911499754,8.899499662,8.937000297,8.906000294,8.910500444,8.972500451,8.904499933,8.897500113,8.911999874,8.902999572]}}
{"pr":33969,"op":"Mamba2Verify","rows":8,"heads":128,"groups":8,"head_dim":64,"state_size":128,"correctness":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"conv_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"ssm_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"xbc":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"readout":{"nrmse":6.08055523951e-08,"cosine":1,"max_abs":7.45058059692e-09,"bitwise":false},"baseline":{"median_us":82.964003086,"batch_p95_us":83.073948324,"trials_us":[83.073496819,82.887999713,82.932502031,82.936502993,82.95699954,82.964003086,83.075001836,83.007998765,82.948498428,83.010002971,83.008997142,82.958996296,83.030998707,82.956500351,82.99549669]},"current":{"median_us":22.196000442,"batch_p95_us":22.263199277,"trials_us":[22.210000083,22.120000795,22.248499095,22.200999781,22.162999958,22.126499563,22.297499701,22.200999781,22.19549939,22.138999775,22.17900008,22.206500173,22.196000442,22.193999961,22.231999785]}}
{"pr":33969,"op":"Mamba2Verify","rows":16,"heads":128,"groups":8,"head_dim":64,"state_size":128,"correctness":{"nrmse":4.23097936912e-09,"cosine":1,"max_abs":2.38418579102e-07,"bitwise":false},"conv_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"ssm_state":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"xbc":{"nrmse":0,"cosine":1,"max_abs":0,"bitwise":true},"readout":{"nrmse":6.15534900543e-08,"cosine":1,"max_abs":1.49011611938e-08,"bitwise":false},"baseline":{"median_us":166.695997119,"batch_p95_us":166.899549961,"trials_us":[166.887998581,166.677996516,166.770502925,166.672497988,166.695997119,166.72000289,166.675001383,166.553497314,166.81599617,166.595995426,166.926503181,166.62299633,166.721999645,166.650503874,166.757002473]},"current":{"median_us":41.581001133,"batch_p95_us":41.67464897,"trials_us":[41.581001133,41.566498578,41.480001062,41.623000056,41.576001793,41.713498533,41.613999754,41.55600071,41.521999985,41.609998792,41.657999158,41.465498507,41.444499046,41.634000838,41.632000357]}}

Benchmark source, synthetic setup graph and validated replay runner: #33962 (comment)

[Copilot]

🟡 Changes recommended

Shape inference incorrectly accepts rank-1 convolution weights that runtime execution rejects; input 3 needs explicit rank-2-or-3 validation and a rank-1 regression test.

1 open issue

• Reject rank-1 convolution weights during shape inference · New

What changed in this PR

Adds com.microsoft::Mamba2Verify as a built-in CUDA contrib operator above the packed-operator stack. It fuses causal convolution, selective recurrence and gated grouped RMS normalization with caller-selected state snapshots.

Changes:

• Adds the operator schema, CUDA registration and implementation for 1–16 token rows.
• Supports post-token snapshots, state-buffer aliasing, trailing-slot preservation and invalid-selector handling.
• Adds provider/internal regression tests and operator documentation; full-model qualification remains separate.

File	Description
onnxruntime/​test/​contrib_ops/​fused_state_space_op_test.cc	Tests outputs, snapshots, aliasing and registration.
onnxruntime/​test/​contrib_ops/​cuda_kernels/​fused_state_space_test.cc	Tests large normalization grids and offset weights.
onnxruntime/​core/​graph/​contrib_ops/​fused_state_space_defs.h	Defines the schema and shape inference.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Registers the contrib schema.
onnxruntime/​contrib_ops/​cuda/​sequence/​FusedVerification.md	Documents contracts, precision and validation scope.
onnxruntime/​contrib_ops/​cuda/​sequence/​fused_state_space.cc	Validates inputs and manages snapshot outputs.
onnxruntime/​contrib_ops/​cuda/​sequence/​fused_state_space_impl.h	Declares CUDA launch interfaces.
onnxruntime/​contrib_ops/​cuda/​sequence/​fused_state_space_impl.cu	Implements convolution, recurrence and normalization kernels.
onnxruntime/​contrib_ops/​cuda/​cuda_contrib_kernels.cc	Registers the CUDA kernel.
docs/​OperatorKernels.md	Lists CUDA kernel support and types.
docs/​ContribOperators.md	Adds the generated operator reference.

_{Review effort: Balanced}

[Rui Ren (rui-ren)]

Stack #33967 review — parent e50e34972b24 -> head 5809f2531dcb.

Reviewed snapshot layouts and selectors, output initialization/preservation, corresponding state-buffer aliasing, BF16/FP32 boundaries, schema/registration, and stream ordering.

Confirmed the existing load-time validation gap in the conv_weight rank thread: input 3 currently accepts ranks {1,2,3}, because rank 1 falls through to the generic rank-one rule. For heads=4/groups=2, an otherwise-valid model with conv_weight[3072] passes that rank predicate but ComputeInternal rejects it as "Invalid convolution parameters". Handle input 3 separately and reject known ranks other than 2/3; add a Graph::Resolve rank-one regression. This is delayed validation, not a demonstrated unsafe CUDA access.

No additional high-confidence correctness finding. The single-row/multi-row readout-order distinction is addressed by #34003's multi-slot replay change. Validation here was source/control-flow review plus a host reproduction of the rank predicates, not an actual Graph::Resolve run or new CUDA/capture test.

[Tianlei Wu (tianleiwu)]

Reviewed this layer (stack #33967, layer 3) at 5809f2531d.

Kernel and host review. I compiled the PR's .cu standalone (CUDA 13.3, H200) and compared it with an independent FP32/BF16 reference over several hundred geometry/row/slot/selector combinations, including the production 128-head / 8-group / 1024-wide-group shape. Aliased and non-aliased runs were bitwise identical, trailing slots were preserved, out-of-range selectors produced NaN rows, compute-sanitizer (memcheck, racecheck, initcheck, synccheck) was clean, and CUDA-graph replay with a device-updated selector worked. I found no correctness defect in the kernels or the host code. (Not run: an ORT build and the OpTester-based tests; only the CUDA translation unit was exercised.)

Open from the existing review. A rank-1 conv_weight still passes shape inference (thread on fused_state_space_defs.h; the checklist item in the PR description is unchecked). I added a concrete predicate and a GPU-independent regression suggestion on that thread. This should be fixed before merge.

Requests (inline, none blocking on its own)

• Test strength: Y/xbc use the default BF16 tolerance; the time_step_min clamp and the raw > 20 branch are never exercised; the production 1024-wide norm branch only sees uniform gate/gamma; and OnlyMambaRemainsAmongProvisionalOperators asserts the absence of names that do not exist in the tree.
• Docs: the schema doc should state slots >= R and the NaN policy for an out-of-range state_row; FusedVerification.md carries provenance that cannot be verified after merge and sits outside the docs/contrib_ops/cuda/ convention.

Minor / design notes

• Validation messages ("Invalid head parameter", "Invalid convolution parameters", "Invalid norm weight or state_row") do not name the input or the expected/actual shapes. With ten inputs and a rigid derived geometry that makes exporter mistakes hard to diagnose.
• xbc and readout are kernel-internal intermediates (the scratch buffers passed between the conv, SSM and norm kernels) but are required outputs of a stable com.microsoft v1 op. Every call allocates and writes them (about 327 KB + 524 KB at R=16 for 128 heads / 8 groups), and any future fused implementation must still materialize them to honor the ABI. If they exist for qualification only, consider making them optional outputs; otherwise accept the ABI freeze consciously.
• Cross-layer: at this head a single-row call runs MambaSsmKernel<2> while multi-row calls run MambaSsmKernel<4>, so the readout reduction order differs between serial single-row replay and an R-row call (a few BF16 Y elements can differ in a standalone comparison). #34003's rows > 1 || slots > 1 dispatch addresses this for multi-slot caches, as the PR description acknowledges. Until then the test references here still encode the old partition (parts = rows == 1 ? 2 : 4) and only keep passing because of the tolerances.
• fused_state_space_impl.h L15 documents snapshots as [rows, heads, state_size, head_size], but the buffers are [slots, heads, d_state, head_dim] with slots >= rows; and MambaInvalidSelectorsProduceNaNWithoutStateReads cannot observe state reads (outputs are NaN either way), so its name over-claims.

[Tianlei Wu (tianleiwu)]

Re-review of 10c6969943.

Verdict: COMMENT — no blocking findings.

Earlier threads, re-checked on this head — all eight are fixed:

• explicit per-input rank predicate + known-dimension validation, with CPU-only schema regressions outside the CUDA guard;
• time_step_min clamp / raw > 20 coverage, and a NaN-propagating clamp (dt < tmin ? tmin : dt) with MambaTimeStepNaNPropagates;
• ULP-aware Y / xbc comparison instead of the default BF16 tolerance;
• nonuniform coverage of the 1024-wide normalization branch;
• the "absent names" test removed;
• slots >= R and the invalid-selector NaN policy are in the public schema / op reference;
• the contract moved to docs/contrib_ops/cuda/mamba2_verify.md, with provenance removed from in-tree docs.

I also read the three kernels for hazards: aliasing is safe by construction (every thread reads its own conv-channel / SSM slice of the selected snapshot before writing any output row, and slices are disjoint across threads and blocks, so state_row < R with in-place state cannot race); the invalid-selector return precedes every __syncthreads() and is block-uniform; shuffles use full masks; 64-bit index math covers the snapshot offsets and heads <= INT_MAX / 1024 bounds the int channel/stride arithmetic.

Independent verification (stack tip 623a725f5c, which contains this layer; Linux, CUDA 13.0, sm_80;sm_90a, Release; A100-SXM4-80GB)

• FusedStateSpaceSchemaTest (3), FusedStateSpaceTest (8) and the internal FusedStateSpaceKernelTest (4) pass; the full stack set (60 provider tests) also passes 3x in shuffled order.
• compute-sanitizer synccheck, racecheck and initcheck including all Mamba tests: no findings.
• CI on this head has no failures so far.

Nits (non-blocking): the constructor ORT_ENFORCE has no message (inline). The op is Mamba2Verify, but the sources, schema header and test suites are named fused_state_space* / FusedStateSpace*, so grepping for the op name misses them; renaming to match (the docs page already does) would help future maintenance.

[Tianlei Wu (tianleiwu)]

Nit: this ORT_ENFORCE has no message, so a bad heads / n_groups fails with just the stringified condition (the shape-inference failure for the same attributes only says "Invalid Mamba2Verify attributes"). Something like "Mamba2Verify requires heads > 0, n_groups > 0, heads % n_groups == 0 and heads <= INT_MAX / 1024, got heads=", heads_, ", n_groups=", groups_ makes it actionable.