[nihalpasham]

What This PR Does

stacked on #1671 (feat/sm100-cute-rs-abstractions).

This PR completes the SM120a (RTX 5090) cute-rs compiler and library surface and migrates every SM120a *_cute example and Qwen model workload onto CuTeDSL 4.7-aligned cute-rs operations:

1. Reachable TiledCopy / ThrCopy (copy_g2s, copy_s2g, copy_r2s), SmemTile::{local_tile, store}, and C-fragment retiling (cute-rs)
2. First-class BFloat16 in CuteCopyG2SOp / CuteLdmatrixOp, StridedTensor dialect lowering, and vectorized Bf16Mma16x8x16 (dialect-cute, mir-importer, cuda-oxide-mlir-export)
3. SM120 NVFP4 (Float4E2M1FN + Float8E4M3FN, SF_VEC = 16) Tensor-Core TiledMma (Mxf4Nvf4TiledMma / Nvfp4TiledMma, CuteNvfp4MmaAtomOp)

BEFORE (source-only wrappers, manual swizzles, CUDA-core NVFP4 prefill):
  • fmha_decode / fmha_prefill:  hand-coded `shared_offset`, raw `*dst.base.add(...)`, direct `ctm_rs::{warp, float}`
  • rope_cache_cute:             `u16` carrier casts for BF16 `copy_g2s`, raw shared-memory pointer offsets
  • ffn_cute / qwen_decoder:     scalar `data[r * stride + c]` reads in `while k < width`
  • qwen3_8b_cute (nvfp4):       hand-coded `(32, 4)` scale bit-shifts, CUDA-core-only prefill (50.71 ms)

AFTER (typed `cute-rs` -> `mir-importer` -> `dialect-cute` -> CUTLASS 4.7 MLIR):
  • fmha_decode / fmha_prefill:  `TiledCopy::get_slice` (`copy_g2s`, `copy_s2g`, `copy_r2s`), `SmemTile::local_tile`
  • rope_cache_cute:             native `SmemTile<BFloat16, ...>` + `GmemMatrix<BFloat16, ...>` + `SmemTile::trailing_views`
  • ffn_cute / qwen_decoder:     `APartition::load_tiled_packed` + `BPartition::load_tiled_packed` -> `cute.copy` (128b)
  • qwen3_8b_cute (nvfp4):       `Sm1xxBlockScaleKMajor<16>` + `Nvfp4TiledCopy` + `cute.ldmatrix` +
                                 `cute_nvgpu.arch.mma.SM120.block_scaled` (`m16n8k64`, `vec_size = 16`) (26.95 ms, 1.88x faster!)

---

1. TiledCopy / ThrCopy, SmemTile, and FMHA Migration

Previously, fmha_decode_sm120a_cute and fmha_prefill_sm120a_cute hand-coded the Swizzle<3, 3, 3> formula (shared_offset) and wrote outputs via raw *mut u32 / *mut u128 pointer casts.

BEFORE (`fmha_prefill_sm120a_cute`):
  fn shared_offset(row: u32, column: u32) -> usize { (row * 64 + (column ^ ((row & 7) << 3))) as usize }
  *dst.base.add(shared_offset(row_0, col)).cast::<u32>() = pack_f16x2(c[0], c[1]);

AFTER (`cute-rs` `TiledCopy` + `SmemTile`):
  let thr_copy_p = SmemTiledCopyP::new().get_slice(lane);
  thr_copy_p.copy_r2s(scores, (warp_id, 0), &mut s_p);
  thr_copy_o.copy_s2g(&part, &mut gmem_o, (q_block * 4 + GROUP, 0));

---

2. First-Class BFloat16 in copy_g2s / ldmatrix & StridedTensor Lowering

• CuteCopyG2SOp and CuteLdmatrixOp now accept BFloat16 directly alongside f16 and u16, removing the u16 carrier cast in rope_cache_cute.
• StridedTensor / StridedTensorMut and Bf16Mma16x8x16::{load_tiled_packed} route aligned global tile reads and writes through Tensor / TensorMut (cute.make_view + cute.zipped_divide + cute.slice + cute.copy), accelerating qwen3_8b_cute BF16 prefill by 2.62x (96.15 ms -> 36.66 ms).

---

3. SM120 NVFP4 (Float4E2M1FN + Float8E4M3FN, SF_VEC = 16) Tensor-Core TiledMma

Nvfp4TiledMma (Mxf4Nvf4TiledMma) pairs ldmatrix.x4 shared-memory loads and canonical Sm120ScaleAtom (E4M3x4) scale quads with CuteNvfp4MmaAtomOp, lowering to cute_nvgpu.arch.mma.SM120.block_scaled (m16n8k64.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue4m3):

Global NVFP4 [M, K] / [N, K]
  ──► Nvfp4TiledCopy::copy_g2s (cp.async.cg 16B) ──► SmemTile<BFloat16, Nvfp4SmemLayout>
  ──► Nvfp4TiledMma::{load_a, load_b_pair}       ──► nvvm.ldmatrix (x4)
  ──► Nvfp4TiledMma::{load_scale_a, load_scale_b}──► Sm1xxBlockScaleKMajor<16> / Sm120ScaleAtom
  ──► Nvfp4TiledMma::gemm                        ──► cute_nvgpu.arch.mma.SM120.block_scaled (vec_size = 16)

---

Verification & Benchmarks (RTX 5090, sm_120a)

Workload	Case	Baseline	After	CuTeDSL Ref
blockscale_gemm_cute	2048x2048x1024 MXFP4	13.376 us	13.376 us	16.832 us
fmha_prefill_sm120a	Runtime Causal	17.952 us	17.952 us	23.200 us
fmha_prefill_sm120a	Runtime Noncausal	18.112 us	17.952 us	23.776 us
fmha_prefill_sm120a	Static Causal	17.952 us	17.792 us	22.848 us
fmha_prefill_sm120a	Static Noncausal	17.856 us	17.888 us	23.808 us
fmha_decode_sm120a	Runtime Noncausal	4.736 us	4.640 us	12.128 us
fmha_decode_sm120a	Static Noncausal	4.384 us	4.480 us	12.544 us
qwen3_8b_cute (36L)	BF16 Prefill (P=32)	96.15 ms	36.66 ms	n/a
qwen3_8b_cute (36L)	NVFP4 Prefill (P=32)	50.71 ms	26.95 ms	53.08 ms
qwen3_8b_cute (36L)	NVFP4 Decode (16 tok)	8.76 ms	8.12 ms	13.20 ms

• scripts/smoketest.sh on RTX 5090 (sm_120a): 241 / 241 examples PASS
• just check & just book: fmt-check, clippy, unit tests, check-guards, check-intrinsics, doc-check, and cuda-oxide-book build all PASS
• All cute-rs-preview examples & models: elementwise_cute, swiglu_cute, nvfp4_gemv_cute, rmsnorm_cute, ffn_cute, rope_cache_cute, blockscale_gemm_cute, fmha_decode_sm120a_cute, fmha_prefill_sm120a_cute, qwen_decoder_cute, and qwen3_8b_cute all PASS

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.