CuTe DSL SSD head_dim=80 for Nemotron Mamba prefill on SM80+ (Blackwell D64 slab)

Detailed description of the requested feature

I would like to contribute CuTe DSL SSD prefill for dim=80, dstate=128 to accelerate Nemotron-3-Nano-4B / Nemotron-Nano-9B-v2 Mamba prefill on SM80+ devices, with a Blackwell D64 slab path on SM100/101/110 (e.g. Jetson AGX Thor).

Today, dim=80, dstate=128 prefill falls back to the serial SSM kernel (selective_state_update_prefill_kernel_simple) because CuteDslSSDRunner::canImplement() only supports (dim, dstate) ∈ {(64,64), (64,128), (128,64), (128,128)}. For Nemotron Nano 4B, this serial kernel is the dominant prefill bottleneck (~74.7 ms/layer at IL=3000 on Thor, ~93% GPU kernel time).

What I implemented

Path A — Generic D80 SSD (SM80+):

• AOT variant ssd_prefill_d80_n128 via existing 5-kernel SSD pipeline (ssd_cutedsl/ssd_prefill.py, --dim 80 --dstate 128)
• canImplement(80, 128), module load/unload, dispatch in CuteDslSSDRunner

Path B — Blackwell D64 slab (SM100–110):

• Reuse native ssd_prefill_blackwell_d64_n128 twice per Mamba layer
• Slab 0: dims [0:64); Slab 1: dims [64:80) zero-padded to D64, unpack tail
• Shared B, C, A, dt, dt_bias, D-vector; workspace pack/unpack via cudaMemcpy2DAsync / cudaMemcpy3DAsync
• getWorkspaceSize() accounts for slab buffers
• No new Blackwell cubin for D=80

Why not native Blackwell D=80 single kernel

My first design extended the proven D64 persistent tile (TMEM O, TMA/WGMMA) to D=80 directly. Single-chunk S=128 passed (rel_err ≈ 0.0006), but production could not be guaranteed:

• D64 uses 512 TMEM columns exactly; D80 direct tile needs ~560 → TMEM/SMEM budget exceeded
• Practical pipelines produced SharedStorage 236–244 KiB > 232 KiB Thor limit
• Multi-chunk S=256 failed (rel_err ≈ 0.12) with reduced stages; D64 regression with same tweaks still passed

I replaced it with generic D80 SSD + D64 slab reuse.

Intended use

Nemotron Mamba2 SSM prefill (head_dim=80, dstate=128, nheads=96, n_groups=8, seq ≥ 128). Decode and unsupported dim/dstate pairs stay on existing paths (decode kernel / serial fallback).

Quantization note

SSD compute uses FP16 tensors for x, B, C, state, output (same as existing SSD path). Nemotron production engines may use NVFP4 for linear layers; that is separate — this SSD kernel does not consume NVFP4 tensors directly.

Measured results (Thor, Nemotron-3-Nano-4B NVFP4 engine, llm_bench, CUDA graph off)

Metric	Serial fallback	Generic D80 SSD	Blackwell D64 slab
E2E IL=3000	1684.6 ms	338.3 ms	222.2 ms (~7.6×)
E2E IL=256	78.8 ms	31.1 ms	27.4 ms (~2.9×)
SSM ms/layer (nsys, IL=3000)	~74.7	~9.8	~3.15
Serial SSM GPU %	93.1%	0%	0%
Throughput IL=3000	~1780 tok/s	~8868 tok/s	~13502 tok/s

Validation

• L2: chunk boundaries S=127/128/129, multi-chunk, ragged varlen, nonzero init state — PASS
• Exact Nemotron shape h96×d80×n128×g8: output norm. max 0.034%, state 0.038%
• C++ suite: 13/13 PASS; NonPaddedNemotronLike plugin test — PASS
• Prefill→decode handoff: finite outputs; CUDA graph capture IL=3000 — PASS
• Not bit-exact vs serial fallback (~0.03% rel error; 1 FP16 ULP max abs)

Proposed upstream change (~6 files)

File	Change
kernelSrcs/build_cutedsl.py	Register ssd_prefill_d80_n128 AOT variant
cpp/kernels/mamba/cuteDslSSDRunner.cpp	canImplement(80,128), D80 dispatch, runPrefillBlackwellD80Slabs, workspace sizing
cpp/kernels/mamba/cuteDslSSDRunner.h	Module declaration, slab API
kernelSrcs/ssd_cutedsl/README.md	Document D80 + Blackwell slab
unittests/ssdCuteDslTests.cpp	D80 boundary/multi-chunk/ragged/init-state/exact-shape tests
tests/python-unittests/test_mamba_plugin.py	Nemotron-realistic dims, test_ssd_nemotron_dim80

No mambaPlugin.cpp changes needed — existing logic prefers SSD when canImplement() succeeds and seq_len >= 128.

No prebuilt cubins committed — CI generates artifacts via build_cutedsl.py (consistent with existing SSD variants).

v1 scope (not included)

• Native Blackwell D=80 persistent single-kernel
• NVFP4 E2E benchmark / generation quality regression study
• SSD + causal_conv1d fusion
• Nemotron 9B D=128 tuning (already covered by ssd_prefill_d128_n128)
• dstate=64 for D=80 (Nemotron uses 128)

I can attach nsys/ncu profiles (baseline vs generic D80 vs Blackwell slab) or share reproduction steps in a follow-up comment.

---

Timeline

Impact if not upstreamed: Should have for Nemotron Nano 4B prefill on Thor — serial SSM remains the fallback and prefill stays ~7.6× slower than the Blackwell slab path at IL=3000 (~93% GPU kernel time on serial SSM).

When: No hard external deadline from my side. I have a working branch ready and can open a PR after this issue is reviewed/approved per CONTRIBUTING.md.

Blocker level: Not a build breaker (serial fallback works), but a significant performance blocker for Nemotron Mamba prefill on Thor-class hardware and any SM80+ deployment where mamba_head_dim=80.

---

Describe alternatives you've considered

1. Keep serial fallback (status quo)
   Works functionally but ~74.7 ms/layer on Thor at IL=3000 and ~93% GPU kernel time. Not acceptable for production Nemotron prefill tuning.

2. Native Blackwell D=80 single kernel
   Attempted first. Single-chunk accurate but TMEM/SMEM limits and multi-chunk correctness failures. Abandoned; D64 slab reuses proven native kernel.

3. Generic D80 only (no Blackwell slab)
   Works on all SM80+ including Thor (~338 ms IL=3000, ~5× vs baseline). Simpler, but ~1.5× slower than slab path (~222 ms). Prefer shipping both: generic for broad SM support, slab for SM100–110.

4. Pad model to D=128 and use ssd_prefill_d128_n128
   Would require ONNX/engine padding changes and wastes compute on padded dims; does not match Nemotron's actual tensor shapes.

5. Optimize serial fallback kernel
   Fundamentally O(seq) recurrence; SSD chunk-parallel scan is the intended Mamba2 prefill path in this codebase.

6. Wait for upstream NVIDIA to implement D80 SSD
   Reasonable long-term, but I already have validated code and profiling evidence and would prefer to contribute it if acceptable.

---

Target hardware/use case

Hardware: DRIVE AGX Thor (SM110); generally SM80+ for generic D80 path; Blackwell SM100/101/110 for native D64 slab (-DENABLE_CUTE_DSL=ssd or ALL).

Use case: Nemotron-3-Nano-4B and Nemotron-Nano-9B-v2 Mamba2 SSM prefill with head_dim=80, dstate=128, nheads=96, n_groups=8, seq ≥ 128, FP16 SSD tensors, batched/ragged prefill via existing plugin workspace.

Not in scope for this feature:

• Mamba decode path (unchanged)
• Unsupported dim/dstate (still serial fallback)
• NVFP4 weight compression itself (only the SSM compute path)
• Nemotron attention / MLP layers (separate kernels)

Start with cpp/kernels/mamba/cuteDslSSDRunner.cpp and its header, then inspect the existing variants registered by kernelSrcs/build_cutedsl.py. Run unittests/ssdCuteDslTests.cpp and tests/python-unittests/test_mamba_plugin.py, focusing on the listed D80 boundary, ragged, multi-chunk, init-state, and Nemotron-shaped cases. Done means the D80 path dispatches correctly, workspace sizing and Blackwell slab handling pass validation, and the README documents the supported variants.