Promote latest kernel artifacts to main

CARD.md

@@ -0,0 +1,55 @@
+# flashrt/fp4-gemm
+
+FlashRT native Blackwell NVFP4 A4W4 GEMM kernels. Both activations and weights
+are packed FP4 inputs; this is not a BF16-activation weight-only operation.
+
+## Functions
+
+- `sfa_size_bytes`
+- `quantize_fp4_sfa_fp16`
+- `quantize_fp4_sfa_bf16`
+- `dequantize_fp4_sfa_fp16`
+- `nvfp4_gemm_bf16`
+- `nvfp4_gemm_bias_bf16`
+- `nvfp4_gemm_bias_residual_bf16`
+- `nvfp4_gemm_residual_bf16`
+- `nvfp4_gemm_bias_gelu_bf16`
+- `nvfp4_gemm_bias_gelu_nvfp4`
+- `nvfp4_gemm_streamk_bf16`
+- `nvfp4_gemm_streamk_bias_bf16`
+- `fp4_w4a16_linear_bf16` (compatibility alias)
+
+## Example
+
+```python
+from kernels import get_kernel
+import torch
+
+ops = get_kernel("flashrt/fp4-gemm", version=1, trust_remote_code=True)
+
+x = torch.randn((32, 256), device="cuda", dtype=torch.float16)
+w = torch.randn((512, 256), device="cuda", dtype=torch.float16)
+
+a, sfa = ops.quantize_fp4_sfa_fp16(x, is_sfb=False)
+b, sfb = ops.quantize_fp4_sfa_fp16(w, is_sfb=True)
+y = ops.nvfp4_gemm_bf16(a, b, sfa, sfb)
+```
+
+BF16 activations should use the direct producer to avoid a separate cast and
+copy before every low-bit projection:
+
+```python
+x = torch.randn((1, 5120), device="cuda", dtype=torch.bfloat16)
+a, sfa = ops.quantize_fp4_sfa_bf16(x)
+```
+
+## Notes
+
+- Blackwell `sm_110a` with CUDA 13+ and `sm_120a` with CUDA 12.8+.
+- Inputs are packed FP4 E2M1 plus CUTLASS Sm1xx SFA/SFB scale buffers.
+- Output is BF16.
+- `variant=-1` is the architecture-aware production auto-dispatch;
+  `variant=0/1/2` expose diagnostic default, widen, and pingpong schedules.
+- The canonical BF16-output GEMM and FP4 pack/unpack helpers support SM110 and
+  SM120. SM110 also supports bias, bias+residual, and bias+GELU-to-FP4
+  production epilogues used by the GROOT N1.7 Thor pipeline.

README.md

@@ -1,9 +1,116 @@
-# flashrt/fp4-gemm
+# fp4-gemm
 
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
+FlashRT native Blackwell NVFP4 A4W4 GEMM kernels.
 
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/fp4-gemm
+This package consumes packed FP4 E2M1 tensors plus CUTLASS Sm1xx SFA/SFB scale
+buffers and produces BF16 output. It is designed to pair with
+`flashrt/fp4-fused-ops` and other static low-bit transformer/diffuser runtime
+paths.
 
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.
+## Available Functions
+
+- `sfa_size_bytes(rows, dim)`
+- `quantize_fp4_sfa_fp16(x, packed=None, sfa=None, is_sfb=False)`
+- `quantize_fp4_sfa_bf16(x, packed=None, sfa=None, is_sfb=False)`
+- `dequantize_fp4_sfa_fp16(packed, sfa, out=None, is_sfb=False)`
+- `nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0, out=None, variant=-1)`
+- `nvfp4_gemm_bias_bf16(a_packed, b_packed, sfa, sfb, bias, out=None)`
+- `nvfp4_gemm_bias_residual_bf16(a_packed, b_packed, sfa, sfb, bias, residual, out=None)`
+- `nvfp4_gemm_residual_bf16(a_packed, b_packed, sfa, sfb, residual, alpha=1.0, out=None)`
+- `nvfp4_gemm_bias_gelu_bf16(a_packed, b_packed, sfa, sfb, bias, alpha=1.0, out=None)`
+- `nvfp4_gemm_bias_gelu_nvfp4(a_packed, b_packed, sfa, sfb, bias, alpha=1.0, out_packed=None, out_sfa=None)`
+- `nvfp4_gemm_streamk_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0, out=None)`
+- `nvfp4_gemm_streamk_bias_bf16(a_packed, b_packed, sfa, sfb, bias, alpha=1.0, out=None)`
+- `fp4_w4a16_linear_bf16(...)` is retained as a compatibility alias
+
+## Tensor Contract
+
+- `a_packed`: `torch.uint8`, shape `(M, K / 2)`.
+- `b_packed`: `torch.uint8`, shape `(N, K / 2)`.
+- `sfa`: `torch.uint8`, CUTLASS SFA layout for `(M, K)`.
+- `sfb`: `torch.uint8`, CUTLASS SFB layout for `(N, K)`.
+- output: `torch.bfloat16`, shape `(M, N)`.
+- `K` must be divisible by 16.
+- Targets: Blackwell `sm_110a` (Jetson AGX Thor, CUDA 13+) and `sm_120a`
+  (RTX Blackwell, CUDA 12.8+).
+
+`variant` selects the CUTLASS schedule:
+
+- `-1`: architecture-aware auto-dispatch (public default).
+- `0`: default `<128,128,256>` cooperative schedule.
+- `1`: widen `<128,256,128>` schedule, intended for very large `N`.
+- `2`: pingpong schedule for A/B testing shape-specific wins.
+
+The canonical linear API and FP4/SFA quantize/dequantize helpers are available
+on both SM110 and SM120. SM110 additionally provides the GROOT N1.7 production
+epilogues `nvfp4_gemm_bias_bf16`, `nvfp4_gemm_bias_residual_bf16`, and
+`nvfp4_gemm_bias_gelu_nvfp4`. The latter emits packed FP4 plus CUTLASS SFA so
+the following projection can consume it without a BF16 materialization and a
+standalone quantization launch. Stream-K and the older BF16 GELU epilogue keep
+their existing SM120 dispatch and reject unsupported architectures explicitly.
+
+The SM110 release gate includes the production `(M,N,K)` shapes
+`(41,4608,1536)`, `(41,6144,1536)`, and `(41,1536,6144)`, plus the legacy
+`M=51` compatibility row. The kernels are the native sources used by FlashRT's
+GROOT N1.7 Thor NVFP4 pipeline.
+
+## Minimal Usage
+
+```python
+from kernels import get_kernel
+import torch
+
+ops = get_kernel("flashrt/fp4-gemm", version=1, trust_remote_code=True)
+
+x = torch.randn((32, 256), device="cuda", dtype=torch.float16)
+w = torch.randn((512, 256), device="cuda", dtype=torch.float16)
+
+a_packed, sfa = ops.quantize_fp4_sfa_fp16(x, is_sfb=False)
+b_packed, sfb = ops.quantize_fp4_sfa_fp16(w, is_sfb=True)
+
+y = ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0)
+```
+
+For BF16 model activations, use the direct producer so the hot path does not
+materialize an intermediate FP16 tensor:
+
+```python
+x_bf16 = torch.randn((1, 5120), device="cuda", dtype=torch.bfloat16)
+a_packed, sfa = ops.quantize_fp4_sfa_bf16(x_bf16)
+```
+
+The BF16 entry writes the same E2M1 bytes and CUTLASS SFA/SFB layout as
+`quantize_fp4_sfa_fp16(x_bf16.to(torch.float16))` for finite FP16-range
+inputs. It is an additive API; the existing FP16 producer remains unchanged.
+
+The quantize/dequantize helpers are included for examples and validation. A
+production runtime should keep weights prepacked and should avoid quantizing in
+the hot path unless that producer kernel is part of the intended low-bit block.
+
+Use the bias/GELU and residual variants to avoid returning to BF16
+elementwise code between low-bit GEMMs. Stream-K variants are selected only
+for the validated large down-projection shapes; unsupported shapes reject
+rather than silently selecting a losing schedule.
+
+## Validation
+
+```bash
+python fp4-gemm/tests/test_fp4_gemm.py --backend source --mode full
+python fp4-gemm/tests/test_fp4_gemm.py --backend installed --mode full \
+  --artifact fp4-gemm/build/torch211-cxx11-cu128-x86_64-linux
+python fp4-gemm/benchmarks/benchmark.py --backend installed --mode headline \
+  --artifact fp4-gemm/build/torch211-cxx11-cu128-x86_64-linux
+
+# Thor model-shape gate
+python fp4-gemm/tests/test_fp4_gemm.py --backend installed \
+  --mode thor-models \
+  --artifact fp4-gemm/build/torch211-cxx11-cu130-aarch64-linux
+```
+
+The correctness reference dequantizes the same FP4/SFA and FP4/SFB inputs used
+by the kernel, then computes the PyTorch GEMM reference from those dequantized
+low-bit values.
+
+The producer gate also checks the BF16 direct entry byte-for-byte against the
+established FP16 compatibility chain at decode widths 5120, 6144 and 17408,
+plus multi-row activation and SFB layouts.

SYNC.md

@@ -0,0 +1,49 @@
+# Source Sync
+
+- Upstream FlashRT source: `../official/FlashRT`
+- Original SM110 sync commit: `132049d7c3a3534fb7d35676cd726f39408b1af6`
+- GROOT N1.7 fused-epilogue sync commit:
+  `24df793f4fa2d50780aea03b644208c6e0cb4162`
+- Initial package date: June 20, 2026
+
+Copied source files:
+
+- `csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm120.cu/.cuh`
+- `csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cu/.cuh`
+- `csrc/gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cu/.cuh`
+- `csrc/quantize/quantize_fp4_sfa.cu/.cuh`
+- `csrc/quantize/quantize_fp4_sfa_bf16.cu/.cuh`
+- `cutlass/util/packed_stride.hpp`, copied from CUTLASS tools util headers
+  into `csrc/cutlass/util/packed_stride.hpp` so the Hub package does not
+  depend on a local `third_party/cutlass/tools/util/include` path.
+
+Packaging helper:
+
+- `csrc/dequantize_fp4_sfa.cu/.cuh` derived from the SFA dequant validation
+  helper used in `fp4-fused-ops`; this package adds `is_sfb` support so tests
+  can dequant both A/SFA and B/SFB.
+
+Local packaging edits:
+
+- Added Tensor-facing PyTorch custom ops in `torch-ext/torch_binding.cpp`.
+- Added Python wrappers and fake registrations in `torch-ext/fp4_gemm`.
+- Added the BF16 direct SFA/SFB producer as an input-type specialization of
+  the existing FP16 producer. Its E2M1 encoding and CUTLASS scale layout are
+  unchanged; the additive entry removes a standalone activation cast.
+- Public APIs accept CUDA tensors only; no raw pointers or stream arguments.
+- CUTLASS SM100/SM120 block-scaled layout support is treated as package scope,
+  not as a test-only compiler define.
+- The Tensor binding dispatches the canonical BF16-output GEMM by runtime
+  compute capability. CUDA 12.8 artifacts link only the SM120 implementation;
+  CUDA 13 artifacts link both SM110 and SM120 implementations.
+- SM110 production auto-dispatch was tiled against PI0.5, GROOT, Cosmos Edge,
+  and LingBot VLA projection shapes. Explicit schedule IDs remain diagnostic.
+
+Architecture limits:
+
+- The canonical BF16-output GEMM and SFA/SFB helpers support SM110 and SM120.
+- Bias, residual, and bias/GELU-to-FP4 epilogues have an independent SM110
+  backend copied from the production GROOT N1.7 path. Stream-K remains an
+  SM120-only API and rejects on SM110.
+- SM110 requires CUDA 13 and the package's pinned CUTLASS 4.4 target; SM120
+  requires CUDA 12.8 and CUTLASS 4.0.

VALIDATION.md

@@ -0,0 +1,83 @@
+# Validation
+
+Local source validation covers NVIDIA GeForce RTX 5090 (SM120) and NVIDIA
+Jetson AGX Thor (SM110).
+
+```bash
+python fp4-gemm/tests/test_fp4_gemm.py \
+  --backend source \
+  --mode full \
+  --json-out internal-tests/fp4-gemm-source-full.json
+```
+
+Result:
+
+- SM120 full gate: `25/25` checks passed, including all fused epilogues and
+  the aggregate BF16 direct-producer layout gate.
+- SM110 model-shape gate: `24/24` checks passed across PI0.5, GROOT, Cosmos
+  Edge, and LingBot VLA projection shapes.
+- Variants `0`, `1`, and `2` were checked.
+- SM110 additionally checks production auto-dispatch (`variant=-1`).
+- `nvfp4_gemm_bf16` is the canonical public API.
+- Correctness reference dequantizes the same FP4/SFA and FP4/SFB inputs used
+  by the kernel, then computes PyTorch GEMM on those dequantized values.
+- The direct BF16 producer is byte-exact against the established
+  BF16-to-FP16 plus FP16-producer contract for packed E2M1, mapped SFA/SFB
+  bytes, and dequantized output. Covered activation shapes are `(1,5120)`,
+  `(1,6144)`, `(1,17408)`, `(16,2048)`, and `(128,512)`; SFB coverage uses
+  `(64,1024)`.
+
+| Shape | Variant | Max abs | Mean abs | P99 abs | Cosine |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| M=16, N=128, K=128 | 0 | 0.0 | 0.0 | 0.0 | 1.0 |
+| M=16, N=128, K=128 | 1 | 0.0 | 0.0 | 0.0 | 1.0 |
+| M=16, N=128, K=128 | 2 | 0.0 | 0.0 | 0.0 | 1.0 |
+| M=32, N=256, K=256 | 0 | 0.0 | 0.0 | 0.0 | 1.0 |
+| M=32, N=256, K=256 | 1 | 0.0 | 0.0 | 0.0 | 1.0 |
+| M=32, N=256, K=256 | 2 | 0.0 | 0.0 | 0.0 | 1.0 |
+| M=64, N=512, K=512 | 0 | 0.0 | 0.0 | 0.0 | 1.0 |
+| M=64, N=512, K=512 | 1 | 0.0 | 0.0 | 0.0 | 1.0 |
+| M=64, N=512, K=512 | 2 | 0.0 | 0.0 | 0.0 | 1.0 |
+
+## Installed Artifact Validation
+
+The local kernel-builder release candidate produced and passed ABI, manylinux,
+layout, and builder `get_kernel` checks for:
+
+- `torch211-cxx11-cu128-x86_64-linux`
+- `torch211-cxx11-cu130-x86_64-linux`
+- `torch212-cxx11-cu130-x86_64-linux`
+- `torch212-cxx11-cu132-x86_64-linux`
+
+The cu128/Torch 2.11 artifact passed `10/10` runtime gates: all nine
+shape/variant correctness rows were exact against the staged reference, and
+the public `nvfp4_gemm_bf16` wrapper was exact under
+`torch.compile(fullgraph=True)`.
+
+The SM110 release flake pins kernel-builder commit
+`d720fa90fb9cd92d1bc60a9dc5c55bef2aafabb8`, which includes CUTLASS 4.5
+support and the corrected CUTLASS 4.5.2 fixed-output hash. HF Jobs, the
+SM110 aarch64 artifact build, and cold Hub loads must pass before the rebuilt
+Hub release is considered complete.
+
+## BF16 Direct Producer
+
+RTX 5090, 100 warmup iterations and 1000 measured iterations:
+
+| Shape | BF16 direct | BF16 cast + FP16 producer | Speedup | Native BF16 producer | Hub/native |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| M=1, K=5120 | 4.098 us | 6.404 us | 1.563x | 6.150 us | 0.666x |
+| M=1, K=6144 | 4.098 us | 6.403 us | 1.562x | 8.190 us | 0.500x |
+| M=1, K=17408 | 4.096 us | 6.413 us | 1.566x | 18.442 us | 0.222x |
+
+The native BF16 producer is included as a latency comparison but uses a
+different FlashRT quantization strategy. Correctness acceptance is therefore
+against this package's established FP16 producer contract, where all tested
+packed and mapped scale bytes are exact.
+
+## Thor Native Parity
+
+The Tensor wrapper was compared against the same native FlashRT launchers on
+Thor with 20 warmup and 100 measured iterations. For production auto-dispatch
+across the six model shapes, wrapper/native latency ratio had median `1.019`
+and maximum `1.086`. Correctness was exact (`max_abs=mean_abs=p99_abs=0`).

benchmarks/RESULTS.md

@@ -0,0 +1,57 @@
+# fp4-gemm Benchmark Results
+
+Installed kernel-builder artifact benchmark on NVIDIA GeForce RTX 5090,
+PyTorch `2.11.0+cu128`.
+
+Command:
+
+```bash
+python fp4-gemm/benchmarks/benchmark.py \
+  --backend installed \
+  --artifact fp4-gemm/build/torch211-cxx11-cu128-x86_64-linux \
+  --mode headline \
+  --warmup 100 \
+  --iterations 500 \
+  --json-out internal-tests/fp4-gemm-installed-benchmark.json
+```
+
+Reference is PyTorch GEMM over the same dequantized FP4/SFA and FP4/SFB inputs
+that the FlashRT kernel consumes.
+
+| Shape | Variant | FlashRT us | Eager us | Compile us | vs eager | vs compile | Max abs |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| M=16, N=128, K=128 | 0 | 6.156 | 15.156 | 27.748 | 2.46x | 4.51x | 0.0 |
+| M=16, N=128, K=128 | 1 | 6.152 | 15.156 | 27.748 | 2.46x | 4.51x | 0.0 |
+| M=16, N=128, K=128 | 2 | 6.145 | 15.156 | 27.748 | 2.47x | 4.52x | 0.0 |
+| M=32, N=256, K=256 | 0 | 6.153 | 16.685 | 35.690 | 2.71x | 5.80x | 0.0 |
+| M=32, N=256, K=256 | 1 | 8.201 | 16.685 | 35.690 | 2.03x | 4.35x | 0.0 |
+| M=32, N=256, K=256 | 2 | 6.147 | 16.685 | 35.690 | 2.71x | 5.81x | 0.0 |
+| M=64, N=512, K=512 | 0 | 6.152 | 16.480 | 36.205 | 2.68x | 5.89x | 0.0 |
+| M=64, N=512, K=512 | 1 | 10.246 | 16.480 | 36.205 | 1.61x | 3.53x | 0.0 |
+| M=64, N=512, K=512 | 2 | 6.152 | 16.480 | 36.205 | 2.68x | 5.89x | 0.0 |
+
+Variant notes:
+
+- `variant=0` is the stable default.
+- `variant=1` is the widen schedule intended for very large `N`; it is not the
+  best choice for these small validation shapes.
+- `variant=2` is competitive on small shapes and remains exposed for explicit
+  A/B testing.
+
+The PyTorch references consume the same already-dequantized FP4 tensors and do
+not include quantization. The compiled reference is warmed before timing.
+
+## BF16 Direct Producer
+
+Source benchmark on RTX 5090 with 100 warmup and 1000 measured iterations:
+
+| Shape | Direct BF16 us | Cast + FP16 producer us | Speedup | Native BF16 us | Wrapper/native |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| M=1, K=5120 | 4.098 | 6.404 | 1.563x | 6.150 | 0.666x |
+| M=1, K=6144 | 4.098 | 6.403 | 1.562x | 8.190 | 0.500x |
+| M=1, K=17408 | 4.096 | 6.413 | 1.566x | 18.442 | 0.222x |
+
+The direct entry is byte-exact against the package's established
+BF16-to-FP16 plus FP16-producer contract. The native timing is reported as a
+performance reference only because that producer uses a distinct quantization
+strategy.

build.toml

@@ -0,0 +1,74 @@
+[general]
+name = "fp4-gemm"
+license = "Apache-2.0"
+version = 1
+backends = ["cuda"]
+
+[general.cuda]
+minver = "12.8"
+
+[general.hub]
+repo-id = "flashrt/fp4-gemm"
+
+[torch]
+include = ["csrc"]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h",
+]
+
+[kernel.fp4_gemm_sm110]
+backend = "cuda"
+depends = ["torch", "cutlass_4_4"]
+include = ["csrc"]
+cuda-minver = "13"
+cuda-capabilities = ["11.0a"]
+cuda-flags = [
+  "--expt-relaxed-constexpr",
+  "-O3",
+  "--use_fast_math",
+  "-DCUTLASS_ARCH_MMA_SM100_SUPPORTED=1",
+]
+src = [
+  "csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cu",
+  "csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cuh",
+  "csrc/gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cu",
+  "csrc/gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cuh",
+  "csrc/quantize/quantize_fp4_sfa_bf16.cu",
+  "csrc/quantize/quantize_fp4_sfa_bf16.cuh",
+  "csrc/gemm/fp4/sm110_dispatch.cu",
+  "csrc/gemm/fp4/sm110_dispatch.cuh",
+]
+
+[kernel.fp4_gemm_common]
+backend = "cuda"
+depends = ["torch"]
+include = ["csrc"]
+cuda-minver = "12.8"
+cuda-capabilities = ["11.0a", "12.0a"]
+cuda-flags = ["--expt-relaxed-constexpr", "-O3", "--use_fast_math"]
+src = [
+  "csrc/quantize/quantize_fp4_sfa.cu",
+  "csrc/quantize/quantize_fp4_sfa.cuh",
+  "csrc/dequantize_fp4_sfa.cu",
+  "csrc/dequantize_fp4_sfa.cuh",
+]
+
+[kernel.fp4_gemm_sm120]
+backend = "cuda"
+depends = ["torch", "cutlass_4_0"]
+include = ["csrc"]
+cuda-minver = "12.8"
+cuda-capabilities = ["12.0a"]
+src = [
+  "csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm120.cu",
+  "csrc/gemm/fp4/fp4_w4a4_mma_warpsplit_sm120.cu",
+  "csrc/gemm/fp4/fp4_w4a4_mma_warpsplit_sm120.cuh",
+  "csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm120.cuh",
+  "csrc/gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_bf16out_sm120.cu",
+  "csrc/gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_bf16out_sm120.cuh",
+  "csrc/gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_fp4out_sm120.cu",
+  "csrc/gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_fp4out_sm120.cuh",
+  "csrc/gemm/fp4/cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cu",
+  "csrc/gemm/fp4/cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cuh",
+]

build/torch211-cxx11-cu130-aarch64-linux/__init__.py

@@ -58,6 +58,16 @@ def _legacy_linear_fake(
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_bf16"))
+def _bias_fake(a, b, sfa, sfb, bias, out) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_residual_bf16"))
+def _bias_residual_fake(a, b, sfa, sfb, bias, residual, out) -> None:
+    return None
+
+
 @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
 def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
     return None
@@ -191,6 +201,45 @@ def fp4_w4a16_linear_bf16(
     )
 
 
+def nvfp4_gemm_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 NVFP4 GEMM with a fused per-column BF16 bias."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_bf16(a_packed, b_packed, sfa, sfb, bias, out)
+    return out
+
+
+def nvfp4_gemm_bias_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 NVFP4 GEMM with fused BF16 bias and residual add."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_bias_residual_bf16(
+        a_packed, b_packed, sfa, sfb, bias, residual, out
+    )
+    return out
+
+
 def nvfp4_gemm_residual_bf16(
     a_packed: torch.Tensor,
     b_packed: torch.Tensor,
@@ -296,8 +345,10 @@ __all__ = [
     "fp4_w4a16_linear_bf16",
     "fp4_w4a4_gemv_warpsplit_bf16",
     "nvfp4_gemm_bf16",
+    "nvfp4_gemm_bias_bf16",
     "nvfp4_gemm_bias_gelu_bf16",
     "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_bias_residual_bf16",
     "nvfp4_gemm_residual_bf16",
     "nvfp4_gemm_streamk_bf16",
     "nvfp4_gemm_streamk_bias_bf16",

build/torch211-cxx11-cu130-aarch64-linux/{_fp4_gemm_cuda_b46a817.abi3.so → _fp4_gemm_cuda_8a66d8b.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:49cad7b4b7a236d14238c99a832bdc106c5e8494c8ce89e000c3bc75b45cc061
-size 751384
+oid sha256:35bd2043a1a1880f362954fe72a22651017d30a585d38ed5d2ffe86de286b45a
+size 1118808

build/torch211-cxx11-cu130-aarch64-linux/_ops.py

@@ -1,6 +1,6 @@
 import torch
-from . import _fp4_gemm_cuda_b46a817
-ops = torch.ops._fp4_gemm_cuda_b46a817
+from . import _fp4_gemm_cuda_8a66d8b
+ops = torch.ops._fp4_gemm_cuda_8a66d8b
 
 def add_op_namespace_prefix(op_name: str):
-    return f"_fp4_gemm_cuda_b46a817::{op_name}"
+    return f"_fp4_gemm_cuda_8a66d8b::{op_name}"

build/torch211-cxx11-cu130-aarch64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "fp4-gemm",
-  "id": "_fp4_gemm_cuda_b46a817",
+  "id": "_fp4_gemm_cuda_8a66d8b",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,20 +13,16 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "pJlGLDlOKnju+m59WyZJjpynMTvkPQTF6w+z4YfCwSY=",
-      "_fp4_gemm_cuda_b46a817.abi3.so": "ScrXtLeiNtFCOMmagyvcEGxehJTIzongAMO8dbRcwGE=",
-      "_ops.py": "PqcxyYzP4Co0y9NUZPOomwASQWuPPNn5QjtKrfTiZ2k=",
+      "__init__.py": "+Kk/VNnIWwe9nszQYvL2qtU3DpKxbqHVfklg5WtVuB4=",
+      "_fp4_gemm_cuda_8a66d8b.abi3.so": "Nb0gQ6GhiA82KVT+cqImUQF9MKWF047V0v/obeKGtFo=",
+      "_ops.py": "/hna1MwmSjR4pewWf12X5smsDRcaiHuoR2Gri/PrBFc=",
       "fp4_gemm/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "b46a81771fdabe26f2c2494dc2ecf8492c88c45a",
+      "sha": "8a66d8b",
       "dirty": false
-    },
-    "validation": {
-      "torch": "2.11.0+cu130",
-      "cuda": "13.0"
     }
   }
 }

csrc/cutlass/util/packed_stride.hpp

@@ -0,0 +1,570 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Utilities for packing constructing canonical CuTe stride types for 3.x mainloop params.
+*/
+
+#pragma once
+
+#include "cute/layout.hpp"
+#include "cute/container/array.hpp"   // cute::array
+#include "cutlass/conv/convolution.h" // cutlass::conv::Operator
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Strides without batch mode
+
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<IntT, cute::Int<1>>
+make_cute_packed_stride(cute::Stride<IntT, cute::Int<1>> s, cute::Shape<int,int,int> shape_MKL) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  auto s_copy = s;
+  cute::get<0>(s_copy) = static_cast<IntT>(cute::get<1>(shape_MKL));
+  return s_copy;
+}
+
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Int<1>, IntT>
+make_cute_packed_stride(cute::Stride<cute::Int<1>, IntT> s, cute::Shape<int,int,int> shape_MKL) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  auto s_copy = s;
+  cute::get<1>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL));
+  return s_copy;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Strides with batch mode
+
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<IntT, cute::Int<1>, int64_t>
+make_cute_packed_stride(cute::Stride<IntT, cute::Int<1>, int64_t> s, cute::Shape<int,int,int> shape_MKL) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  auto s_copy = s;
+  cute::get<0>(s_copy) = static_cast<IntT>(cute::get<1>(shape_MKL));
+  int batch_count =  cute::get<2>(shape_MKL);
+  if (batch_count > 1) {
+    cute::get<2>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL) * cute::get<1>(shape_MKL));
+  }
+  else {
+    cute::get<2>(s_copy) = static_cast<IntT>(0);
+  }
+  return s_copy;
+}
+
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Int<1>, IntT, int64_t>
+make_cute_packed_stride(cute::Stride<cute::Int<1>, IntT, int64_t> s, cute::Shape<int,int,int> shape_MKL) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  auto s_copy = s;
+  cute::get<1>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL));
+  int batch_count =  cute::get<2>(shape_MKL);
+  if (batch_count > 1) {
+    cute::get<2>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL) * cute::get<1>(shape_MKL));
+  }
+  else {
+    cute::get<2>(s_copy) = static_cast<IntT>(0);
+  }
+  return s_copy;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Strides with group mode
+
+template <class StrideIntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<StrideIntT, cute::Int<1>, cute::Int<0>>
+make_cute_packed_stride(cute::Stride<StrideIntT, cute::Int<1>, cute::Int<0>> s, cute::Shape<int,int,int> shape_MKL) {
+  static_assert(std::is_integral_v<StrideIntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  auto s_copy = s;
+  cute::get<0>(s_copy) = static_cast<StrideIntT>(cute::get<1>(shape_MKL));
+  return s_copy;
+}
+
+template <class StrideIntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Int<1>, StrideIntT, cute::Int<0>>
+make_cute_packed_stride(cute::Stride<cute::Int<1>, StrideIntT, cute::Int<0>> s, cute::Shape<int,int,int> shape_MKL) {
+  static_assert(std::is_integral_v<StrideIntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  auto s_copy = s;
+  cute::get<1>(s_copy) = static_cast<StrideIntT>(cute::get<0>(shape_MKL));
+  return s_copy;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Strides for convolutions
+
+// Output cutlass::layout::TensorNDHWC -> rank-3 stride (InT,_1,_0)
+// Note: For fprop/dgrad kernel, strides are assumed to be layout right in NZPQK/NDHWC order
+// and therefore can be coalesced to just q/w. For wgrad kernel, strides are assumed to be layout
+// right in KTRSC order and can be coalesced to just k.
+// We enforce this condition here with asserts.
+template <class IntT, size_t RankT_>
+CUTLASS_HOST_DEVICE
+cute::Stride<IntT, cute::Int<1>, cute::Int<0>>
+make_cute_packed_stride(
+    cute::Stride<IntT, cute::Int<1>, cute::Int<0>> s,
+    cute::array<int32_t, RankT_> shape_output,
+    cute::array<IntT, RankT_> stride_output,
+    cutlass::conv::Operator conv_op) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  static_assert(RankT_ >= 3u);
+  constexpr static int RankT = static_cast<int>(RankT_);
+
+  assert(stride_output[RankT-1] == 1);
+  cute::for_each(cute::make_seq<RankT-2>{}, [&](auto i) {
+    assert(stride_output[i] == shape_output[i+1] * stride_output[i+1]);
+  });
+
+  auto s_copy = s;
+  cute::get<0>(s_copy) = (conv_op == cutlass::conv::Operator::kWgrad) ?
+      stride_output[0] :
+      stride_output[RankT-2];
+  return s_copy;
+}
+
+//
+// Activation tensor ((w, h, d, n), _1) for fprop kernel
+//
+
+// Activation cutlass::layout::TensorNWC -> rank-2 stride ((W,N),_1)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Stride<IntT, IntT>, cute::Int<1>>
+make_cute_packed_stride(
+    cute::Stride<cute::Stride<IntT, IntT>, cute::Int<1>> s,
+    cute::array<IntT, 3> stride_nwc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  assert(stride_nwc[2] == 1);
+  auto s_copy = s;
+  cute::get<0,0>(s_copy) = stride_nwc[1];
+  cute::get<0,1>(s_copy) = stride_nwc[0];
+  return s_copy;
+}
+
+// Activation cutlass::layout::TensorNHWC -> rank-2 stride ((W,H,N),_1)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Stride<IntT, IntT, IntT>, cute::Int<1>>
+make_cute_packed_stride(
+    cute::Stride<cute::Stride<IntT, IntT, IntT>, cute::Int<1>> s,
+    cute::array<IntT, 4> stride_nhwc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  assert(stride_nhwc[3] == 1);
+  auto s_copy = s;
+  cute::for_each(cute::make_seq<3>{}, [&](auto i) {
+    cute::get<0,i>(s_copy) = stride_nhwc[2-i];
+  });
+  return s_copy;
+}
+
+// Activation cutlass::layout::TensorNDHWC -> rank-2 stride ((W,H,D,N),_1)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Stride<IntT, IntT, IntT, IntT>, cute::Int<1>>
+make_cute_packed_stride(
+    cute::Stride<cute::Stride<IntT, IntT, IntT, IntT>, cute::Int<1>> s,
+    cute::array<IntT, 5> stride_ndhwc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_ndhwc[4] == 1);
+  auto s_copy = s;
+  cute::for_each(cute::make_seq<4>{}, [&](auto i) {
+    cute::get<0,i>(s_copy) = stride_ndhwc[3-i];
+  });
+  return s_copy;
+}
+
+//
+// Filter tensor (k, (_1, s, r, t)) for fprop kernel
+//
+
+// Filter cutlass::layout::TensorNWC -> rank-2 stride (k, (_1, s))
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>>
+make_cute_packed_stride(
+    cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>> s,
+    cute::array<IntT, 3> stride_ksc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_ksc[2] == 1);
+  auto s_copy = s;
+  cute::get<0,0>(s_copy) = stride_ksc[0];
+  cute::get<1,1>(s_copy) = stride_ksc[1];
+  return s_copy;
+}
+
+// Filter cutlass::layout::TensorNHWC -> rank-2 stride (k, (_1, s, r))
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>>
+make_cute_packed_stride(
+    cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>> s,
+    cute::array<IntT, 4> stride_krsc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_krsc[3] == 1);
+  auto s_copy = s;
+  cute::get<0,0>(s_copy) = stride_krsc[0];
+  cute::for_each(cute::make_seq<2>{}, [&](auto i) {
+    cute::get<1,2-i>(s_copy) = stride_krsc[i+1];
+  });
+  return s_copy;
+}
+
+// Filter cutlass::layout::TensorNDHWC -> rank-2 stride (k, (_1, s, r, t))
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>>
+make_cute_packed_stride(
+    cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>> s,
+    cute::array<IntT, 5> stride_ktrsc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_ktrsc[4] == 1);
+  auto s_copy = s;
+  cute::get<0,0>(s_copy) = stride_ktrsc[0];
+  cute::for_each(cute::make_seq<3>{}, [&](auto i) {
+    cute::get<1,3-i>(s_copy) = stride_ktrsc[i+1];
+  });
+  return s_copy;
+}
+
+//
+// Activation tensor (_1, (w, h, d, n)) for wgrad kernel
+//
+// It is also Filter tensor ((_1), (k, s, r, t)) for dgrad kernel
+//
+
+// Activation cutlass::layout::TensorNWC -> rank-2 stride (_1, (W,N)) in wgrad
+// Filter cutlass::layout::TensorNWC -> rank-2 stride ((_1), (k, s)) in dgrad
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT>>
+make_cute_packed_stride(
+    cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT>> s,
+    cute::array<IntT, 3> stride_nwc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_nwc[2] == 1);
+  auto s_copy = s;
+  if (ConvOp == cutlass::conv::Operator::kWgrad) {
+    cute::get<1,0>(s_copy) = stride_nwc[1];
+    cute::get<1,1>(s_copy) = stride_nwc[0];
+  }
+  else if (ConvOp == cutlass::conv::Operator::kDgrad) {
+    // stride_nwc in dgrad is ksc.
+    cute::get<1,0>(s_copy) = stride_nwc[0];
+    cute::get<1,1>(s_copy) = stride_nwc[1];
+  }
+  return s_copy;
+}
+
+// Activation cutlass::layout::TensorNHWC -> rank-2 stride (_1, (W,H,N)) in wgrad
+// Filter cutlass::layout::TensorNHWC -> rank-2 stride ((_1), (k, s, r)) in dgrad
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT>>
+make_cute_packed_stride(
+    cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT>> s,
+    cute::array<IntT, 4> stride_nhwc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_nhwc[3] == 1);
+  auto s_copy = s;
+  if (ConvOp == cutlass::conv::Operator::kWgrad) {
+    cute::for_each(cute::make_seq<3>{}, [&](auto i) {
+      cute::get<1,i>(s_copy) = stride_nhwc[2-i];
+    });
+  }
+  else if (ConvOp == cutlass::conv::Operator::kDgrad) {
+    // stride_nhwc in dgrad is krsc.
+    cute::get<1,0>(s_copy) = stride_nhwc[0];
+    cute::for_each(cute::make_seq<2>{}, [&](auto i) {
+      cute::get<1,2-i>(s_copy) = stride_nhwc[i+1];
+    });
+  }
+  return s_copy;
+}
+
+// Activation cutlass::layout::TensorNDHWC -> rank-2 stride (_1, (W,H,D,N)) in wgrad
+// Filter cutlass::layout::TensorNDHWC -> rank-2 stride ((_1), (k, s, r, t)) in dgrad
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT, IntT>>
+make_cute_packed_stride(
+    cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT, IntT>> s,
+    cute::array<IntT, 5> stride_ndhwc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_ndhwc[4] == 1);
+  auto s_copy = s;
+  if (ConvOp == cutlass::conv::Operator::kWgrad) {
+    cute::for_each(cute::make_seq<4>{}, [&](auto i) {
+      cute::get<1,i>(s_copy) = stride_ndhwc[3-i];
+    });
+  }
+  else if (ConvOp == cutlass::conv::Operator::kDgrad) {
+    // stride_ndhwc in dgrad is ktrsc.
+    cute::get<1,0>(s_copy) = stride_ndhwc[0];
+    cute::for_each(cute::make_seq<3>{}, [&](auto i) {
+      cute::get<1,3-i>(s_copy) = stride_ndhwc[i+1];
+    });
+  }
+  return s_copy;
+}
+
+//
+// NZPQ tensor (_1, nzpq) for wgrad kernel
+//
+
+// cutlass::layout::TensorNWC -> rank-2 stride (_1, nzpq)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Int<1>, IntT>
+make_cute_packed_stride(
+    cute::Stride<cute::Int<1>, IntT> s,
+    cute::array<IntT, 3> stride_nqk,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_nqk[2] == 1);
+  auto s_copy = s;
+  cute::get<1>(s_copy) = stride_nqk[1];
+  return s_copy;
+}
+
+// cutlass::layout::TensorNHWC -> rank-2 stride (_1, nzpq)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Int<1>, IntT>
+make_cute_packed_stride(
+    cute::Stride<cute::Int<1>, IntT> s,
+    cute::array<IntT, 4> stride_npqk,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_npqk[3] == 1);
+  auto s_copy = s;
+  cute::get<1>(s_copy) = stride_npqk[2];
+  return s_copy;
+}
+
+// cutlass::layout::TensorNDHWC -> rank-2 stride (_1, nzpq)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Int<1>, IntT>
+make_cute_packed_stride(
+    cute::Stride<cute::Int<1>, IntT> s,
+    cute::array<IntT, 5> stride_nzpqk,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_nzpqk[4] == 1);
+  auto s_copy = s;
+  cute::get<1>(s_copy) = stride_nzpqk[3];
+  return s_copy;
+}
+
+
+
+//
+// Wgrad output tensor (k, (_1, s, r, t), _0)
+//
+
+// Filter cutlass::layout::TensorKCS -> rank-3 stride (k, (_1, s), _0)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>, cute::Int<0>>
+make_cute_packed_stride(
+    cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>, cute::Int<0>> s,
+    [[maybe_unused]] cute::array<int32_t, 3> shape_output,
+    cute::array<IntT, 3> stride_ksc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_ksc[2] == 1);
+  auto s_copy = s;
+  cute::get<0,0>(s_copy) = stride_ksc[0];
+  cute::get<1,1>(s_copy) = stride_ksc[1];
+  return s_copy;
+}
+
+// Filter cutlass::layout::TensorKCSR -> rank-3 stride (k, (_1, s, r), _0)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>, cute::Int<0>>
+make_cute_packed_stride(
+    cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>, cute::Int<0>> s,
+    [[maybe_unused]] cute::array<int32_t, 4> shape_output,
+    cute::array<IntT, 4> stride_krsc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_krsc[3] == 1);
+  auto s_copy = s;
+  cute::get<0,0>(s_copy) = stride_krsc[0];
+  cute::for_each(cute::make_seq<2>{}, [&](auto i) {
+    cute::get<1,2-i>(s_copy) = stride_krsc[i+1];
+  });
+  return s_copy;
+}
+
+// Filter cutlass::layout::TensorKCSRT -> rank-3 stride (k, (_1, s, r, t), _0)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>, cute::Int<0>>
+make_cute_packed_stride(
+    cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>, cute::Int<0>> s,
+    [[maybe_unused]] cute::array<int32_t, 5> shape_output,
+    cute::array<IntT, 5> stride_ktrsc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_ktrsc[4] == 1);
+  auto s_copy = s;
+  cute::get<0,0>(s_copy) = stride_ktrsc[0];
+  cute::for_each(cute::make_seq<3>{}, [&](auto i) {
+    cute::get<1,3-i>(s_copy) = stride_ktrsc[i+1];
+  });
+  return s_copy;
+}
+
+
+//
+// Wgrad output tensor ((_1, s, r, t), k, _0)
+//
+
+// Filter cutlass::layout::TensorCSK -> rank-3 stride ((_1, s), k, _0)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Stride<cute::Int<1>, IntT>, IntT, cute::Int<0>>
+make_cute_packed_stride(
+    cute::Stride<cute::Stride<cute::Int<1>, IntT>, IntT, cute::Int<0>> s,
+    [[maybe_unused]] cute::array<int32_t, 3> shape_output,
+    cute::array<IntT, 3> stride_ksc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_ksc[2] == 1);
+  auto s_copy = s;
+  cute::get<1,0>(s_copy) = stride_ksc[0];
+  cute::get<0,1>(s_copy) = stride_ksc[1];
+  return s_copy;
+}
+
+// Filter cutlass::layout::TensorCSRK -> rank-3 stride ((_1, s, r), k, _0)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT>, IntT, cute::Int<0>>
+make_cute_packed_stride(
+    cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT>, IntT, cute::Int<0>> s,
+    [[maybe_unused]] cute::array<int32_t, 4> shape_output,
+    cute::array<IntT, 4> stride_krsc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_krsc[3] == 1);
+  auto s_copy = s;
+  cute::get<1,0>(s_copy) = stride_krsc[0];
+  cute::for_each(cute::make_seq<2>{}, [&](auto i) {
+    cute::get<0,2-i>(s_copy) = stride_krsc[i+1];
+  });
+  return s_copy;
+}
+
+// Filter cutlass::layout::TensorCSRTK -> rank-3 stride ((_1, s, r, t), k, _0)
+template <class IntT>
+CUTLASS_HOST_DEVICE
+cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT, IntT>, IntT, cute::Int<0>>
+make_cute_packed_stride(
+    cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT, IntT>, IntT, cute::Int<0>> s,
+    [[maybe_unused]] cute::array<int32_t, 5> shape_output,
+    cute::array<IntT, 5> stride_ktrsc,
+    conv::Operator ConvOp) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+
+  assert(stride_ktrsc[4] == 1);
+  auto s_copy = s;
+  cute::get<1,0>(s_copy) = stride_ktrsc[0];
+  cute::for_each(cute::make_seq<3>{}, [&](auto i) {
+    cute::get<0,3-i>(s_copy) = stride_ktrsc[i+1];
+  });
+  return s_copy;
+}
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace cutlass

csrc/dequantize_fp4_sfa.cu

@@ -0,0 +1,89 @@
+// SPDX-License-Identifier: Apache-2.0
+#include "dequantize_fp4_sfa.cuh"
+
+#include <cuda_fp8.h>
+
+#ifndef CUTLASS_ARCH_MMA_SM100_SUPPORTED
+#  define CUTLASS_ARCH_MMA_SM100_SUPPORTED 1
+#endif
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(__CUDA_ARCH__)
+#  include "cutlass/cutlass.h"
+#  include "cutlass/detail/sm100_blockscaled_layout.hpp"
+#  include "cute/tensor.hpp"
+#  define FV_HAVE_CUTLASS 1
+#else
+#  define FV_HAVE_CUTLASS 0
+#endif
+
+namespace flash_rt {
+namespace fused_fp4 {
+
+#if FV_HAVE_CUTLASS
+
+using CfgDequant = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+__device__ __forceinline__ float e2m1_to_fp32_dequant(uint8_t value) {
+  static constexpr float mags[8] = {0.f, 0.5f, 1.f, 1.5f, 2.f, 3.f, 4.f, 6.f};
+  float mag = mags[value & 0x7];
+  return (value & 0x8) ? -mag : mag;
+}
+
+template <class LayoutSF>
+__global__ void dequantize_fp4_sfa_kernel(
+    const uint8_t* __restrict__ packed,
+    const uint8_t* __restrict__ sfa,
+    __half* __restrict__ out,
+    LayoutSF layout,
+    int dim) {
+  int row = blockIdx.y;
+  int block_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  int n_blocks = dim / 16;
+  if (block_idx >= n_blocks) return;
+
+  int col_base = block_idx * 16;
+  int sfa_off = layout(row, col_base, 0);
+  __nv_fp8_e4m3 scale_q;
+  *reinterpret_cast<uint8_t*>(&scale_q) = sfa[sfa_off];
+  float scale = static_cast<float>(scale_q);
+
+  const uint8_t* packed_block = packed + row * (dim / 2) + block_idx * 8;
+  __half* out_block = out + row * dim + col_base;
+#pragma unroll
+  for (int p = 0; p < 8; ++p) {
+    uint8_t byte = packed_block[p];
+    out_block[2 * p] = __float2half(e2m1_to_fp32_dequant(byte & 0xF) * scale);
+    out_block[2 * p + 1] = __float2half(e2m1_to_fp32_dequant(byte >> 4) * scale);
+  }
+}
+
+#endif
+
+void dequantize_fp4_sfa_fp16(
+    const uint8_t* packed,
+    const uint8_t* sfa,
+    __half* out,
+    int rows,
+    int dim,
+    bool is_sfb,
+    cudaStream_t stream) {
+#if FV_HAVE_CUTLASS
+  int n_blocks = dim / 16;
+  dim3 block(256);
+  dim3 grid((n_blocks + block.x - 1) / block.x, rows);
+  auto shape = cute::make_shape(is_sfb ? 1 : rows, is_sfb ? rows : 1, dim, 1);
+  if (is_sfb) {
+    auto layout = CfgDequant::tile_atom_to_shape_SFB(shape);
+    dequantize_fp4_sfa_kernel<<<grid, block, 0, stream>>>(
+        packed, sfa, out, layout, dim);
+  } else {
+    auto layout = CfgDequant::tile_atom_to_shape_SFA(shape);
+    dequantize_fp4_sfa_kernel<<<grid, block, 0, stream>>>(
+        packed, sfa, out, layout, dim);
+  }
+#else
+  (void)packed; (void)sfa; (void)out; (void)rows; (void)dim; (void)is_sfb; (void)stream;
+#endif
+}
+
+}  // namespace fused_fp4
+}  // namespace flash_rt

csrc/dequantize_fp4_sfa.cuh

@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <cstdint>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+
+namespace flash_rt {
+namespace fused_fp4 {
+
+void dequantize_fp4_sfa_fp16(
+    const uint8_t* packed,
+    const uint8_t* sfa,
+    __half* out,
+    int rows,
+    int dim,
+    bool is_sfb,
+    cudaStream_t stream);
+
+}  // namespace fused_fp4
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cu

@@ -0,0 +1,208 @@
+// ============================================================================
+//  NVFP4 GEMMs with bf16 fused-bias epilogues. See header for the contract.
+//
+//  Three fusions over one shared skinny-M mainloop config (Sm100,
+//  tile 128x64x256, cluster 1x1x1 — the narrow-N + wide-K shape that wins
+//  at small M where the GEMM is weight-bandwidth-bound):
+//    bias:      LinCombPerColBias  (bf16 out, beta = 0)
+//    bias+res:  LinCombPerColBias  (bf16 out, beta = 1, C = residual)
+//    bias+gelu: LinCombPerColBiasEltActBlockScaleFactor<GELU_taylor>
+//               (fp4 + SFA out for the following NVFP4 GEMM)
+// ============================================================================
+#include "gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cuh"
+
+#include "cutlass/cutlass.h"
+#include "cutlass/epilogue/thread/activation.h"
+#include "cutlass/epilogue/dispatch_policy.hpp"
+#include "cutlass/epilogue/fusion/operations.hpp"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+#include "cutlass/util/packed_stride.hpp"
+#include "cutlass/detail/sm100_blockscaled_layout.hpp"
+#include "cute/tensor.hpp"
+
+namespace flash_rt {
+namespace fp4 {
+
+namespace bias_bf16_gemm {
+
+using namespace cute;
+
+using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+using LayoutATag = cutlass::layout::RowMajor;
+constexpr int AlignmentA = 32;
+
+using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+using LayoutBTag = cutlass::layout::ColumnMajor;
+constexpr int AlignmentB = 32;
+
+using ElementAccumulator = float;
+using ElementCompute = float;
+using ArchTag = cutlass::arch::Sm100;
+using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
+constexpr int SFVecSize = 16;
+
+using MmaTileShape = Shape<_128, _64, _256>;
+using ClusterShape = Shape<_1, _1, _1>;
+
+template <class FusionOp, class ElemC, class ElemD, int AlignCD>
+struct BiasGemm {
+  using CollectiveEpilogue =
+      typename cutlass::epilogue::collective::CollectiveBuilder<
+          ArchTag, OperatorClass, MmaTileShape, ClusterShape,
+          cutlass::epilogue::collective::EpilogueTileAuto,
+          ElementAccumulator, ElementAccumulator,
+          ElemC, cutlass::layout::RowMajor, AlignCD,
+          ElemD, cutlass::layout::RowMajor, AlignCD,
+          cutlass::epilogue::collective::EpilogueScheduleAuto,
+          FusionOp>::CollectiveOp;
+
+  using CollectiveMainloop =
+      typename cutlass::gemm::collective::CollectiveBuilder<
+          ArchTag, OperatorClass,
+          ElementA, LayoutATag, AlignmentA,
+          ElementB, LayoutBTag, AlignmentB,
+          ElementAccumulator, MmaTileShape, ClusterShape,
+          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+              sizeof(typename CollectiveEpilogue::SharedStorage))>,
+          cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
+
+  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+      Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>;
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+};
+
+// ── bias / bias+res: bf16 out ──────────────────────────────────────────────
+using ElementCD = cutlass::bfloat16_t;
+using FusionBias = cutlass::epilogue::fusion::LinCombPerColBias<
+    ElementCD, ElementCompute, ElementCD, ElementCD, ElementCompute>;
+using GemmBias = BiasGemm<FusionBias, ElementCD, ElementCD, 8>::Gemm;
+
+// ── bias+gelu: fp4 + SFA out ───────────────────────────────────────────────
+using ElementDQ = cutlass::float_e2m1_t;
+using ElementSFD = cutlass::float_ue4m3_t;
+using FusionGelu =
+    cutlass::epilogue::fusion::LinCombPerColBiasEltActBlockScaleFactor<
+        cutlass::epilogue::thread::GELU_taylor, SFVecSize,
+        ElementDQ, ElementCompute, ElementSFD, cutlass::layout::RowMajor,
+        ElementCD, ElementDQ, ElementCompute>;
+using GemmGelu = BiasGemm<FusionGelu, ElementDQ, ElementDQ, 32>::Gemm;
+
+template <class Gemm>
+static int run_gemm(typename Gemm::Arguments& args, cudaStream_t stream) {
+  Gemm gemm;
+  auto st = gemm.can_implement(args);
+  if (st != cutlass::Status::kSuccess) return static_cast<int>(st) | 0x10000;
+  size_t ws_sz = Gemm::get_workspace_size(args);
+  void* ws = nullptr;
+  if (ws_sz > 0 && cudaMalloc(&ws, ws_sz) != cudaSuccess) return -1;
+  st = gemm.initialize(args, ws, stream);
+  if (st != cutlass::Status::kSuccess) {
+    if (ws) cudaFree(ws);
+    return static_cast<int>(st) | 0x20000;
+  }
+  st = gemm.run(stream);
+  if (ws) cudaFree(ws);
+  return (st == cutlass::Status::kSuccess) ? 0
+                                           : (static_cast<int>(st) | 0x30000);
+}
+
+template <class Gemm, class ElemC, class ElemD>
+static typename Gemm::Arguments make_args(
+    void const* A, void const* SFA, void const* B, void const* SFB,
+    void const* C, void* D, int M, int N, int K) {
+  auto stride_A = cutlass::make_cute_packed_stride(
+      typename Gemm::GemmKernel::StrideA{}, {M, K, 1});
+  auto stride_B = cutlass::make_cute_packed_stride(
+      typename Gemm::GemmKernel::StrideB{}, {N, K, 1});
+  auto stride_C = cutlass::make_cute_packed_stride(
+      typename Gemm::GemmKernel::StrideC{}, {M, N, 1});
+  auto stride_D = cutlass::make_cute_packed_stride(
+      typename Gemm::GemmKernel::StrideD{}, {M, N, 1});
+  using Cfg =
+      typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
+  auto layout_SFA = Cfg::tile_atom_to_shape_SFA(make_shape(M, N, K, 1));
+  auto layout_SFB = Cfg::tile_atom_to_shape_SFB(make_shape(M, N, K, 1));
+
+  using EA = typename ElementA::DataType;
+  using SA = typename ElementA::ScaleFactorType;
+
+  return typename Gemm::Arguments{
+      cutlass::gemm::GemmUniversalMode::kGemm, {M, N, K, 1},
+      {reinterpret_cast<EA const*>(A), stride_A,
+       reinterpret_cast<EA const*>(B), stride_B,
+       reinterpret_cast<SA const*>(SFA), layout_SFA,
+       reinterpret_cast<SA const*>(SFB), layout_SFB},
+      {{},
+       reinterpret_cast<ElemC const*>(C), stride_C,
+       reinterpret_cast<ElemD*>(D), stride_D}};
+}
+
+}  // namespace bias_bf16_gemm
+
+int cutlass_fp4_gemm_bias_bf16(
+    void const* A_packed, void const* SFA,
+    void const* B_packed, void const* SFB,
+    void const* bias_bf16,
+    void* D_bf16,
+    int M, int N, int K, cudaStream_t stream) {
+  using namespace bias_bf16_gemm;
+  auto args = make_args<GemmBias, ElementCD, ElementCD>(
+      A_packed, SFA, B_packed, SFB, D_bf16, D_bf16, M, N, K);
+  args.epilogue.thread.alpha = 1.0f;
+  args.epilogue.thread.beta = 0.0f;
+  args.epilogue.thread.bias_ptr =
+      reinterpret_cast<ElementCD const*>(bias_bf16);
+  return run_gemm<GemmBias>(args, stream);
+}
+
+int cutlass_fp4_gemm_bias_res_bf16(
+    void const* A_packed, void const* SFA,
+    void const* B_packed, void const* SFB,
+    void const* bias_bf16,
+    void const* C_bf16, void* D_bf16,
+    int M, int N, int K, cudaStream_t stream) {
+  using namespace bias_bf16_gemm;
+  auto args = make_args<GemmBias, ElementCD, ElementCD>(
+      A_packed, SFA, B_packed, SFB, C_bf16, D_bf16, M, N, K);
+  args.epilogue.thread.alpha = 1.0f;
+  args.epilogue.thread.beta = 1.0f;
+  args.epilogue.thread.bias_ptr =
+      reinterpret_cast<ElementCD const*>(bias_bf16);
+  return run_gemm<GemmBias>(args, stream);
+}
+
+int cutlass_fp4_gemm_bias_gelu_fp4out_bf16(
+    void const* A_packed, void const* SFA,
+    void const* B_packed, void const* SFB,
+    void const* bias_bf16,
+    void* D_packed, void* D_SFD,
+    int M, int N, int K, cudaStream_t stream) {
+  using namespace bias_bf16_gemm;
+  auto args = make_args<GemmGelu, ElementDQ, ElementDQ>(
+      A_packed, SFA, B_packed, SFB, D_packed, D_packed, M, N, K);
+  args.epilogue.thread.alpha = 1.0f;
+  args.epilogue.thread.beta = 0.0f;
+  args.epilogue.thread.bias_ptr =
+      reinterpret_cast<ElementCD const*>(bias_bf16);
+  // The block-scale epilogue divides by a device-resident norm constant;
+  // 1.0 keeps the native per-16 dynamic scale. Allocated once, first call
+  // must happen before any CUDA Graph capture (warmup covers this).
+  static float* d_norm = nullptr;
+  if (!d_norm) {
+    if (cudaMalloc(&d_norm, sizeof(float)) != cudaSuccess) return -1;
+    float h = 1.0f;
+    cudaMemcpyAsync(d_norm, &h, sizeof(float), cudaMemcpyHostToDevice,
+                    stream);
+  }
+  args.epilogue.thread.block_scale_factor_ptr =
+      reinterpret_cast<bias_bf16_gemm::ElementSFD*>(D_SFD);
+  args.epilogue.thread.norm_constant_ptr = d_norm;
+  return run_gemm<GemmGelu>(args, stream);
+}
+
+}  // namespace fp4
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cuh

@@ -0,0 +1,47 @@
+// ============================================================================
+//  FlashRT — NVFP4 GEMMs with bf16 fused-bias epilogues (SM100/SM110).
+//
+//  bf16 companions of the fp16 fused-epilogue NVFP4 GEMMs, for pipelines
+//  whose activations and biases are bf16 (GR00T N1.7 DiT). All three share
+//  the proven skinny-M block-scaled mainloop (tile 128x64x256, cluster
+//  1x1x1). A is row-major [M, K] packed e2m1, B is column-major [N, K]
+//  packed e2m1 (A @ B^T, nn.Linear convention); SFA/SFB use the CUTLASS
+//  Sm1xx tile-interleaved UE4M3 layout.
+//
+//  Additive: new symbols only; no existing kernel is modified.
+// ============================================================================
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace fp4 {
+
+// D_bf16[M,N] = A @ B^T + bias[N]
+int cutlass_fp4_gemm_bias_bf16(
+    void const* A_packed, void const* SFA,
+    void const* B_packed, void const* SFB,
+    void const* bias_bf16,
+    void* D_bf16,
+    int M, int N, int K, cudaStream_t stream);
+
+// D_bf16[M,N] = A @ B^T + bias[N] + C_bf16[M,N]   (residual; C may alias D)
+int cutlass_fp4_gemm_bias_res_bf16(
+    void const* A_packed, void const* SFA,
+    void const* B_packed, void const* SFB,
+    void const* bias_bf16,
+    void const* C_bf16, void* D_bf16,
+    int M, int N, int K, cudaStream_t stream);
+
+// D_fp4[M,N], SFD = blockscale(gelu_tanh(A @ B^T + bias[N]))
+// SFD is written in the SFA tile-interleaved layout over (M, N) so the
+// output can feed the K side of a following NVFP4 GEMM directly.
+int cutlass_fp4_gemm_bias_gelu_fp4out_bf16(
+    void const* A_packed, void const* SFA,
+    void const* B_packed, void const* SFB,
+    void const* bias_bf16,
+    void* D_packed, void* D_SFD,
+    int M, int N, int K, cudaStream_t stream);
+
+}  // namespace fp4
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_bf16out_sm120.cu

@@ -0,0 +1,212 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS NVFP4 W4A16 GEMM with fused per-col-bias + GELU(tanh) epilogue,
+// BF16 output, SM120a. Recipe C step 1.
+//
+// Replaces the (cutlass NVFP4 GEMM_up + bias_gelu_inplace_bf16) 2-launch
+// chain segment in the Wan FFN forward (motus). At M=360 K=3072 N=14336
+// the fused kernel ships at ~32 µs/call vs ~41 µs for the 2-launch chain
+// (1.28× standalone, ~-0.5 ms E2E wall per replay in CUDA graph mode).
+//
+// Schedule: KernelTmaWarpSpecializedPingpong + PersistentScheduler — picked
+// from the empirical sweep over {coop, pingpong} × {persistent, streamk}
+// at production shape (pingpong wins by ~0.6 µs/call).
+//
+// TileShape <128,128,256> ClusterShape <1,1,1>: locked by cutlass v4.4
+// NVFP4 sm_120 BlockScaled (all unit tests use this tile; other tiles
+// fail TMA atom constraints).
+
+#include "cutlass_nvfp4_gemm_bias_gelu_bf16out_sm120.cuh"
+
+#include "cute/tensor.hpp"
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/detail/sm100_blockscaled_layout.hpp"
+
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/thread/activation.h"
+#include "cutlass/epilogue/fusion/operations.hpp"
+
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+
+#include "cutlass/util/packed_stride.hpp"
+
+#include <cstdio>
+#include <mutex>
+#include <unordered_map>
+
+namespace flash_rt {
+namespace gemm {
+
+namespace {
+using namespace cute;
+
+using ElementA           = cutlass::float_e2m1_t;
+using ElementB           = cutlass::float_e2m1_t;
+using ElementC           = cutlass::bfloat16_t;
+using ElementD           = cutlass::bfloat16_t;
+using ElementBias        = cutlass::bfloat16_t;
+using ElementAccumulator = float;
+using ElementCompute     = float;
+using ElementSF          = cutlass::float_ue4m3_t;
+
+using LayoutA = cutlass::layout::RowMajor;
+using LayoutB = cutlass::layout::ColumnMajor;
+using LayoutC = cutlass::layout::RowMajor;
+using LayoutD = cutlass::layout::RowMajor;
+
+using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+
+constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value;
+constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value;
+constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+
+// TileShape <128,128,256>: E2E winner. Standalone bench on random inputs
+// shows Tile<128,128,128>+coop is 1.06 µs faster per call (32.84 vs 33.90)
+// but in CUDA graph mode E2E the 256-K tile + pingpong is 0.3 ms wall
+// faster — graph scheduler reshapes the cost picture vs standalone.
+using TileShape    = Shape<_128, _128, _256>;
+using ClusterShape = Shape<_1, _1, _1>;
+
+using Sm1xxBlkScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+// D = GELU_tanh(alpha * acc + per_col_bias).
+using FusionOperation = cutlass::epilogue::fusion::LinCombPerColBiasEltAct<
+    cutlass::epilogue::thread::GELU_taylor,
+    ElementD, ElementCompute, ElementBias, ElementC>;
+
+using CollectiveEpilogue =
+    typename cutlass::epilogue::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
+        TileShape, ClusterShape,
+        cutlass::epilogue::collective::EpilogueTileAuto,
+        ElementAccumulator, ElementCompute,
+        ElementC, LayoutC, AlignmentC,
+        ElementD, LayoutD, AlignmentD,
+        cutlass::epilogue::collective::EpilogueScheduleAuto,
+        FusionOperation
+    >::CollectiveOp;
+
+using CollectiveMainloop =
+    typename cutlass::gemm::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
+        ElementPairA, LayoutA, AlignmentA,
+        ElementPairB, LayoutB, AlignmentB,
+        ElementAccumulator,
+        TileShape, ClusterShape,
+        cutlass::gemm::collective::StageCountAutoCarveout<
+            static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+        cutlass::gemm::KernelTmaWarpSpecializedPingpong
+    >::CollectiveOp;
+
+using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+    Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
+    cutlass::gemm::PersistentScheduler>;
+
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+struct ShapeKey {
+  int M, N, K;
+  bool operator==(const ShapeKey& o) const {
+    return M == o.M && N == o.N && K == o.K;
+  }
+};
+struct SHash {
+  size_t operator()(const ShapeKey& k) const noexcept {
+    return (size_t(k.M) * 1315423911u) ^ (size_t(k.N) * 2654435761u)
+         ^ size_t(k.K);
+  }
+};
+struct CachedWs { void* ptr = nullptr; size_t size = 0; };
+std::unordered_map<ShapeKey, CachedWs, SHash> g_ws;
+std::mutex g_mu;
+
+void* get_ws(int M, int N, int K, size_t need) {
+  std::lock_guard<std::mutex> lk(g_mu);
+  ShapeKey k{M, N, K};
+  auto it = g_ws.find(k);
+  if (it != g_ws.end() && it->second.size >= need) return it->second.ptr;
+  if (it != g_ws.end()) { cudaFree(it->second.ptr); g_ws.erase(it); }
+  CachedWs w; w.size = need;
+  if (need > 0) cudaMalloc(&w.ptr, need);
+  g_ws[k] = w;
+  return w.ptr;
+}
+
+}  // namespace
+
+void fp4_w4a16_gemm_bias_gelu_bf16out_sm120(
+    const void* A_packed, const void* B_packed,
+    const void* SFA,      const void* SFB,
+    const void* bias_bf16,
+    void*       D_bf16,
+    int M, int N, int K,
+    float alpha,
+    cudaStream_t stream)
+{
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+  StrideA strA = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1));
+  StrideB strB = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1));
+  StrideC strC = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1));
+  StrideD strD = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1));
+  auto problem = cute::make_shape(M, N, K, 1);
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem);
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem);
+
+  using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
+  using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
+
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ArrayElementA const*>(A_packed), strA,
+          reinterpret_cast<ArrayElementB const*>(B_packed), strB,
+          reinterpret_cast<ElementSF const*>(SFA), layout_SFA,
+          reinterpret_cast<ElementSF const*>(SFB), layout_SFB
+      },
+      {
+          {alpha, 0.0f},
+          nullptr, strC,
+          reinterpret_cast<ElementD*>(D_bf16), strD
+      }
+  };
+  args.epilogue.thread.bias_ptr =
+      reinterpret_cast<ElementBias const*>(bias_bf16);
+
+  Gemm gemm;
+  size_t ws_size = Gemm::get_workspace_size(args);
+  void* ws_ptr = get_ws(M, N, K, ws_size);
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_bias_gelu_bf16out_sm120] can_implement FAIL status=%d\n",
+        int(status));
+    return;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_bias_gelu_bf16out_sm120] initialize FAIL status=%d\n",
+        int(status));
+    return;
+  }
+  status = gemm.run(stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_bias_gelu_bf16out_sm120] run FAIL status=%d\n",
+        int(status));
+  }
+}
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_bf16out_sm120.cuh

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS NVFP4 W4A16 GEMM with fused per-col bias + GELU(tanh) epilogue,
+// BF16 output, SM120a. Header for csrc/gemm/fp4/cutlass_nvfp4_gemm_bias_
+// gelu_bf16out_sm120.cu (Recipe C step 1).
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace gemm {
+
+// D = GELU_tanh(alpha * (A_fp4 @ B_fp4^T scaled by SFA/SFB) + bias_per_col)
+//
+//   A_packed : (M, K/2)  uint8   NVFP4 packed (cutlass-swizzled)
+//   B_packed : (N, K/2)  uint8   NVFP4 packed (cutlass-swizzled)
+//   SFA      : (M*K/16)  e4m3    NVFP4 SF for A
+//   SFB      : (N*K/16)  e4m3    NVFP4 SF for B
+//   bias_bf16: (N,)      bf16    per-col bias (added before GELU)
+//   D_bf16   : (M, N)    bf16    output
+//   alpha    : float32           = sf_global_a * sf_global_b
+//
+// Stream-safe; per-shape workspace cached internally.
+void fp4_w4a16_gemm_bias_gelu_bf16out_sm120(
+    const void*  A_packed,
+    const void*  B_packed,
+    const void*  SFA,
+    const void*  SFB,
+    const void*  bias_bf16,
+    void*        D_bf16,
+    int M, int N, int K,
+    float        alpha,
+    cudaStream_t stream);
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_fp4out_sm120.cu

@@ -0,0 +1,234 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS NVFP4 W4A16 GEMM with fused per-col bias + GELU(tanh) +
+// per-block-16 NVFP4 quantization epilogue, FP4 packed output, SM120a.
+//
+// Replaces the 3-launch chain
+//   cutlass NVFP4 GEMM_up    (~33 µs)
+// + bias_gelu_inplace_bf16   (~8  µs)
+// + quantize_bf16_to_nvfp4   (~20 µs)
+// with a single cutlass-fork kernel (~33 µs).
+//
+// Same TileShape <128,128,256> ClusterShape <1,1,1>
+// KernelTmaWarpSpecializedPingpong as the bf16-out fork; the
+// FusionOperation is swapped to LinCombPerColBiasEltActBlockScaleFactor
+// which produces packed NVFP4 + UE4M3 SF in cutlass-swizzled layout
+// (consumable directly by the downstream NVFP4 W4A16 GEMM_dn that reads
+// the same SF layout).
+
+#include "cutlass_nvfp4_gemm_bias_gelu_fp4out_sm120.cuh"
+
+#include "cute/tensor.hpp"
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/detail/sm100_blockscaled_layout.hpp"
+
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/thread/activation.h"
+#include "cutlass/epilogue/fusion/operations.hpp"
+
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+
+#include "cutlass/util/packed_stride.hpp"
+
+#include <cstdio>
+#include <mutex>
+#include <unordered_map>
+
+namespace flash_rt {
+namespace gemm {
+
+namespace {
+using namespace cute;
+
+using ElementA           = cutlass::float_e2m1_t;
+using ElementB           = cutlass::float_e2m1_t;
+using ElementC           = cutlass::bfloat16_t;
+using ElementD           = cutlass::float_e2m1_t;
+using ElementSFD         = cutlass::float_ue4m3_t;
+using ElementBias        = cutlass::bfloat16_t;
+using ElementAccumulator = float;
+using ElementCompute     = float;
+using ElementSF          = cutlass::float_ue4m3_t;
+
+using LayoutA      = cutlass::layout::RowMajor;
+using LayoutB      = cutlass::layout::ColumnMajor;
+using LayoutC      = cutlass::layout::ColumnMajor;
+using LayoutD      = cutlass::layout::RowMajor;
+using LayoutSFDTag = cutlass::layout::RowMajor;
+
+using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+
+constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value;
+constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value;
+constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+
+using TileShape    = Shape<_128, _128, _256>;
+using ClusterShape = Shape<_1, _1, _1>;
+using Sm1xxBlkScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+constexpr int OutputSFVectorSize = 16;
+
+using FusionOperation = cutlass::epilogue::fusion::LinCombPerColBiasEltActBlockScaleFactor<
+    cutlass::epilogue::thread::GELU_taylor,
+    OutputSFVectorSize,
+    ElementD,
+    ElementCompute,
+    ElementSFD,
+    LayoutSFDTag,
+    ElementBias>;
+
+using CollectiveEpilogue =
+    typename cutlass::epilogue::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
+        TileShape, ClusterShape,
+        cutlass::epilogue::collective::EpilogueTileAuto,
+        ElementAccumulator, ElementCompute,
+        ElementC, LayoutC, AlignmentC,
+        ElementD, LayoutD, AlignmentD,
+        cutlass::epilogue::collective::EpilogueScheduleAuto,
+        FusionOperation
+    >::CollectiveOp;
+
+using CollectiveMainloop =
+    typename cutlass::gemm::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
+        ElementPairA, LayoutA, AlignmentA,
+        ElementPairB, LayoutB, AlignmentB,
+        ElementAccumulator,
+        TileShape, ClusterShape,
+        cutlass::gemm::collective::StageCountAutoCarveout<
+            static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+        cutlass::gemm::KernelTmaWarpSpecializedPingpong
+    >::CollectiveOp;
+
+using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+    Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
+    cutlass::gemm::PersistentScheduler>;
+
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+using SfdOutputCfg = cutlass::detail::Sm1xxBlockScaledOutputConfig<OutputSFVectorSize>;
+
+struct ShapeKey {
+  int M, N, K;
+  bool operator==(const ShapeKey& o) const {
+    return M == o.M && N == o.N && K == o.K;
+  }
+};
+struct SHash {
+  size_t operator()(const ShapeKey& k) const noexcept {
+    return (size_t(k.M) * 1315423911u) ^ (size_t(k.N) * 2654435761u)
+         ^ size_t(k.K);
+  }
+};
+struct CachedWs { void* ptr = nullptr; size_t size = 0; };
+std::unordered_map<ShapeKey, CachedWs, SHash> g_ws;
+std::mutex g_mu;
+
+void* get_ws(int M, int N, int K, size_t need) {
+  std::lock_guard<std::mutex> lk(g_mu);
+  ShapeKey k{M, N, K};
+  auto it = g_ws.find(k);
+  if (it != g_ws.end() && it->second.size >= need) return it->second.ptr;
+  if (it != g_ws.end()) { cudaFree(it->second.ptr); g_ws.erase(it); }
+  CachedWs w; w.size = need;
+  if (need > 0) cudaMalloc(&w.ptr, need);
+  g_ws[k] = w;
+  return w.ptr;
+}
+
+float* get_norm_const_one() {
+  static float* p = nullptr;
+  if (p == nullptr) {
+    cudaMalloc(&p, sizeof(float));
+    float one = 1.0f;
+    cudaMemcpy(p, &one, sizeof(float), cudaMemcpyHostToDevice);
+  }
+  return p;
+}
+
+}  // namespace
+
+void fp4_w4a16_gemm_bias_gelu_fp4out_sm120(
+    const void* A_packed, const void* B_packed,
+    const void* SFA,      const void* SFB,
+    const void* bias_bf16,
+    void*       D_packed,
+    void*       SFD,
+    int M, int N, int K,
+    float alpha,
+    cudaStream_t stream)
+{
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+  StrideA strA = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1));
+  StrideB strB = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1));
+  StrideC strC = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1));
+  StrideD strD = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1));
+
+  auto problem_MNKL = cute::make_shape(M, N, K, 1);
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_MNKL);
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_MNKL);
+
+  using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
+  using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
+
+  float* norm_const_dev = get_norm_const_one();
+
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ArrayElementA const*>(A_packed), strA,
+          reinterpret_cast<ArrayElementB const*>(B_packed), strB,
+          reinterpret_cast<ElementSF const*>(SFA), layout_SFA,
+          reinterpret_cast<ElementSF const*>(SFB), layout_SFB
+      },
+      {
+          {alpha, 0.0f},
+          nullptr, strC,
+          reinterpret_cast<ElementD*>(D_packed), strD
+      }
+  };
+  args.epilogue.thread.bias_ptr =
+      reinterpret_cast<ElementBias const*>(bias_bf16);
+  args.epilogue.thread.block_scale_factor_ptr =
+      reinterpret_cast<ElementSFD*>(SFD);
+  args.epilogue.thread.norm_constant_ptr = norm_const_dev;
+
+  Gemm gemm;
+  size_t ws_size = Gemm::get_workspace_size(args);
+  void* ws_ptr = get_ws(M, N, K, ws_size);
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_bias_gelu_fp4out_sm120] can_implement FAIL "
+        "M=%d N=%d K=%d status=%d\n", M, N, K, int(status));
+    return;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_bias_gelu_fp4out_sm120] initialize FAIL "
+        "M=%d N=%d K=%d status=%d\n", M, N, K, int(status));
+    return;
+  }
+  status = gemm.run(stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_bias_gelu_fp4out_sm120] run FAIL status=%d\n",
+        int(status));
+  }
+}
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_fp4out_sm120.cuh

@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS NVFP4 W4A16 GEMM with fused per-col bias + GELU(tanh) +
+// per-block-16 NVFP4 quantization epilogue, FP4 packed output, SM120a.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace gemm {
+
+// D_packed[m, n/2] = pack_FP4(
+//   GELU_tanh(alpha * (A_fp4 @ B_fp4^T scaled by SFA/SFB) + bias_per_col)
+// ) with per-16-block NVFP4 SFD in cutlass-swizzled UE4M3 layout.
+//
+//   A_packed : (M, K/2)  uint8   NVFP4 packed (cutlass-swizzled SF)
+//   B_packed : (N, K/2)  uint8   NVFP4 packed (cutlass-swizzled SF)
+//   SFA      : (M*K/16)  e4m3
+//   SFB      : (N*K/16)  e4m3
+//   bias_bf16: (N,)      bf16    per-col bias
+//   D_packed : (M, N/2)  uint8   NVFP4 packed
+//   SFD      : (M*N/16)  e4m3    output SF, cutlass-swizzled layout
+//   alpha    : float32           = sf_global_a * sf_global_b
+//
+// Stream-safe; per-shape workspace cached internally.
+void fp4_w4a16_gemm_bias_gelu_fp4out_sm120(
+    const void*  A_packed,
+    const void*  B_packed,
+    const void*  SFA,
+    const void*  SFB,
+    const void*  bias_bf16,
+    void*        D_packed,
+    void*        SFD,
+    int M, int N, int K,
+    float        alpha,
+    cudaStream_t stream);
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cu

@@ -0,0 +1,307 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS NVFP4 W4A16 GEMM with fused per-col bias epilogue, BF16 output,
+// **StreamK scheduler**, SM120a.
+//
+// fvk's default fp4_w4a16_gemm_sm120_bf16out uses
+// KernelTmaWarpSpecializedCooperative + PersistentScheduler. At motus
+// Wan FFN GEMM_dn shape (M=360, N=K_motus=3072, K=F=14336), TileShape
+// <128,128,256> yields 3 × 24 = 72 CTAs on 170 SMs = 0.42 wave, leaving
+// the GPU under-utilized. StreamK partitions the K-axis to issue more
+// CTAs and converges to ~1.7 waves, recovering 1.277× speedup standalone
+// (47 µs → 37 µs).
+//
+// Per-col bias is absorbed into the epilogue via LinCombPerColBiasEltAct
+// with Identity activation — eliminates the separate add_bias_bf16 launch.
+
+#include "cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cuh"
+
+#include "cute/tensor.hpp"
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/detail/sm100_blockscaled_layout.hpp"
+
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/thread/activation.h"
+#include "cutlass/epilogue/thread/linear_combination.h"
+#include "cutlass/epilogue/fusion/operations.hpp"
+
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+
+#include "cutlass/util/packed_stride.hpp"
+
+#include <cstdio>
+#include <mutex>
+#include <unordered_map>
+
+namespace flash_rt {
+namespace gemm {
+
+namespace {
+using namespace cute;
+
+using ElementA           = cutlass::float_e2m1_t;
+using ElementB           = cutlass::float_e2m1_t;
+using ElementC           = cutlass::bfloat16_t;
+using ElementD           = cutlass::bfloat16_t;
+using ElementBias        = cutlass::bfloat16_t;
+using ElementAccumulator = float;
+using ElementCompute     = float;
+using ElementSF          = cutlass::float_ue4m3_t;
+
+using LayoutA = cutlass::layout::RowMajor;
+using LayoutB = cutlass::layout::ColumnMajor;
+using LayoutC = cutlass::layout::RowMajor;
+using LayoutD = cutlass::layout::RowMajor;
+
+using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+
+constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value;
+constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value;
+constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+
+using TileShape    = Shape<_128, _128, _256>;
+using ClusterShape = Shape<_1, _1, _1>;
+using Sm1xxBlkScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+using FusionOperation = cutlass::epilogue::fusion::LinCombPerColBiasEltAct<
+    cutlass::epilogue::thread::Identity,
+    ElementD, ElementCompute, ElementBias, ElementC>;
+
+using CollectiveEpilogue =
+    typename cutlass::epilogue::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
+        TileShape, ClusterShape,
+        cutlass::epilogue::collective::EpilogueTileAuto,
+        ElementAccumulator, ElementCompute,
+        ElementC, LayoutC, AlignmentC,
+        ElementD, LayoutD, AlignmentD,
+        cutlass::epilogue::collective::EpilogueScheduleAuto,
+        FusionOperation
+    >::CollectiveOp;
+
+using CollectiveMainloop =
+    typename cutlass::gemm::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
+        ElementPairA, LayoutA, AlignmentA,
+        ElementPairB, LayoutB, AlignmentB,
+        ElementAccumulator,
+        TileShape, ClusterShape,
+        cutlass::gemm::collective::StageCountAutoCarveout<
+            static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+        cutlass::gemm::KernelTmaWarpSpecializedCooperative
+    >::CollectiveOp;
+
+using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+    Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
+    cutlass::gemm::StreamKScheduler>;
+
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+using NoBiasFusionOperation = cutlass::epilogue::fusion::LinearCombination<
+    ElementD, ElementCompute, ElementC, ElementCompute>;
+
+using NoBiasCollectiveEpilogue =
+    typename cutlass::epilogue::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
+        TileShape, ClusterShape,
+        cutlass::epilogue::collective::EpilogueTileAuto,
+        ElementAccumulator, ElementCompute,
+        ElementC, LayoutC, AlignmentC,
+        ElementD, LayoutD, AlignmentD,
+        cutlass::epilogue::collective::EpilogueScheduleAuto,
+        NoBiasFusionOperation
+    >::CollectiveOp;
+
+using NoBiasCollectiveMainloop =
+    typename cutlass::gemm::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
+        ElementPairA, LayoutA, AlignmentA,
+        ElementPairB, LayoutB, AlignmentB,
+        ElementAccumulator,
+        TileShape, ClusterShape,
+        cutlass::gemm::collective::StageCountAutoCarveout<
+            static_cast<int>(sizeof(typename NoBiasCollectiveEpilogue::SharedStorage))>,
+        cutlass::gemm::KernelTmaWarpSpecializedCooperative
+    >::CollectiveOp;
+
+using NoBiasGemmKernel = cutlass::gemm::kernel::GemmUniversal<
+    Shape<int, int, int, int>, NoBiasCollectiveMainloop,
+    NoBiasCollectiveEpilogue, cutlass::gemm::StreamKScheduler>;
+
+using NoBiasGemm =
+    cutlass::gemm::device::GemmUniversalAdapter<NoBiasGemmKernel>;
+
+struct ShapeKey {
+  int M, N, K;
+  bool operator==(const ShapeKey& o) const {
+    return M == o.M && N == o.N && K == o.K;
+  }
+};
+struct SHash {
+  size_t operator()(const ShapeKey& k) const noexcept {
+    return (size_t(k.M) * 1315423911u) ^ (size_t(k.N) * 2654435761u)
+         ^ size_t(k.K);
+  }
+};
+struct CachedWs { void* ptr = nullptr; size_t size = 0; };
+std::unordered_map<ShapeKey, CachedWs, SHash> g_ws;
+std::mutex g_mu;
+
+void* get_ws(int M, int N, int K, size_t need) {
+  std::lock_guard<std::mutex> lk(g_mu);
+  ShapeKey k{M, N, K};
+  auto it = g_ws.find(k);
+  if (it != g_ws.end() && it->second.size >= need) return it->second.ptr;
+  if (it != g_ws.end()) { cudaFree(it->second.ptr); g_ws.erase(it); }
+  CachedWs w; w.size = need;
+  if (need > 0) cudaMalloc(&w.ptr, need);
+  g_ws[k] = w;
+  return w.ptr;
+}
+
+}  // namespace
+
+void fp4_w4a16_gemm_dn_streamk_bf16out_sm120(
+    const void* A_packed, const void* B_packed,
+    const void* SFA,      const void* SFB,
+    void*       D_bf16,
+    int M, int N, int K,
+    float alpha,
+    cudaStream_t stream)
+{
+  using StrideA = typename NoBiasGemm::GemmKernel::StrideA;
+  using StrideB = typename NoBiasGemm::GemmKernel::StrideB;
+  using StrideC = typename NoBiasGemm::GemmKernel::StrideC;
+  using StrideD = typename NoBiasGemm::GemmKernel::StrideD;
+  StrideA strA = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1));
+  StrideB strB = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1));
+  StrideC strC = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1));
+  StrideD strD = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1));
+
+  auto problem = cute::make_shape(M, N, K, 1);
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem);
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem);
+
+  using ArrayElementA =
+      typename NoBiasGemm::GemmKernel::CollectiveMainloop::ArrayElementA;
+  using ArrayElementB =
+      typename NoBiasGemm::GemmKernel::CollectiveMainloop::ArrayElementB;
+
+  typename NoBiasGemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ArrayElementA const*>(A_packed), strA,
+          reinterpret_cast<ArrayElementB const*>(B_packed), strB,
+          reinterpret_cast<ElementSF const*>(SFA), layout_SFA,
+          reinterpret_cast<ElementSF const*>(SFB), layout_SFB
+      },
+      {
+          {alpha, 0.0f},
+          nullptr, strC,
+          reinterpret_cast<ElementD*>(D_bf16), strD
+      }
+  };
+
+  NoBiasGemm gemm;
+  size_t ws_size = NoBiasGemm::get_workspace_size(args);
+  void* ws_ptr = get_ws(M, N, K, ws_size);
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_dn_streamk_bf16out_sm120] can_implement FAIL "
+        "M=%d N=%d K=%d status=%d\n", M, N, K, int(status));
+    return;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_dn_streamk_bf16out_sm120] initialize FAIL "
+        "M=%d N=%d K=%d status=%d\n", M, N, K, int(status));
+    return;
+  }
+  status = gemm.run(stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_dn_streamk_bf16out_sm120] run FAIL status=%d\n",
+        int(status));
+  }
+}
+
+void fp4_w4a16_gemm_dn_streamk_bias_bf16out_sm120(
+    const void* A_packed, const void* B_packed,
+    const void* SFA,      const void* SFB,
+    const void* bias_bf16,
+    void*       D_bf16,
+    int M, int N, int K,
+    float alpha,
+    cudaStream_t stream)
+{
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+  StrideA strA = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1));
+  StrideB strB = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1));
+  StrideC strC = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1));
+  StrideD strD = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1));
+
+  auto problem = cute::make_shape(M, N, K, 1);
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem);
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem);
+
+  using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
+  using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
+
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ArrayElementA const*>(A_packed), strA,
+          reinterpret_cast<ArrayElementB const*>(B_packed), strB,
+          reinterpret_cast<ElementSF const*>(SFA), layout_SFA,
+          reinterpret_cast<ElementSF const*>(SFB), layout_SFB
+      },
+      {
+          {alpha, 0.0f},
+          nullptr, strC,
+          reinterpret_cast<ElementD*>(D_bf16), strD
+      }
+  };
+  args.epilogue.thread.bias_ptr =
+      reinterpret_cast<ElementBias const*>(bias_bf16);
+
+  Gemm gemm;
+  size_t ws_size = Gemm::get_workspace_size(args);
+  void* ws_ptr = get_ws(M, N, K, ws_size);
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_dn_streamk_bias_bf16out_sm120] can_implement FAIL "
+        "M=%d N=%d K=%d status=%d\n", M, N, K, int(status));
+    return;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_dn_streamk_bias_bf16out_sm120] initialize FAIL "
+        "M=%d N=%d K=%d status=%d\n", M, N, K, int(status));
+    return;
+  }
+  status = gemm.run(stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_dn_streamk_bias_bf16out_sm120] run FAIL status=%d\n",
+        int(status));
+  }
+}
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cuh

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS NVFP4 W4A16 GEMM_dn with fused per-col bias epilogue, BF16
+// output, **StreamK scheduler**, SM120a.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace gemm {
+
+// D = (alpha * (A_fp4 @ B_fp4^T scaled by SFA/SFB) + per_col_bias) → bf16
+//
+//   A_packed : (M, K/2)  uint8   NVFP4 packed (cutlass-swizzled SF)
+//   B_packed : (N, K/2)  uint8   NVFP4 packed
+//   SFA      : (M*K/16)  e4m3
+//   SFB      : (N*K/16)  e4m3
+//   bias_bf16: (N,)      bf16    per-col bias added in epilogue
+//   D_bf16   : (M, N)    bf16    output
+//   alpha    : float32           = sf_global_a * sf_global_b
+//
+// Stream-safe; per-shape workspace cached internally. Uses
+// StreamKScheduler to recover SM utilization at the motus Wan FFN
+// GEMM_dn shape (M=360, N=3072, K=14336): 1.277× over default
+// PersistentScheduler standalone.
+void fp4_w4a16_gemm_dn_streamk_bias_bf16out_sm120(
+    const void*  A_packed,
+    const void*  B_packed,
+    const void*  SFA,
+    const void*  SFB,
+    const void*  bias_bf16,
+    void*        D_bf16,
+    int M, int N, int K,
+    float        alpha,
+    cudaStream_t stream);
+
+// D = alpha * (A_fp4 @ B_fp4^T scaled by SFA/SFB) -> bf16
+//
+// Same StreamK schedule as the bias variant, but with a pure linear-combine
+// epilogue. This matches Motus down-only sites whose down bias is skipped.
+void fp4_w4a16_gemm_dn_streamk_bf16out_sm120(
+    const void*  A_packed,
+    const void*  B_packed,
+    const void*  SFA,
+    const void*  SFB,
+    void*        D_bf16,
+    int M, int N, int K,
+    float        alpha,
+    cudaStream_t stream);
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cu

@@ -0,0 +1,411 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Sm100 NVFP4 W4A16 block-scaled GEMM. BF16 output.
+//
+// Header: cutlass_nvfp4_w4a16_gemm_sm100.cuh.
+//
+// Template structure is a translation of the verified Sm120 path
+// (cutlass_nvfp4_w4a16_gemm_sm120.cu) onto the SM100 dispatch:
+//   - arch::Sm120                          -> arch::Sm100
+//   - KernelTmaWarpSpecializedCooperative  -> KernelScheduleAuto
+//   - KernelTmaWarpSpecializedPingpong     -> KernelScheduleAuto
+// All other types and layouts (FP4 e2m1 A/B, ue4m3 group scales,
+// row-major D in bf16, group_size=16) match the Sm120 variant byte
+// for byte. The wire-format contract (activation quantizer SFA layout,
+// loader SFB layout, alpha = sf_global_a * sf_global_b) is identical
+// so the Qwen3.6 frontend re-uses the same calls.
+//
+// Built only when GPU_ARCH==110 (Thor). The Sm100 dispatch reaches the
+// correct sm_110a tcgen05 mainloop without any per-arch macro.
+
+#include "cutlass_nvfp4_w4a16_gemm_sm100.cuh"
+
+#include "cute/tensor.hpp"
+#include "cute/atom/mma_atom.hpp"
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/detail/sm100_blockscaled_layout.hpp"
+
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/default_epilogue.hpp"
+#include "cutlass/epilogue/thread/linear_combination.h"
+
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+
+#include "cutlass/util/packed_stride.hpp"
+
+#include <cstdio>
+#include <mutex>
+#include <unordered_map>
+
+namespace flash_rt {
+namespace gemm {
+
+// ─────────────────────────────────────────────────────────────────
+// Default tile <128,128,256>, cluster <1,1,1>, schedule Auto.
+// ─────────────────────────────────────────────────────────────────
+namespace sm100_default {
+
+using namespace cute;
+
+using ElementA           = cutlass::float_e2m1_t;
+using ElementB           = cutlass::float_e2m1_t;
+using ElementC           = cutlass::bfloat16_t;
+using ElementD           = cutlass::bfloat16_t;
+using ElementAccumulator = float;
+using ElementCompute     = float;
+using ElementSF          = cutlass::float_ue4m3_t;
+
+using LayoutA = cutlass::layout::RowMajor;
+using LayoutB = cutlass::layout::ColumnMajor;
+using LayoutC = cutlass::layout::RowMajor;
+using LayoutD = cutlass::layout::RowMajor;
+
+using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+
+constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value;  // 32
+constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value;  // 32
+constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;     // 8
+constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;     // 8
+
+using TileShape    = Shape<_128, _128, _256>;
+using ClusterShape = Shape<_1, _1, _1>;
+
+using Sm1xxBlkScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+using CollectiveEpilogue =
+    typename cutlass::epilogue::collective::CollectiveBuilder<
+        cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp,
+        TileShape, ClusterShape,
+        cutlass::epilogue::collective::EpilogueTileAuto,
+        ElementAccumulator, ElementCompute,
+        ElementC, LayoutC, AlignmentC,
+        ElementD, LayoutD, AlignmentD,
+        cutlass::epilogue::collective::EpilogueScheduleAuto
+    >::CollectiveOp;
+
+using CollectiveMainloop =
+    typename cutlass::gemm::collective::CollectiveBuilder<
+        cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp,
+        ElementPairA, LayoutA, AlignmentA,
+        ElementPairB, LayoutB, AlignmentB,
+        ElementAccumulator,
+        TileShape, ClusterShape,
+        cutlass::gemm::collective::StageCountAutoCarveout<
+            static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+        cutlass::gemm::collective::KernelScheduleAuto
+    >::CollectiveOp;
+
+using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+    Shape<int, int, int, int>,
+    CollectiveMainloop,
+    CollectiveEpilogue>;
+
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+struct ShapeKey {
+  int M, N, K;
+  bool operator==(const ShapeKey& o) const {
+    return M == o.M && N == o.N && K == o.K;
+  }
+};
+struct ShapeKeyHash {
+  size_t operator()(const ShapeKey& k) const noexcept {
+    return (static_cast<size_t>(k.M) * 1315423911u)
+         ^ (static_cast<size_t>(k.N) * 2654435761u)
+         ^ static_cast<size_t>(k.K);
+  }
+};
+struct CachedWorkspace { void* ptr = nullptr; size_t size = 0; };
+
+std::unordered_map<ShapeKey, CachedWorkspace, ShapeKeyHash> g_ws_cache;
+std::mutex g_ws_mu;
+
+void* get_workspace(int M, int N, int K, size_t needed) {
+  std::lock_guard<std::mutex> lk(g_ws_mu);
+  ShapeKey key{M, N, K};
+  auto it = g_ws_cache.find(key);
+  if (it != g_ws_cache.end() && it->second.size >= needed) return it->second.ptr;
+  if (it != g_ws_cache.end()) { cudaFree(it->second.ptr); g_ws_cache.erase(it); }
+  CachedWorkspace w; w.size = needed;
+  if (needed > 0) cudaMalloc(&w.ptr, needed);
+  g_ws_cache[key] = w;
+  return w.ptr;
+}
+
+cutlass::Status run_gemm(
+    const void* A_packed, const void* B_packed, void* D_bf16,
+    int M, int N, int K,
+    const void* SFA, const void* SFB,
+    float alpha,
+    cudaStream_t stream)
+{
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+
+  StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1));
+  StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1));
+  StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1));
+  StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1));
+
+  auto problem_shape_MNKL = cute::make_shape(M, N, K, 1);
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL);
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL);
+
+  using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
+  using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
+
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ArrayElementA const*>(A_packed), stride_A,
+          reinterpret_cast<ArrayElementB const*>(B_packed), stride_B,
+          reinterpret_cast<ElementSF const*>(SFA), layout_SFA,
+          reinterpret_cast<ElementSF const*>(SFB), layout_SFB
+      },
+      {
+          {alpha, 0.0f},
+          nullptr, stride_C,
+          reinterpret_cast<ElementD*>(D_bf16), stride_D
+      }
+  };
+
+  Gemm gemm;
+  size_t ws_size = Gemm::get_workspace_size(args);
+  void* ws_ptr = get_workspace(M, N, K, ws_size);
+
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm100_bf16out] can_implement FAIL M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm100_bf16out] initialize FAIL M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  return gemm.run(stream);
+}
+
+}  // namespace sm100_default
+
+void fp4_w4a16_gemm_sm100_bf16out(
+    const void*  A_packed, const void*  B_packed, void*  D_bf16,
+    int M, int N, int K,
+    const void*  SFA, const void*  SFB,
+    float alpha, cudaStream_t stream)
+{
+  cutlass::Status status = sm100_default::run_gemm(
+      A_packed, B_packed, D_bf16, M, N, K, SFA, SFB, alpha, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm100_bf16out] run FAIL M=%d N=%d K=%d (status=%d); D output undefined\n",
+        M, N, K, static_cast<int>(status));
+  }
+}
+
+// ─────────────────────────────────────────────────────────────────
+// Wide-N tile <128,256,128>, cluster <1,1,1>, schedule Auto.
+// ─────────────────────────────────────────────────────────────────
+namespace sm100_widen {
+
+using namespace cute;
+
+using ElementA           = cutlass::float_e2m1_t;
+using ElementB           = cutlass::float_e2m1_t;
+using ElementC           = cutlass::bfloat16_t;
+using ElementD           = cutlass::bfloat16_t;
+using ElementAccumulator = float;
+using ElementCompute     = float;
+using ElementSF          = cutlass::float_ue4m3_t;
+using LayoutA = cutlass::layout::RowMajor;
+using LayoutB = cutlass::layout::ColumnMajor;
+using LayoutC = cutlass::layout::RowMajor;
+using LayoutD = cutlass::layout::RowMajor;
+using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value;
+constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value;
+constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+
+using TileShape    = Shape<_128, _256, _128>;
+using ClusterShape = Shape<_1, _1, _1>;
+using Sm1xxBlkScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+using CollectiveEpilogue =
+    typename cutlass::epilogue::collective::CollectiveBuilder<
+        cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp,
+        TileShape, ClusterShape,
+        cutlass::epilogue::collective::EpilogueTileAuto,
+        ElementAccumulator, ElementCompute,
+        ElementC, LayoutC, AlignmentC,
+        ElementD, LayoutD, AlignmentD,
+        cutlass::epilogue::collective::EpilogueScheduleAuto
+    >::CollectiveOp;
+
+using CollectiveMainloop =
+    typename cutlass::gemm::collective::CollectiveBuilder<
+        cutlass::arch::Sm100, cutlass::arch::OpClassBlockScaledTensorOp,
+        ElementPairA, LayoutA, AlignmentA,
+        ElementPairB, LayoutB, AlignmentB,
+        ElementAccumulator,
+        TileShape, ClusterShape,
+        cutlass::gemm::collective::StageCountAutoCarveout<
+            static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+        cutlass::gemm::collective::KernelScheduleAuto
+    >::CollectiveOp;
+
+using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+    Shape<int, int, int, int>,
+    CollectiveMainloop,
+    CollectiveEpilogue>;
+
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+struct ShapeKey {
+  int M, N, K;
+  bool operator==(const ShapeKey& o) const {
+    return M == o.M && N == o.N && K == o.K;
+  }
+};
+struct ShapeKeyHash {
+  size_t operator()(const ShapeKey& k) const noexcept {
+    return (static_cast<size_t>(k.M) * 1315423911u)
+         ^ (static_cast<size_t>(k.N) * 2654435761u)
+         ^ static_cast<size_t>(k.K);
+  }
+};
+struct CachedWorkspace { void* ptr = nullptr; size_t size = 0; };
+
+std::unordered_map<ShapeKey, CachedWorkspace, ShapeKeyHash> g_ws_cache_widen;
+std::mutex g_ws_mu_widen;
+
+void* get_workspace_widen(int M, int N, int K, size_t needed) {
+  std::lock_guard<std::mutex> lk(g_ws_mu_widen);
+  ShapeKey key{M, N, K};
+  auto it = g_ws_cache_widen.find(key);
+  if (it != g_ws_cache_widen.end() && it->second.size >= needed) return it->second.ptr;
+  if (it != g_ws_cache_widen.end()) { cudaFree(it->second.ptr); g_ws_cache_widen.erase(it); }
+  CachedWorkspace w; w.size = needed;
+  if (needed > 0) cudaMalloc(&w.ptr, needed);
+  g_ws_cache_widen[key] = w;
+  return w.ptr;
+}
+
+cutlass::Status run_gemm(
+    const void* A_packed, const void* B_packed, void* D_bf16,
+    int M, int N, int K,
+    const void* SFA, const void* SFB,
+    float alpha,
+    cudaStream_t stream)
+{
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+
+  StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1));
+  StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1));
+  StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1));
+  StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1));
+
+  auto problem_shape_MNKL = cute::make_shape(M, N, K, 1);
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL);
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL);
+
+  using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
+  using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
+
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ArrayElementA const*>(A_packed), stride_A,
+          reinterpret_cast<ArrayElementB const*>(B_packed), stride_B,
+          reinterpret_cast<ElementSF const*>(SFA), layout_SFA,
+          reinterpret_cast<ElementSF const*>(SFB), layout_SFB
+      },
+      {
+          {alpha, 0.0f},
+          nullptr, stride_C,
+          reinterpret_cast<ElementD*>(D_bf16), stride_D
+      }
+  };
+
+  Gemm gemm;
+  size_t ws_size = Gemm::get_workspace_size(args);
+  void* ws_ptr = get_workspace_widen(M, N, K, ws_size);
+
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm100_bf16out_widen] can_implement FAIL M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm100_bf16out_widen] initialize FAIL M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  return gemm.run(stream);
+}
+
+}  // namespace sm100_widen
+
+void fp4_w4a16_gemm_sm100_bf16out_widen(
+    const void*  A_packed, const void*  B_packed, void*  D_bf16,
+    int M, int N, int K,
+    const void*  SFA, const void*  SFB,
+    float alpha, cudaStream_t stream)
+{
+  cutlass::Status status = sm100_widen::run_gemm(
+      A_packed, B_packed, D_bf16, M, N, K, SFA, SFB, alpha, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm100_bf16out_widen] run FAIL M=%d N=%d K=%d (status=%d); D output undefined\n",
+        M, N, K, static_cast<int>(status));
+  }
+}
+
+// ─────────────────────────────────────────────────────────────────
+// Pingpong placeholder: same default tile + Auto schedule. The Sm100
+// dispatch under KernelScheduleAuto already exercises a 2SM pingpong-
+// style schedule, so this entry exists for binding parity with the
+// Sm120 surface. A dedicated alternate schedule may replace this
+// after the Thor tile sweep.
+// ─────────────────────────────────────────────────────────────────
+void fp4_w4a16_gemm_sm100_bf16out_pingpong(
+    const void*  A_packed, const void*  B_packed, void*  D_bf16,
+    int M, int N, int K,
+    const void*  SFA, const void*  SFB,
+    float alpha, cudaStream_t stream)
+{
+  // Routes through the default tile until the Thor sweep adds a
+  // distinct pingpong-equivalent schedule.
+  cutlass::Status status = sm100_default::run_gemm(
+      A_packed, B_packed, D_bf16, M, N, K, SFA, SFB, alpha, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm100_bf16out_pingpong] run FAIL M=%d N=%d K=%d (status=%d); D output undefined\n",
+        M, N, K, static_cast<int>(status));
+  }
+}
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cuh

@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS-based NVFP4 W4A16 GEMM for sm_100-class Blackwell (datacenter
+// SM100 / Jetson AGX Thor SM110). Block-scaled FP4 GEMM matching the
+// Qwen3.6 NVFP4 ckpt schema (compressed-tensors `nvfp4-pack-quantized`).
+//
+// Sibling of cutlass_nvfp4_w4a16_gemm_sm120.cuh. The two differ only in
+// CUTLASS arch dispatch and kernel-schedule policy:
+//   - sm120: arch::Sm120 + KernelTmaWarpSpecializedCooperative /
+//            KernelTmaWarpSpecializedPingpong
+//   - sm100: arch::Sm100 + KernelScheduleAuto
+// On Thor (sm_110a) the Sm100 dispatch path produces the correct
+// blockscaled tcgen05 mainloop.
+//
+// Wire-format contract is identical to the sm120 variant, so the
+// Python-side weight/scale layout and the activation quantizer are
+// reused unchanged. The pybind layer binds these Thor symbols under
+// the existing public names ``fp4_w4a16_gemm_sm120_bf16out*`` so the
+// Qwen3.6 frontend code path does not need any hardware fork.
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace gemm {
+
+// Default tile <128,128,256>, cluster <1,1,1>, KernelScheduleAuto.
+void fp4_w4a16_gemm_sm100_bf16out(
+    const void*  A_packed,    // (M, K/2)        u8  row-major
+    const void*  B_packed,    // (N, K/2)        u8  row-major (read as ColMajor (K,N))
+    void*        D_bf16,      // (M, N)          bf16 row-major
+    int M, int N, int K,
+    const void*  SFA,         // (M, K/16)       e4m3 (Sm1xx blockscaled atom layout)
+    const void*  SFB,         // (N, K/16)       e4m3 (Sm1xx blockscaled atom layout)
+    float        alpha,       // = sf_global_a * sf_global_b
+    cudaStream_t stream);
+
+// Wide-N tile <128,256,128>, cluster <1,1,1>, KernelScheduleAuto.
+// For shapes with very large N (lm_head, MLP gate/up).
+void fp4_w4a16_gemm_sm100_bf16out_widen(
+    const void*  A_packed,
+    const void*  B_packed,
+    void*        D_bf16,
+    int M, int N, int K,
+    const void*  SFA,
+    const void*  SFB,
+    float        alpha,
+    cudaStream_t stream);
+
+// Default tile <128,128,256>, cluster <1,1,1>, KernelScheduleAuto.
+// Kept as a separate symbol so callers can A/B against the default
+// variant after the tile sweep produces a Thor-tuned schedule.
+void fp4_w4a16_gemm_sm100_bf16out_pingpong(
+    const void*  A_packed,
+    const void*  B_packed,
+    void*        D_bf16,
+    int M, int N, int K,
+    const void*  SFA,
+    const void*  SFB,
+    float        alpha,
+    cudaStream_t stream);
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm120.cu

@@ -0,0 +1,690 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS NVFP4 W4A16 block-scaled GEMM, SM120a, BF16 output.
+// Header: cutlass_nvfp4_w4a16_gemm_sm120.cuh
+//
+// Template is a direct port of NVIDIA's verified unit test at
+//   third_party/cutlass/test/unit/gemm/device/
+//     sm120_blockscaled_tensorop_gemm/sm120_bs_gemm_nvf4_nvf4_f32_bf16.cu
+// (one config: TileShape <128,128,256>, ClusterShape <1,1,1>,
+//  KernelTmaWarpSpecializedCooperative, OpClassBlockScaledTensorOp).
+//
+// Why this is the "vendor-best" path (not hand-written):
+//   * Uses CUTLASS 4.x's `CollectiveBuilder` for SM120 BlockScaled
+//     mainloop + epilogue. NVIDIA tunes the mainloop schedule.
+//   * Same family as the existing FP8 SM120 GEMM (cutlass_sm120_block128_
+//     fp8_gemm.cu) — only the element types and `OpClass` differ.
+//   * No handwritten PTX or custom layout swizzling.
+//
+// Per-shape argument cache + workspace match the FP8 path so the
+// hot-path is launch-only.
+
+#include "cutlass_nvfp4_w4a16_gemm_sm120.cuh"
+
+#include "cute/tensor.hpp"
+#include "cute/atom/mma_atom.hpp"
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/detail/sm100_blockscaled_layout.hpp"
+
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/default_epilogue.hpp"
+#include "cutlass/epilogue/thread/linear_combination.h"
+
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+
+#include "cutlass/util/packed_stride.hpp"
+
+#include <cstdio>
+#include <mutex>
+#include <unordered_map>
+
+namespace flash_rt {
+namespace gemm {
+
+namespace {
+
+using namespace cute;
+
+// ── Element / layout types (copy of unit test "kernel_1") ────────
+using ElementA           = cutlass::float_e2m1_t;
+using ElementB           = cutlass::float_e2m1_t;
+using ElementC           = cutlass::bfloat16_t;
+using ElementD           = cutlass::bfloat16_t;
+using ElementAccumulator = float;
+using ElementCompute     = float;
+using ElementSF          = cutlass::float_ue4m3_t;
+
+using LayoutA = cutlass::layout::RowMajor;
+using LayoutB = cutlass::layout::ColumnMajor;
+// Public API hands us D in row-major (matches the FP8 sm120 kernel
+// and what HF / our pipeline downstream expects for (M, N) tensors).
+using LayoutC = cutlass::layout::RowMajor;
+using LayoutD = cutlass::layout::RowMajor;
+
+using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+
+// 16-byte alignments in bits = 16 * 8 = 128. Convert to element count.
+constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value;  // 32
+constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value;  // 32
+constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;     // 8
+constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;     // 8
+
+using TileShape    = Shape<_128, _128, _256>;
+using ClusterShape = Shape<_1, _1, _1>;
+
+// SF tensor layout helper. Vector size = 16 (NVFP4 group) per ckpt.
+// Sm1xxBlockScaledConfig generates the (M-blk, K-blk) atom layout
+// CUTLASS expects on-device. Our weight loader and act quantizer
+// must produce SF in this exact layout (transformed once).
+using Sm1xxBlkScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+using CollectiveEpilogue =
+    typename cutlass::epilogue::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
+        TileShape, ClusterShape,
+        cutlass::epilogue::collective::EpilogueTileAuto,
+        ElementAccumulator, ElementCompute,
+        ElementC, LayoutC, AlignmentC,
+        ElementD, LayoutD, AlignmentD,
+        cutlass::epilogue::collective::EpilogueScheduleAuto
+    >::CollectiveOp;
+
+using CollectiveMainloop =
+    typename cutlass::gemm::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
+        ElementPairA, LayoutA, AlignmentA,
+        ElementPairB, LayoutB, AlignmentB,
+        ElementAccumulator,
+        TileShape, ClusterShape,
+        cutlass::gemm::collective::StageCountAutoCarveout<
+            static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+        cutlass::gemm::KernelTmaWarpSpecializedCooperative
+    >::CollectiveOp;
+
+using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+    Shape<int, int, int, int>,
+    CollectiveMainloop,
+    CollectiveEpilogue,
+    cutlass::gemm::PersistentScheduler>;
+
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+// ── Per-shape workspace cache (mirrors FP8 path) ─────────────────
+struct ShapeKey {
+  int M, N, K;
+  bool operator==(const ShapeKey& o) const {
+    return M == o.M && N == o.N && K == o.K;
+  }
+};
+struct ShapeKeyHash {
+  size_t operator()(const ShapeKey& k) const noexcept {
+    return (static_cast<size_t>(k.M) * 1315423911u)
+         ^ (static_cast<size_t>(k.N) * 2654435761u)
+         ^ static_cast<size_t>(k.K);
+  }
+};
+
+struct CachedWorkspace {
+  void*  ptr  = nullptr;
+  size_t size = 0;
+};
+
+std::unordered_map<ShapeKey, CachedWorkspace, ShapeKeyHash> g_ws_cache;
+std::mutex g_ws_mu;
+
+void* get_workspace(int M, int N, int K, size_t needed) {
+  std::lock_guard<std::mutex> lk(g_ws_mu);
+  ShapeKey key{M, N, K};
+  auto it = g_ws_cache.find(key);
+  if (it != g_ws_cache.end() && it->second.size >= needed) {
+    return it->second.ptr;
+  }
+  if (it != g_ws_cache.end()) {
+    cudaFree(it->second.ptr);
+    g_ws_cache.erase(it);
+  }
+  CachedWorkspace w;
+  w.size = needed;
+  if (needed > 0) {
+    cudaMalloc(&w.ptr, needed);
+  }
+  g_ws_cache[key] = w;
+  return w.ptr;
+}
+
+cutlass::Status run_gemm(
+    const void* A_packed, const void* B_packed, void* D_bf16,
+    int M, int N, int K,
+    const void* SFA, const void* SFB,
+    float alpha,
+    cudaStream_t stream)
+{
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+
+  StrideA stride_A = cutlass::make_cute_packed_stride(
+      StrideA{}, cute::make_shape(M, K, 1));
+  StrideB stride_B = cutlass::make_cute_packed_stride(
+      StrideB{}, cute::make_shape(N, K, 1));
+  StrideC stride_C = cutlass::make_cute_packed_stride(
+      StrideC{}, cute::make_shape(M, N, 1));
+  StrideD stride_D = cutlass::make_cute_packed_stride(
+      StrideD{}, cute::make_shape(M, N, 1));
+
+  auto problem_shape_MNKL = cute::make_shape(M, N, K, 1);
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL);
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL);
+
+  using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
+  using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
+
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ArrayElementA const*>(A_packed), stride_A,
+          reinterpret_cast<ArrayElementB const*>(B_packed), stride_B,
+          reinterpret_cast<ElementSF const*>(SFA), layout_SFA,
+          reinterpret_cast<ElementSF const*>(SFB), layout_SFB
+      },
+      {
+          {alpha, 0.0f},                                      // (alpha, beta)
+          nullptr, stride_C,                                  // C unused (beta=0)
+          reinterpret_cast<ElementD*>(D_bf16), stride_D
+      }
+  };
+
+  Gemm gemm;
+  size_t ws_size = Gemm::get_workspace_size(args);
+  void* ws_ptr = get_workspace(M, N, K, ws_size);
+
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm120_bf16out] can_implement FAIL "
+        "M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm120_bf16out] initialize FAIL "
+        "M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  return gemm.run(stream);
+}
+
+// Same default tile, but with a per-element residual C addend folded into the
+// epilogue: D = alpha*(A*B) + C. The epilogue already carries the C operand
+// (ElementC=bf16) — here we just feed C (beta=1) instead of nullptr (beta=0).
+// Lets o_proj/down fuse their residual add, so the following rms_norm reads ONE
+// tensor (D) instead of two (gemm_out + residual). C must be bf16 (M,N) row-major.
+cutlass::Status run_gemm_residual(
+    const void* A_packed, const void* B_packed,
+    const void* C_residual, void* D_bf16,
+    int M, int N, int K,
+    const void* SFA, const void* SFB,
+    float alpha,
+    cudaStream_t stream)
+{
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+
+  StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1));
+  StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1));
+  StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1));
+  StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1));
+
+  auto problem_shape_MNKL = cute::make_shape(M, N, K, 1);
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL);
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL);
+
+  using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
+  using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
+
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ArrayElementA const*>(A_packed), stride_A,
+          reinterpret_cast<ArrayElementB const*>(B_packed), stride_B,
+          reinterpret_cast<ElementSF const*>(SFA), layout_SFA,
+          reinterpret_cast<ElementSF const*>(SFB), layout_SFB
+      },
+      {
+          {alpha, 1.0f},                                       // (alpha, beta=1)
+          reinterpret_cast<ElementC const*>(C_residual), stride_C,
+          reinterpret_cast<ElementD*>(D_bf16), stride_D
+      }
+  };
+
+  Gemm gemm;
+  size_t ws_size = Gemm::get_workspace_size(args);
+  void* ws_ptr = get_workspace(M, N, K, ws_size);
+
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr, "[fp4_w4a16_gemm_residual] can_implement FAIL "
+        "M=%d N=%d K=%d (status=%d)\n", M, N, K, static_cast<int>(status));
+    return status;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) return status;
+  return gemm.run(stream);
+}
+
+}  // namespace
+
+void fp4_w4a16_gemm_residual_sm120_bf16out(
+    const void* A_packed, const void* B_packed,
+    const void* C_residual, void* D_bf16,
+    int M, int N, int K,
+    const void* SFA, const void* SFB,
+    float alpha, cudaStream_t stream)
+{
+  cutlass::Status status = run_gemm_residual(
+      A_packed, B_packed, C_residual, D_bf16, M, N, K, SFA, SFB, alpha, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr, "[fp4_w4a16_gemm_residual_sm120_bf16out] run FAIL "
+        "M=%d N=%d K=%d (status=%d)\n", M, N, K, static_cast<int>(status));
+  }
+}
+
+void fp4_w4a16_gemm_sm120_bf16out(
+    const void*  A_packed,
+    const void*  B_packed,
+    void*        D_bf16,
+    int M, int N, int K,
+    const void*  SFA,
+    const void*  SFB,
+    float        alpha,
+    cudaStream_t stream)
+{
+  cutlass::Status status = run_gemm(
+      A_packed, B_packed, D_bf16, M, N, K, SFA, SFB, alpha, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm120_bf16out] run FAIL "
+        "M=%d N=%d K=%d (status=%d); D output undefined\n",
+        M, N, K, static_cast<int>(status));
+  }
+}
+
+// ============================================================
+// WIDEN variant: TileShape <128, 256, 128>. Same kernel template
+// machinery as above but a wider N tile + narrower K. Profiled
+// faster on shapes with very large N (lm_head N=248320: 88% BW vs
+// 64% baseline; MLP gate/up N=17408: 66% vs 56%). Slower on small/
+// medium-N shapes (k/v_proj N=1024, lin in_proj_z N=6144, etc.) so
+// callers dispatch by shape.
+// ============================================================
+namespace widen {
+
+using namespace cute;
+
+using ElementA           = cutlass::float_e2m1_t;
+using ElementB           = cutlass::float_e2m1_t;
+using ElementC           = cutlass::bfloat16_t;
+using ElementD           = cutlass::bfloat16_t;
+using ElementAccumulator = float;
+using ElementCompute     = float;
+using ElementSF          = cutlass::float_ue4m3_t;
+using LayoutA = cutlass::layout::RowMajor;
+using LayoutB = cutlass::layout::ColumnMajor;
+using LayoutC = cutlass::layout::RowMajor;
+using LayoutD = cutlass::layout::RowMajor;
+using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value;
+constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value;
+constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+
+using TileShape    = Shape<_128, _256, _128>;
+using ClusterShape = Shape<_1, _1, _1>;
+
+using Sm1xxBlkScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+using CollectiveEpilogue =
+    typename cutlass::epilogue::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
+        TileShape, ClusterShape,
+        cutlass::epilogue::collective::EpilogueTileAuto,
+        ElementAccumulator, ElementCompute,
+        ElementC, LayoutC, AlignmentC,
+        ElementD, LayoutD, AlignmentD,
+        cutlass::epilogue::collective::EpilogueScheduleAuto
+    >::CollectiveOp;
+
+using CollectiveMainloop =
+    typename cutlass::gemm::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
+        ElementPairA, LayoutA, AlignmentA,
+        ElementPairB, LayoutB, AlignmentB,
+        ElementAccumulator,
+        TileShape, ClusterShape,
+        cutlass::gemm::collective::StageCountAutoCarveout<
+            static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+        cutlass::gemm::KernelTmaWarpSpecializedCooperative
+    >::CollectiveOp;
+
+using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+    Shape<int, int, int, int>,
+    CollectiveMainloop,
+    CollectiveEpilogue,
+    cutlass::gemm::PersistentScheduler>;
+
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+struct ShapeKey {
+  int M, N, K;
+  bool operator==(const ShapeKey& o) const {
+    return M == o.M && N == o.N && K == o.K;
+  }
+};
+struct ShapeKeyHash {
+  size_t operator()(const ShapeKey& k) const noexcept {
+    return (static_cast<size_t>(k.M) * 1315423911u)
+         ^ (static_cast<size_t>(k.N) * 2654435761u)
+         ^ static_cast<size_t>(k.K);
+  }
+};
+struct CachedWorkspace { void* ptr = nullptr; size_t size = 0; };
+std::unordered_map<ShapeKey, CachedWorkspace, ShapeKeyHash> g_ws_cache_widen;
+std::mutex g_ws_mu_widen;
+
+void* get_workspace_widen(int M, int N, int K, size_t needed) {
+  std::lock_guard<std::mutex> lk(g_ws_mu_widen);
+  ShapeKey key{M, N, K};
+  auto it = g_ws_cache_widen.find(key);
+  if (it != g_ws_cache_widen.end() && it->second.size >= needed) {
+    return it->second.ptr;
+  }
+  if (it != g_ws_cache_widen.end()) {
+    cudaFree(it->second.ptr);
+    g_ws_cache_widen.erase(it);
+  }
+  CachedWorkspace w;
+  w.size = needed;
+  if (needed > 0) cudaMalloc(&w.ptr, needed);
+  g_ws_cache_widen[key] = w;
+  return w.ptr;
+}
+
+cutlass::Status run_gemm_widen(
+    const void* A_packed, const void* B_packed, void* D_bf16,
+    int M, int N, int K,
+    const void* SFA, const void* SFB,
+    float alpha,
+    cudaStream_t stream)
+{
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+
+  StrideA stride_A = cutlass::make_cute_packed_stride(
+      StrideA{}, cute::make_shape(M, K, 1));
+  StrideB stride_B = cutlass::make_cute_packed_stride(
+      StrideB{}, cute::make_shape(N, K, 1));
+  StrideC stride_C = cutlass::make_cute_packed_stride(
+      StrideC{}, cute::make_shape(M, N, 1));
+  StrideD stride_D = cutlass::make_cute_packed_stride(
+      StrideD{}, cute::make_shape(M, N, 1));
+
+  auto problem_shape_MNKL = cute::make_shape(M, N, K, 1);
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL);
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL);
+
+  using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
+  using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
+
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ArrayElementA const*>(A_packed), stride_A,
+          reinterpret_cast<ArrayElementB const*>(B_packed), stride_B,
+          reinterpret_cast<ElementSF const*>(SFA), layout_SFA,
+          reinterpret_cast<ElementSF const*>(SFB), layout_SFB
+      },
+      {
+          {alpha, 0.0f},
+          nullptr, stride_C,
+          reinterpret_cast<ElementD*>(D_bf16), stride_D
+      }
+  };
+
+  Gemm gemm;
+  size_t ws_size = Gemm::get_workspace_size(args);
+  void* ws_ptr = get_workspace_widen(M, N, K, ws_size);
+
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm120_bf16out_widen] can_implement FAIL "
+        "M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm120_bf16out_widen] initialize FAIL "
+        "M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  return gemm.run(stream);
+}
+
+}  // namespace widen
+
+void fp4_w4a16_gemm_sm120_bf16out_widen(
+    const void*  A_packed,
+    const void*  B_packed,
+    void*        D_bf16,
+    int M, int N, int K,
+    const void*  SFA,
+    const void*  SFB,
+    float        alpha,
+    cudaStream_t stream)
+{
+  cutlass::Status status = widen::run_gemm_widen(
+      A_packed, B_packed, D_bf16, M, N, K, SFA, SFB, alpha, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm120_bf16out_widen] run FAIL "
+        "M=%d N=%d K=%d (status=%d); D output undefined\n",
+        M, N, K, static_cast<int>(status));
+  }
+}
+
+// ============================================================
+// PINGPONG variant: default <128,128,256> tile with
+// KernelTmaWarpSpecializedPingpong. This is intentionally separate from
+// the production entrypoint so Qwen can A/B per shape before any dispatch
+// policy is changed.
+// ============================================================
+namespace pingpong {
+
+using namespace cute;
+
+using ElementA           = cutlass::float_e2m1_t;
+using ElementB           = cutlass::float_e2m1_t;
+using ElementC           = cutlass::bfloat16_t;
+using ElementD           = cutlass::bfloat16_t;
+using ElementAccumulator = float;
+using ElementCompute     = float;
+using ElementSF          = cutlass::float_ue4m3_t;
+using LayoutA = cutlass::layout::RowMajor;
+using LayoutB = cutlass::layout::ColumnMajor;
+using LayoutC = cutlass::layout::RowMajor;
+using LayoutD = cutlass::layout::RowMajor;
+using ElementPairA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+using ElementPairB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+constexpr int AlignmentA = 16 * 8 / cutlass::sizeof_bits<ElementA>::value;
+constexpr int AlignmentB = 16 * 8 / cutlass::sizeof_bits<ElementB>::value;
+constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+
+using TileShape    = Shape<_128, _128, _256>;
+using ClusterShape = Shape<_1, _1, _1>;
+using Sm1xxBlkScaledConfig = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+using CollectiveEpilogue =
+    typename cutlass::epilogue::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassTensorOp,
+        TileShape, ClusterShape,
+        cutlass::epilogue::collective::EpilogueTileAuto,
+        ElementAccumulator, ElementCompute,
+        ElementC, LayoutC, AlignmentC,
+        ElementD, LayoutD, AlignmentD,
+        cutlass::epilogue::collective::EpilogueScheduleAuto
+    >::CollectiveOp;
+
+using CollectiveMainloop =
+    typename cutlass::gemm::collective::CollectiveBuilder<
+        cutlass::arch::Sm120, cutlass::arch::OpClassBlockScaledTensorOp,
+        ElementPairA, LayoutA, AlignmentA,
+        ElementPairB, LayoutB, AlignmentB,
+        ElementAccumulator,
+        TileShape, ClusterShape,
+        cutlass::gemm::collective::StageCountAutoCarveout<
+            static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+        cutlass::gemm::KernelTmaWarpSpecializedPingpong
+    >::CollectiveOp;
+
+using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+    Shape<int, int, int, int>,
+    CollectiveMainloop,
+    CollectiveEpilogue,
+    cutlass::gemm::PersistentScheduler>;
+
+using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+std::unordered_map<ShapeKey, CachedWorkspace, ShapeKeyHash> g_ws_cache_pingpong;
+std::mutex g_ws_mu_pingpong;
+
+void* get_workspace_pingpong(int M, int N, int K, size_t needed) {
+  std::lock_guard<std::mutex> lk(g_ws_mu_pingpong);
+  ShapeKey key{M, N, K};
+  auto it = g_ws_cache_pingpong.find(key);
+  if (it != g_ws_cache_pingpong.end() && it->second.size >= needed) {
+    return it->second.ptr;
+  }
+  if (it != g_ws_cache_pingpong.end()) {
+    cudaFree(it->second.ptr);
+    g_ws_cache_pingpong.erase(it);
+  }
+  CachedWorkspace w;
+  w.size = needed;
+  if (needed > 0) cudaMalloc(&w.ptr, needed);
+  g_ws_cache_pingpong[key] = w;
+  return w.ptr;
+}
+
+cutlass::Status run_gemm_pingpong(
+    const void* A_packed, const void* B_packed, void* D_bf16,
+    int M, int N, int K,
+    const void* SFA, const void* SFB,
+    float alpha,
+    cudaStream_t stream)
+{
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+
+  StrideA stride_A = cutlass::make_cute_packed_stride(
+      StrideA{}, cute::make_shape(M, K, 1));
+  StrideB stride_B = cutlass::make_cute_packed_stride(
+      StrideB{}, cute::make_shape(N, K, 1));
+  StrideC stride_C = cutlass::make_cute_packed_stride(
+      StrideC{}, cute::make_shape(M, N, 1));
+  StrideD stride_D = cutlass::make_cute_packed_stride(
+      StrideD{}, cute::make_shape(M, N, 1));
+
+  auto problem_shape_MNKL = cute::make_shape(M, N, K, 1);
+  auto layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(problem_shape_MNKL);
+  auto layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(problem_shape_MNKL);
+
+  using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA;
+  using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB;
+
+  typename Gemm::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGemm,
+      {M, N, K, 1},
+      {
+          reinterpret_cast<ArrayElementA const*>(A_packed), stride_A,
+          reinterpret_cast<ArrayElementB const*>(B_packed), stride_B,
+          reinterpret_cast<ElementSF const*>(SFA), layout_SFA,
+          reinterpret_cast<ElementSF const*>(SFB), layout_SFB
+      },
+      {
+          {alpha, 0.0f},
+          nullptr, stride_C,
+          reinterpret_cast<ElementD*>(D_bf16), stride_D
+      }
+  };
+
+  Gemm gemm;
+  size_t ws_size = Gemm::get_workspace_size(args);
+  void* ws_ptr = get_workspace_pingpong(M, N, K, ws_size);
+
+  auto status = gemm.can_implement(args);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm120_bf16out_pingpong] can_implement FAIL "
+        "M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  status = gemm.initialize(args, ws_ptr, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm120_bf16out_pingpong] initialize FAIL "
+        "M=%d N=%d K=%d (status=%d)\n",
+        M, N, K, static_cast<int>(status));
+    return status;
+  }
+  return gemm.run(stream);
+}
+
+}  // namespace pingpong
+
+void fp4_w4a16_gemm_sm120_bf16out_pingpong(
+    const void*  A_packed,
+    const void*  B_packed,
+    void*        D_bf16,
+    int M, int N, int K,
+    const void*  SFA,
+    const void*  SFB,
+    float        alpha,
+    cudaStream_t stream)
+{
+  cutlass::Status status = pingpong::run_gemm_pingpong(
+      A_packed, B_packed, D_bf16, M, N, K, SFA, SFB, alpha, stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::fprintf(stderr,
+        "[fp4_w4a16_gemm_sm120_bf16out_pingpong] run FAIL "
+        "M=%d N=%d K=%d (status=%d); D output undefined\n",
+        M, N, K, static_cast<int>(status));
+  }
+}
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm120.cuh

@@ -0,0 +1,112 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// CUTLASS-based NVFP4 W4A16 GEMM for SM120a (RTX 5090 / Blackwell
+// consumer GeForce). Native block-scaled FP4 GEMM matching the Qwen3.6
+// NVFP4 ckpt schema (compressed-tensors `nvfp4-pack-quantized` format).
+//
+// Wraps NVIDIA's verified template from
+// third_party/cutlass/test/unit/gemm/device/sm120_blockscaled_tensorop_gemm/
+// sm120_bs_gemm_nvf4_nvf4_f32_bf16.cu — using `OpClassBlockScaledTensorOp`
+// + `nv_float4_t<float_e2m1_t>` + `float_ue4m3_t` group scales, BF16
+// output. SM_120 + SM_121 (RTX 5090 / 5080) gated.
+//
+// Schema (matches both A=act and B=weight after per-token NVFP4 quant):
+//   * elements    : 4-bit FP e2m1, packed two per byte
+//   * group scale : FP8 ue4m3, one scale per 16-element group
+//                   (`group_size = 16` per the ckpt config.json)
+//   * global scale: a single FP32 per tensor (fed via the epilogue's
+//                   alpha so we get D = sf_global_a * sf_global_b *
+//                   (A * B) with a single multiply instead of a
+//                   per-tile rescale)
+//
+// Caller responsibilities:
+//   * A_packed, B_packed are u8 arrays viewing FP4 e2m1 (2x packed).
+//     A is row-major (M, K/2 byte-pairs). B is column-major weight
+//     view; we accept the natural HF row-major (N, K/2) layout and
+//     reinterpret as ColumnMajor (K, N) — same memory.
+//   * SFA, SFB are FP8 ue4m3 with the CUTLASS Sm1xx blockscaled tile
+//     atom layout. The activation quantizer (`quantize_fp4_dynamic_*`)
+//     is responsible for emitting SFA in this layout; the weight
+//     loader does the same SFB transform once at load.
+//
+// Constraints (verified by `can_implement` at runtime):
+//   * K must be a multiple of 16 (group size).
+//   * Pointer alignments: A/B 16 bytes (32 FP4 elements), C/D 16 bytes
+//     (8 BF16 elements).
+//   * M unrestricted. We pick a tile shape by M (small-M variant
+//     coming once profiled — first cut uses the unit test's
+//     <128,128,256> for all M and lets CUTLASS handle padding).
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace gemm {
+
+// NVFP4 W4A16 GEMM, BF16 output, SM120a (RTX 5090).
+//
+//   A_packed  : (M, K/2)        u8  row-major   (FP4 e2m1, 2 per byte)
+//   B_packed  : (N, K/2)        u8  row-major   (FP4 e2m1, 2 per byte)
+//                                              — read as ColumnMajor (K, N)
+//   D_bf16    : (M, N)          bf16 row-major
+//   SFA       : (M, K/16)       e4m3 (CUTLASS blockscaled atom layout)
+//   SFB       : (N, K/16)       e4m3 (CUTLASS blockscaled atom layout)
+//   alpha     : fp32 scalar = act_global_scale * w_global_scale
+//
+// Stream-safe; per-shape arguments + workspace cached internally
+// (mirrors the FP8 sm_120 kernel).
+void fp4_w4a16_gemm_sm120_bf16out(
+    const void*  A_packed,    // (M, K/2)        u8
+    const void*  B_packed,    // (N, K/2)        u8
+    void*        D_bf16,      // (M, N)          bf16
+    int M, int N, int K,
+    const void*  SFA,         // (M, K/16)       e4m3 (Sm1xx blockscaled layout)
+    const void*  SFB,         // (N, K/16)       e4m3 (Sm1xx blockscaled layout)
+    float        alpha,       // = sf_global_a * sf_global_b
+    cudaStream_t stream);
+
+// Residual variant: D = alpha*(A*B) + C, C a per-element bf16 (M,N) addend.
+// Folds the post-GEMM residual add (o_proj/down) into the epilogue so the
+// following rms_norm reads one tensor (D) not two. Default tile (same as above).
+void fp4_w4a16_gemm_residual_sm120_bf16out(
+    const void*  A_packed,
+    const void*  B_packed,
+    const void*  C_residual,  // (M, N) bf16 row-major
+    void*        D_bf16,
+    int M, int N, int K,
+    const void*  SFA,
+    const void*  SFB,
+    float        alpha,
+    cudaStream_t stream);
+
+// Wide-N variant: TileShape <128, 256, 128>. For shapes with very
+// large N (lm_head N=248320, MLP gate/up N=17408) where the wider N
+// tile uses fewer waves and hits ~88%/66% peak BW vs ~64%/56% for
+// the default <128,128,256> tile. For small/medium N (<= 6144) the
+// default kernel is faster — caller dispatches by shape.
+void fp4_w4a16_gemm_sm120_bf16out_widen(
+    const void*  A_packed,
+    const void*  B_packed,
+    void*        D_bf16,
+    int M, int N, int K,
+    const void*  SFA,
+    const void*  SFB,
+    float        alpha,
+    cudaStream_t stream);
+
+// Same tile shape as the default kernel, but with
+// KernelTmaWarpSpecializedPingpong. Kept as an explicit opt-in variant so
+// callers can A/B schedule effects per shape without perturbing the default.
+void fp4_w4a16_gemm_sm120_bf16out_pingpong(
+    const void*  A_packed,
+    const void*  B_packed,
+    void*        D_bf16,
+    int M, int N, int K,
+    const void*  SFA,
+    const void*  SFB,
+    float        alpha,
+    cudaStream_t stream);
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/fp4_w4a4_mma_warpsplit_sm120.cu

@@ -0,0 +1,147 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Warp-split-K NVFP4 W4A4 M=1 GEMV for sm_120 — for the long-K / small-N
+// decode shapes (mlp_down K=17408, out_proj) where the single-warp full_n
+// kernel underfills the SMs. Instead of splitting K across BLOCKS (which
+// needs a cross-block fp32 reduce that is fragile under CUDA-graph replay),
+// this splits K across WARPS WITHIN one block: 8 N-cols/block, WARPS warps,
+// each warp streams K/WARPS, and the warp partials are summed in SHARED
+// MEMORY (intra-block) before the bf16 write. Single kernel, direct output,
+// no cross-kernel intermediate -> graph-replay safe. More warps/SM (occupancy)
+// + shorter per-warp streams give the same fill-the-SM win as block split-K.
+// Additive: new file + new entry point.
+//
+// Header: fp4_w4a4_mma_warpsplit_sm120.cuh.
+#include "gemm/fp4/fp4_w4a4_mma_warpsplit_sm120.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+#include "cute/arch/mma_sm120.hpp"
+#include "cutlass/numeric_types.h"
+
+namespace flash_rt {
+namespace gemm {
+namespace {
+
+using AtomType = cute::SM120::BLOCKSCALED::SM120_16x8x64_TN_VS<
+    cutlass::float_e2m1_t, cutlass::float_e2m1_t, float,
+    cutlass::float_ue4m3_t, 16>;
+
+__device__ __forceinline__ uint32_t fa(const uint8_t* s, int t0, int t1, int r) {
+  int ro = ((r & 1) ? (t1 + 8) : t1) * 32;
+  return *reinterpret_cast<const uint32_t*>(s + ro + t0 * 4 + ((r >> 1) & 1) * 16);
+}
+__device__ __forceinline__ uint32_t fb(const uint8_t* s, int t0, int t1, int r) {
+  return *reinterpret_cast<const uint32_t*>(s + t1 * 32 + t0 * 4 + r * 16);
+}
+__device__ __forceinline__ uint32_t fsa(const uint8_t* p, int u) {
+  return *reinterpret_cast<const uint32_t*>(p + u * 4);
+}
+__device__ __forceinline__ void cpa(uint8_t* d, const uint8_t* s) {
+  uint32_t i = __cvta_generic_to_shared(d);
+  asm volatile("cp.async.ca.shared.global.L2::128B [%0], [%1], 4;\n" :: "r"(i), "l"(s));
+}
+__device__ __forceinline__ void commit() { asm volatile("cp.async.commit_group;\n" ::); }
+template <int N> __device__ __forceinline__ void waitg() {
+  asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
+}
+
+template <int STAGES, int WARPS>
+__global__ void warpsplit_kernel(
+    const uint8_t* __restrict__ A, const uint8_t* __restrict__ B,
+    const uint8_t* __restrict__ SFA, const uint8_t* __restrict__ SFB,
+    __nv_bfloat16* __restrict__ D, float alpha, int N, int K) {
+  // per-warp pipeline buffers
+  __shared__ uint8_t sA[WARPS][STAGES][16 * 32];
+  __shared__ uint8_t sSFA[WARPS][STAGES][16 * 4];
+  __shared__ uint8_t sB[WARPS][STAGES][8 * 32];
+  __shared__ uint8_t sSFB[WARPS][STAGES][8 * 4];
+  __shared__ float s_red[WARPS][8];   // each warp's 8 col partials
+
+  int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31;
+  int my_n = blockIdx.x * 8;
+  const int KI = K / 64, KIw = KI / WARPS;   // K-tiles per warp
+  const int kt0 = warp * KIw;
+  const int KH = K / 2, ncs = (K / 16 + 3) / 4;
+  int t0 = lane & 3, t1 = lane >> 2, sau = (lane & 1) * 8 + (lane >> 2), sbu = lane >> 2;
+  float c0 = 0, c1 = 0, c2 = 0, c3 = 0;
+
+  uint8_t (*mA)[16 * 32] = sA[warp];
+  uint8_t (*mSFA)[16 * 4] = sSFA[warp];
+  uint8_t (*mB)[8 * 32] = sB[warp];
+  uint8_t (*mSFB)[8 * 4] = sSFB[warp];
+
+  if (lane >= 1 && lane < 16) {
+    #pragma unroll
+    for (int st = 0; st < STAGES; ++st) {
+      int4* av = reinterpret_cast<int4*>(mA[st]); int4 z{0, 0, 0, 0};
+      av[lane * 2] = z; av[lane * 2 + 1] = z;
+    }
+    if (lane < 4) for (int st = 0; st < STAGES; ++st)
+      for (int i = 4 + lane; i < 64; i += 4) mSFA[st][i] = 0;
+  }
+  auto ld = [&](int bf, int kt) {
+    int bo = kt * 32;
+    if (lane < 8) cpa(mA[bf] + lane * 4, A + bo + lane * 4);
+    if (lane == 0) cpa(mSFA[bf], SFA + kt * 512);
+    for (int c = 0; c < 2; ++c) { int ch = lane + c * 32, col = ch >> 3, off = ch & 7;
+      cpa(mB[bf] + ch * 4, B + (my_n + col) * KH + bo + off * 4); }
+    if (lane < 8) { int col = my_n + lane, rb = col >> 7, ri = col & 127;
+      int si = rb * ncs + kt, ib = (ri & 31) * 16 + ((ri >> 5) & 3) * 4;
+      cpa(mSFB[bf] + lane * 4, SFB + si * 512 + ib); }
+  };
+  #pragma unroll
+  for (int st = 0; st < STAGES - 1; ++st) { if (st < KIw) ld(st, kt0 + st); commit(); }
+  for (int j = 0; j < KIw; ++j) {
+    int cb = j % STAGES, jp = j + STAGES - 1;
+    if (jp < KIw) ld(jp % STAGES, kt0 + jp);
+    commit(); waitg<STAGES - 1>(); __syncwarp();
+    uint32_t a0 = fa(mA[cb], t0, t1, 0), a1 = fa(mA[cb], t0, t1, 1);
+    uint32_t a2 = fa(mA[cb], t0, t1, 2), a3 = fa(mA[cb], t0, t1, 3);
+    uint32_t b0 = fb(mB[cb], t0, t1, 0), b1 = fb(mB[cb], t0, t1, 1);
+    uint32_t sfa = fsa(mSFA[cb], sau), sfb = fsa(mSFB[cb], sbu);
+    float d0, d1, d2, d3;
+    AtomType::fma(d0, d1, d2, d3, a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, sfa, sfb);
+    c0 = d0; c1 = d1; c2 = d2; c3 = d3;
+  }
+  // each warp: lanes 0..3 hold row-0 partials c0 (col 2r) / c1 (col 2r+1)
+  int q = lane >> 2, r = lane & 3;
+  if (q == 0) { s_red[warp][r * 2] = c0; s_red[warp][r * 2 + 1] = c1; }
+  __syncthreads();
+  // warp 0 sums the WARPS partials per col and writes the bf16 output
+  if (warp == 0 && lane < 8) {
+    float acc = 0.f;
+    #pragma unroll
+    for (int w = 0; w < WARPS; ++w) acc += s_red[w][lane];
+    int col = my_n + lane;
+    if (col < N) D[col] = __float2bfloat16(acc * alpha);
+  }
+}
+
+}  // namespace
+
+int fp4_w4a4_mma_sm120_warpsplit_bf16out(
+    const void* A_packed, const void* B_packed, void* D_bf16, int N, int K,
+    const void* SFA, const void* SFB, float alpha, int warps, int stages,
+    cudaStream_t stream) {
+  if (!A_packed || !B_packed || !D_bf16 || !SFA || !SFB) return 1;
+  if (K <= 0 || (K % 64) != 0 || ((K / 64) % warps) != 0) return 2;
+  if (N <= 0 || (N % 8) != 0) return 3;
+  dim3 grid(N / 8);
+  auto a = reinterpret_cast<const uint8_t*>(A_packed);
+  auto b = reinterpret_cast<const uint8_t*>(B_packed);
+  auto sa = reinterpret_cast<const uint8_t*>(SFA);
+  auto sb = reinterpret_cast<const uint8_t*>(SFB);
+  auto d = reinterpret_cast<__nv_bfloat16*>(D_bf16);
+  #define WS_L(ST, WP) warpsplit_kernel<ST, WP><<<grid, WP * 32, 0, stream>>>(a, b, sa, sb, d, alpha, N, K)
+  if (warps == 2) { if (stages == 3) WS_L(3, 2); else if (stages == 4) WS_L(4, 2); else if (stages == 6) WS_L(6, 2); else return 5; }
+  else if (warps == 4) { if (stages == 3) WS_L(3, 4); else if (stages == 4) WS_L(4, 4); else if (stages == 6) WS_L(6, 4); else return 5; }
+  else if (warps == 8) { if (stages == 3) WS_L(3, 8); else if (stages == 4) WS_L(4, 8); else return 5; }
+  else return 6;
+  return 0;
+}
+
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/fp4_w4a4_mma_warpsplit_sm120.cuh

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Warp-split-K NVFP4 W4A4 M=1 GEMV for sm_120: 8 N-cols/block, `warps` warps
+// each streaming K/warps, partials summed in shared memory (intra-block) then
+// written bf16. Graph-replay safe (no cross-block/cross-kernel intermediate).
+// For long-K/small-N decode shapes (mlp_down, out_proj) the full_n kernel
+// underfills. Additive.
+#pragma once
+#include <cuda_runtime.h>
+namespace flash_rt {
+namespace gemm {
+// A_packed (K/2,), B_packed (N,K/2), D_bf16 (N,). SFA (K/16,), SFB (N,K/16)
+// swizzled. warps in {2,4,8}, stages in {3,4,6}. N%8==0, K%64==0,
+// (K/64)%warps==0. Returns 0 on success.
+int fp4_w4a4_mma_sm120_warpsplit_bf16out(
+    const void* A_packed, const void* B_packed, void* D_bf16, int N, int K,
+    const void* SFA, const void* SFB, float alpha, int warps, int stages,
+    cudaStream_t stream);
+}  // namespace gemm
+}  // namespace flash_rt

csrc/gemm/fp4/sm110_dispatch.cu

@@ -0,0 +1,50 @@
+#include "gemm/fp4/sm110_dispatch.cuh"
+
+#include "gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm100.cuh"
+#include "gemm/fp4/cutlass_fp4_gemm_bias_bf16_sm100.cuh"
+#include "quantize/quantize_fp4_sfa_bf16.cuh"
+
+namespace flash_rt::hub {
+namespace {
+
+void launch_sm110(
+    const void* a,
+    const void* b,
+    void* out,
+    int m,
+    int n,
+    int k,
+    const void* sfa,
+    const void* sfb,
+    float alpha,
+    int variant,
+    cudaStream_t stream) {
+  if (variant == 1) {
+    gemm::fp4_w4a16_gemm_sm100_bf16out_widen(
+        a, b, out, m, n, k, sfa, sfb, alpha, stream);
+  } else if (variant == 2) {
+    gemm::fp4_w4a16_gemm_sm100_bf16out_pingpong(
+        a, b, out, m, n, k, sfa, sfb, alpha, stream);
+  } else {
+    gemm::fp4_w4a16_gemm_sm100_bf16out(
+        a, b, out, m, n, k, sfa, sfb, alpha, stream);
+  }
+}
+
+struct Sm110DispatchRegistration {
+  Sm110DispatchRegistration() {
+    sm110_gemm_dispatch = &launch_sm110;
+    sm110_gemm_bias_dispatch = &fp4::cutlass_fp4_gemm_bias_bf16;
+    sm110_gemm_bias_residual_dispatch =
+        &fp4::cutlass_fp4_gemm_bias_res_bf16;
+    sm110_gemm_bias_gelu_fp4_dispatch =
+        &fp4::cutlass_fp4_gemm_bias_gelu_fp4out_bf16;
+    sm110_quantize_bf16_dispatch =
+        &fp4::quantize_fp4_dynamic_sfa_bf16_vec;
+  }
+};
+
+Sm110DispatchRegistration registration;
+
+}  // namespace
+}  // namespace flash_rt::hub

csrc/gemm/fp4/sm110_dispatch.cuh

@@ -0,0 +1,45 @@
+#pragma once
+
+#include <cuda_runtime_api.h>
+
+namespace flash_rt::hub {
+
+using Sm110GemmDispatch = void (*)(
+    const void* a,
+    const void* b,
+    void* out,
+    int m,
+    int n,
+    int k,
+    const void* sfa,
+    const void* sfb,
+    float alpha,
+    int variant,
+    cudaStream_t stream);
+
+using Sm110GemmBiasDispatch = int (*)(
+    const void* a, const void* sfa, const void* b, const void* sfb,
+    const void* bias, void* out, int m, int n, int k,
+    cudaStream_t stream);
+
+using Sm110GemmBiasResidualDispatch = int (*)(
+    const void* a, const void* sfa, const void* b, const void* sfb,
+    const void* bias, const void* residual, void* out,
+    int m, int n, int k, cudaStream_t stream);
+
+using Sm110GemmBiasGeluFp4Dispatch = int (*)(
+    const void* a, const void* sfa, const void* b, const void* sfb,
+    const void* bias, void* out_packed, void* out_sfa,
+    int m, int n, int k, cudaStream_t stream);
+
+using Sm110QuantizeBf16Dispatch = int (*)(
+    const void* x, void* packed, void* sfa, int rows, int dim,
+    bool is_sfb, cudaStream_t stream);
+
+extern Sm110GemmDispatch sm110_gemm_dispatch;
+extern Sm110GemmBiasDispatch sm110_gemm_bias_dispatch;
+extern Sm110GemmBiasResidualDispatch sm110_gemm_bias_residual_dispatch;
+extern Sm110GemmBiasGeluFp4Dispatch sm110_gemm_bias_gelu_fp4_dispatch;
+extern Sm110QuantizeBf16Dispatch sm110_quantize_bf16_dispatch;
+
+}  // namespace flash_rt::hub

csrc/quantize/quantize_fp4_sfa.cu

@@ -0,0 +1,194 @@
+// ============================================================================
+//  Fused FP4 quantize + CUTLASS SFA/SFB tile-interleaved scale write.
+//
+//  Implementation = kernel_quantize_fp4 (quantize_fp4_dynamic.cu) with the
+//  scale-store address replaced by the CUTLASS layout functor. Packed fp4
+//  elements layout is UNCHANGED (still linear [N, D/2]), only the scale
+//  byte goes to a different location.
+// ============================================================================
+#include "quantize_fp4_sfa.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+
+#ifndef CUTLASS_ARCH_MMA_SM100_SUPPORTED
+#  define CUTLASS_ARCH_MMA_SM100_SUPPORTED 1
+#endif
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(__CUDA_ARCH__)
+#  include "cutlass/cutlass.h"
+#  include "cutlass/detail/sm100_blockscaled_layout.hpp"
+#  include "cute/tensor.hpp"
+#  define FV_HAVE_CUTLASS 1
+#else
+#  define FV_HAVE_CUTLASS 0
+#endif
+
+namespace flash_rt {
+namespace fp4 {
+
+#if FV_HAVE_CUTLASS
+
+using Cfg = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+// ── Device helpers (duplicated locally to stay additive — not linking against
+//    quantize_fp4_dynamic.cu so we don't risk ODR issues). Identical logic. ──
+__device__ __forceinline__ uint8_t fp32_to_e2m1_sfa(float x) {
+    uint8_t sign = (x < 0.f) ? 0x8u : 0x0u;
+    float ax = fabsf(x);
+    uint8_t mant;
+    if      (ax <= 0.25f) mant = 0u;
+    else if (ax <= 0.75f) mant = 1u;
+    else if (ax <= 1.25f) mant = 2u;
+    else if (ax <= 1.75f) mant = 3u;
+    else if (ax <= 2.5f)  mant = 4u;
+    else if (ax <= 3.5f)  mant = 5u;
+    else if (ax <= 5.0f)  mant = 6u;
+    else                  mant = 7u;
+    return sign | mant;
+}
+
+__device__ __forceinline__ __nv_fp8_e4m3 quantize_ue4m3_sfa(float x) {
+    float v = fmaxf(x, 0.f);
+    return __nv_fp8_e4m3(v);
+}
+
+__device__ __forceinline__ float dequantize_ue4m3_sfa(__nv_fp8_e4m3 s) {
+    return static_cast<float>(s);
+}
+
+__device__ __forceinline__ float input_to_float(__half value) {
+    return __half2float(value);
+}
+
+__device__ __forceinline__ float input_to_float(__nv_bfloat16 value) {
+    return __bfloat162float(value);
+}
+
+// ── Fused kernel ──
+// One thread per (row, 16-element block). Scale byte goes to
+// dst_sfa[layout(row, block_idx*16, 0)].
+template <typename Input, class LayoutSF>
+__global__ void kernel_quantize_fp4_sfa(
+    const Input* __restrict__ src,
+    uint8_t* __restrict__ dst_packed,
+    uint8_t* __restrict__ dst_sfa,   // raw byte view of the CUTLASS SFA/SFB buffer
+    LayoutSF layout,
+    int N, int D) {
+  const int block_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  const int row       = blockIdx.y;
+  const int n_blocks  = D / 16;
+  if (row >= N || block_idx >= n_blocks) return;
+
+  const int base = row * D + block_idx * 16;
+  float vals[16];
+  float amax = 0.f;
+  #pragma unroll
+  for (int i = 0; i < 16; ++i) {
+    vals[i] = input_to_float(src[base + i]);
+    float a = fabsf(vals[i]);
+    if (a > amax) amax = a;
+  }
+
+  float desired = amax / 6.f;
+  if (desired < 1e-12f) desired = 1e-12f;
+  __nv_fp8_e4m3 bs_q = quantize_ue4m3_sfa(desired);
+  float bs_dq        = dequantize_ue4m3_sfa(bs_q);
+
+  // ── CORE FUSION: direct SFA tile-layout write ──
+  // LayoutSF maps (row, k, L=0) → byte offset. k is the full-K coordinate;
+  // SFVecSize=16 is baked in so any k in [block*16, block*16+15] hits the
+  // same offset. Use block_idx*16 (same convention as reshape_scales_sfa.cu).
+  int sfa_off = layout(row, block_idx * 16, 0);
+  dst_sfa[sfa_off] = *reinterpret_cast<uint8_t*>(&bs_q);
+
+  // Packed fp4 elements: layout unchanged.
+  const int out_base = row * (D / 2) + block_idx * 8;
+  const float inv_bs = 1.f / bs_dq;
+  #pragma unroll
+  for (int p = 0; p < 8; ++p) {
+    float v_lo = vals[2 * p    ] * inv_bs;
+    float v_hi = vals[2 * p + 1] * inv_bs;
+    uint8_t lo = fp32_to_e2m1_sfa(v_lo);
+    uint8_t hi = fp32_to_e2m1_sfa(v_hi);
+    dst_packed[out_base + p] = lo | (hi << 4);
+  }
+}
+
+#endif  // FV_HAVE_CUTLASS
+
+int quantize_fp4_dynamic_sfa_fp16(
+    const void* src_fp16, void* dst_packed, void* dst_sfa,
+    int N, int D, bool is_sfb, cudaStream_t stream) {
+#if FV_HAVE_CUTLASS
+  if (D % 16 != 0) return -1;
+  const int n_blocks = D / 16;
+  const int threads = 128;
+  dim3 grid((n_blocks + threads - 1) / threads, N);
+  dim3 block(threads);
+
+  // Shape: SFA uses (M=N, 1, K=D, L=1); SFB uses (1, N=N, K=D, L=1).
+  auto shape = cute::make_shape(
+      is_sfb ? 1 : N,
+      is_sfb ? N : 1,
+      D, 1);
+
+  if (is_sfb) {
+    auto layout = Cfg::tile_atom_to_shape_SFB(shape);
+    kernel_quantize_fp4_sfa<__half><<<grid, block, 0, stream>>>(
+        reinterpret_cast<const __half*>(src_fp16),
+        reinterpret_cast<uint8_t*>(dst_packed),
+        reinterpret_cast<uint8_t*>(dst_sfa),
+        layout, N, D);
+  } else {
+    auto layout = Cfg::tile_atom_to_shape_SFA(shape);
+    kernel_quantize_fp4_sfa<__half><<<grid, block, 0, stream>>>(
+        reinterpret_cast<const __half*>(src_fp16),
+        reinterpret_cast<uint8_t*>(dst_packed),
+        reinterpret_cast<uint8_t*>(dst_sfa),
+        layout, N, D);
+  }
+  cudaError_t e = cudaGetLastError();
+  return (e == cudaSuccess) ? 0 : -static_cast<int>(e);
+#else
+  (void)src_fp16; (void)dst_packed; (void)dst_sfa;
+  (void)N; (void)D; (void)is_sfb; (void)stream;
+  return -2;
+#endif
+}
+
+int quantize_fp4_dynamic_sfa_bf16(
+    const void* src_bf16, void* dst_packed, void* dst_sfa,
+    int N, int D, bool is_sfb, cudaStream_t stream) {
+#if FV_HAVE_CUTLASS
+  if (D % 16 != 0) return -1;
+  const int n_blocks = D / 16;
+  const int threads = 128;
+  dim3 grid((n_blocks + threads - 1) / threads, N);
+  dim3 block(threads);
+  auto shape = cute::make_shape(is_sfb ? 1 : N, is_sfb ? N : 1, D, 1);
+
+  if (is_sfb) {
+    auto layout = Cfg::tile_atom_to_shape_SFB(shape);
+    kernel_quantize_fp4_sfa<__nv_bfloat16><<<grid, block, 0, stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(src_bf16),
+        reinterpret_cast<uint8_t*>(dst_packed),
+        reinterpret_cast<uint8_t*>(dst_sfa), layout, N, D);
+  } else {
+    auto layout = Cfg::tile_atom_to_shape_SFA(shape);
+    kernel_quantize_fp4_sfa<__nv_bfloat16><<<grid, block, 0, stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(src_bf16),
+        reinterpret_cast<uint8_t*>(dst_packed),
+        reinterpret_cast<uint8_t*>(dst_sfa), layout, N, D);
+  }
+  cudaError_t e = cudaGetLastError();
+  return (e == cudaSuccess) ? 0 : -static_cast<int>(e);
+#else
+  (void)src_bf16; (void)dst_packed; (void)dst_sfa;
+  (void)N; (void)D; (void)is_sfb; (void)stream;
+  return -2;
+#endif
+}
+
+}  // namespace fp4
+}  // namespace flash_rt

csrc/quantize/quantize_fp4_sfa.cuh

@@ -0,0 +1,41 @@
+// ============================================================================
+//  FlashRT — fused (FP4 quantize + CUTLASS SFA/SFB tile-interleave) kernel.
+//
+//  Equivalent to:
+//      quantize_fp4_dynamic_fp16(src, packed, linear_scales, N, D)
+//      reshape_linear_scales_to_sfa(linear_scales, sfa, N, D, is_sfb)
+//  in a SINGLE kernel launch. Scale byte is written directly to the CUTLASS
+//  tile-interleaved offset — linear_scales intermediate buffer is gone.
+//
+//  Additive: does NOT modify quantize_fp4_dynamic.* or reshape_scales_sfa.*.
+//  Both remain callable for existing paths.
+// ============================================================================
+#pragma once
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace fp4 {
+
+// fp16 [N, D] → packed [N, D/2] (e2m1) + SFA/SFB tile-interleaved UE4M3 scales.
+//   is_sfb = false  → SFA layout (use for A = activation, shape [M=N, K=D])
+//   is_sfb = true   → SFB layout (use for B = weight,     shape [N=N, K=D])
+// Returns 0 on success.
+int quantize_fp4_dynamic_sfa_fp16(
+    const void* src_fp16,
+    void* dst_packed,
+    void* dst_sfa,
+    int N, int D, bool is_sfb,
+    cudaStream_t stream);
+
+// BF16 [N, D] -> the exact same packed E2M1 + CUTLASS SFA/SFB layout as the
+// FP16 entry. This avoids a standalone BF16-to-FP16 conversion in decode
+// pipelines whose activations are already BF16.
+int quantize_fp4_dynamic_sfa_bf16(
+    const void* src_bf16,
+    void* dst_packed,
+    void* dst_sfa,
+    int N, int D, bool is_sfb,
+    cudaStream_t stream);
+
+}  // namespace fp4
+}  // namespace flash_rt

csrc/quantize/quantize_fp4_sfa_bf16.cu

@@ -0,0 +1,142 @@
+// ============================================================================
+//  bf16-input vectorized fused FP4 quantize + CUTLASS SFA/SFB scales.
+//
+//  Same per-block scale selection and e2m1 rounding as
+//  quantize_fp4_dynamic_sfa_fp16_vec, with bf16 source elements. Each
+//  thread quantizes one 16-element block: two 16-byte loads, one 8-byte
+//  packed store, one SFA byte at the tile-interleaved offset.
+// ============================================================================
+#include "quantize_fp4_sfa_bf16.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(__CUDA_ARCH__)
+#  include "cutlass/cutlass.h"
+#  include "cutlass/detail/sm100_blockscaled_layout.hpp"
+#  include "cute/tensor.hpp"
+#  define FV_HAVE_CUTLASS 1
+#else
+#  define FV_HAVE_CUTLASS 0
+#endif
+
+namespace flash_rt {
+namespace fp4 {
+
+#if FV_HAVE_CUTLASS
+
+namespace {
+
+using CfgVecB = cutlass::detail::Sm1xxBlockScaledConfig<16>;
+
+__device__ __forceinline__ uint8_t fp32_to_e2m1_bvec(float x) {
+    uint8_t sign = (x < 0.f) ? 0x8u : 0x0u;
+    float ax = fabsf(x);
+    uint8_t mant;
+    if      (ax <= 0.25f) mant = 0u;
+    else if (ax <= 0.75f) mant = 1u;
+    else if (ax <= 1.25f) mant = 2u;
+    else if (ax <= 1.75f) mant = 3u;
+    else if (ax <= 2.5f)  mant = 4u;
+    else if (ax <= 3.5f)  mant = 5u;
+    else if (ax <= 5.0f)  mant = 6u;
+    else                  mant = 7u;
+    return sign | mant;
+}
+
+template <class LayoutSF>
+__global__ void kernel_quantize_fp4_sfa_bf16_vec(
+    const int4* __restrict__ src,      // bf16 [N, D] as int4 (8 elements)
+    uint2* __restrict__ dst_packed,    // [N, D/2] bytes as uint2 (1 block)
+    uint8_t* __restrict__ dst_sfa,
+    LayoutSF layout,
+    int N, int D8) {                   // D8 = D / 8 int4 chunks per row
+  const int block_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  const int row = blockIdx.y;
+  const int n_blocks = D8 >> 1;        // 16 elements per block
+  if (row >= N || block_idx >= n_blocks) return;
+
+  const int4 raw0 = src[row * D8 + 2 * block_idx];
+  const int4 raw1 = src[row * D8 + 2 * block_idx + 1];
+  const __nv_bfloat16* h0 = reinterpret_cast<const __nv_bfloat16*>(&raw0);
+  const __nv_bfloat16* h1 = reinterpret_cast<const __nv_bfloat16*>(&raw1);
+
+  float vals[16];
+  float amax = 0.f;
+  #pragma unroll
+  for (int i = 0; i < 8; ++i) {
+    vals[i] = __bfloat162float(h0[i]);
+    vals[8 + i] = __bfloat162float(h1[i]);
+  }
+  #pragma unroll
+  for (int i = 0; i < 16; ++i) {
+    const float a = fabsf(vals[i]);
+    if (a > amax) amax = a;
+  }
+
+  float desired = amax / 6.f;
+  if (desired < 1e-12f) desired = 1e-12f;
+  __nv_fp8_e4m3 bs_q = __nv_fp8_e4m3(fmaxf(desired, 0.f));
+  const float bs_dq = static_cast<float>(bs_q);
+
+  dst_sfa[layout(row, block_idx * 16, 0)] =
+      *reinterpret_cast<uint8_t*>(&bs_q);
+
+  const float inv_bs = 1.f / bs_dq;
+  uint2 out;
+  uint8_t* ob = reinterpret_cast<uint8_t*>(&out);
+  #pragma unroll
+  for (int p = 0; p < 8; ++p) {
+    const uint8_t lo = fp32_to_e2m1_bvec(vals[2 * p] * inv_bs);
+    const uint8_t hi = fp32_to_e2m1_bvec(vals[2 * p + 1] * inv_bs);
+    ob[p] = static_cast<uint8_t>(lo | (hi << 4));
+  }
+  dst_packed[row * n_blocks + block_idx] = out;
+}
+
+}  // namespace
+
+#endif  // FV_HAVE_CUTLASS
+
+int quantize_fp4_dynamic_sfa_bf16_vec(
+    const void* src_bf16, void* dst_packed, void* dst_sfa,
+    int N, int D, bool is_sfb, cudaStream_t stream) {
+#if FV_HAVE_CUTLASS
+  if (D % 16 != 0) return -1;
+  if ((reinterpret_cast<uintptr_t>(src_bf16) & 15) ||
+      (reinterpret_cast<uintptr_t>(dst_packed) & 7)) return -1;
+  const int n_blocks = D / 16;
+  const int threads = 128;
+  dim3 grid((n_blocks + threads - 1) / threads, N);
+
+  auto shape = cute::make_shape(
+      is_sfb ? 1 : N,
+      is_sfb ? N : 1,
+      D, 1);
+
+  if (is_sfb) {
+    auto layout = CfgVecB::tile_atom_to_shape_SFB(shape);
+    kernel_quantize_fp4_sfa_bf16_vec<<<grid, threads, 0, stream>>>(
+        reinterpret_cast<const int4*>(src_bf16),
+        reinterpret_cast<uint2*>(dst_packed),
+        reinterpret_cast<uint8_t*>(dst_sfa),
+        layout, N, D >> 3);
+  } else {
+    auto layout = CfgVecB::tile_atom_to_shape_SFA(shape);
+    kernel_quantize_fp4_sfa_bf16_vec<<<grid, threads, 0, stream>>>(
+        reinterpret_cast<const int4*>(src_bf16),
+        reinterpret_cast<uint2*>(dst_packed),
+        reinterpret_cast<uint8_t*>(dst_sfa),
+        layout, N, D >> 3);
+  }
+  const cudaError_t e = cudaGetLastError();
+  return (e == cudaSuccess) ? 0 : -static_cast<int>(e);
+#else
+  (void)src_bf16; (void)dst_packed; (void)dst_sfa;
+  (void)N; (void)D; (void)is_sfb; (void)stream;
+  return -2;
+#endif
+}
+
+}  // namespace fp4
+}  // namespace flash_rt

csrc/quantize/quantize_fp4_sfa_bf16.cuh

@@ -0,0 +1,24 @@
+// ============================================================================
+//  FlashRT — bf16-input fused NVFP4 quantize + CUTLASS SFA/SFB scale write.
+//
+//  bf16 companion of quantize_fp4_dynamic_sfa_fp16 for pipelines whose
+//  activations are bf16 (GR00T N1.7 DiT). Additive: new symbols only.
+// ============================================================================
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace fp4 {
+
+// Quantize a bf16 [N, D] row-major tensor to packed e2m1 [N, D/2] plus
+// UE4M3 per-16-element scales written directly at the CUTLASS
+// tile-interleaved SFA/SFB offsets. Vectorized (16-byte loads, 8-byte
+// packed stores). Returns 0 on success, -1 on unsupported shape or
+// misaligned buffers, -2 when built without CUTLASS.
+int quantize_fp4_dynamic_sfa_bf16_vec(
+    const void* src_bf16, void* dst_packed, void* dst_sfa,
+    int N, int D, bool is_sfb, cudaStream_t stream);
+
+}  // namespace fp4
+}  // namespace flash_rt

examples/README.md

@@ -0,0 +1,10 @@
+# fp4-gemm Examples
+
+These examples show direct Hub-style usage of `flashrt/fp4-gemm`.
+
+```bash
+python fp4-gemm/examples/fp4_gemm_linear.py
+```
+
+The quantization helper is included for validation and small examples. In a
+runtime, weights should normally be prepacked and loaded as FP4/SFA/SFB buffers.

examples/fp4_gemm_linear.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+"""Minimal Hub-style call for flashrt/fp4-gemm."""
+
+from __future__ import annotations
+
+import torch
+from kernels import get_kernel
+
+
+def main() -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+
+    ops = get_kernel("flashrt/fp4-gemm", version=1, trust_remote_code=True)
+
+    x = torch.randn((32, 256), device="cuda", dtype=torch.float16)
+    w = torch.randn((512, 256), device="cuda", dtype=torch.float16)
+
+    a_packed, sfa = ops.quantize_fp4_sfa_fp16(x, is_sfb=False)
+    b_packed, sfb = ops.quantize_fp4_sfa_fp16(w, is_sfb=True)
+    y = ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0)
+
+    print("a_packed", tuple(a_packed.shape), a_packed.dtype)
+    print("b_packed", tuple(b_packed.shape), b_packed.dtype)
+    print("output", tuple(y.shape), y.dtype)
+
+
+if __name__ == "__main__":
+    main()

flake.nix

@@ -0,0 +1,18 @@
+{
+  description = "Flake for FlashRT FP4 GEMM kernels";
+
+  inputs = {
+    # huggingface/kernels#741: CUTLASS 4.4.2 plus the corrected 4.5.2 hash.
+    kernel-builder.url = "github:huggingface/kernels/870e825d881664e39f9287a27a74ef63ff3c545e";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

tests/test_fp4_gemm.py

@@ -0,0 +1,668 @@
+#!/usr/bin/env python3
+"""Correctness tests for fp4-gemm."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import os
+import sys
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import torch
+
+
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "fp4-gemm"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
+)
+DEFAULT_CUTLASS_INCLUDE = (
+    ROOT.parent
+    / "flashrt_pr31_review"
+    / "third_party"
+    / "cutlass"
+    / "include"
+)
+
+SHAPES = {
+    "small_m16_n128_k128": (16, 128, 128),
+    "small_m32_n256_k256": (32, 256, 256),
+    "mlp_tile_m64_n512_k512": (64, 512, 512),
+}
+
+SM110_SHAPES = {
+    "pi05_action_gate_up": (51, 16384, 2048),
+    "pi05_action_down": (51, 2048, 8192),
+    "groot_n17_dit_qkv": (41, 4608, 1536),
+    "groot_n17_dit_ffn_up": (41, 6144, 1536),
+    "groot_n17_dit_ffn_down": (41, 1536, 6144),
+    "groot_legacy_dit_qkv": (51, 4608, 1536),
+    "groot_backbone_gate_up": (277, 16384, 2048),
+    "cosmos_edge_action": (64, 9216, 2048),
+    "lingbot_action_gate_up": (105, 16384, 2048),
+}
+
+EPILOGUE_SHAPES = {
+    "epilogue_tile": (64, 512, 512),
+    "motus_up": (360, 14336, 3072),
+    "motus_down": (360, 3072, 14336),
+}
+
+MODES = {
+    "smoke": ["small_m16_n128_k128"],
+    "full": list(SHAPES),
+    "thor-models": list(SM110_SHAPES),
+}
+
+
+@dataclass
+class Metrics:
+    shape: str
+    M: int
+    N: int
+    K: int
+    workload: str
+    variant: int | None
+    max_abs: float
+    mean_abs: float
+    p99_abs: float
+    cosine: float
+    passed: bool
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    @staticmethod
+    def sfa_size_bytes(rows: int, dim: int) -> int:
+        n_blocks = dim // 16
+        n_row_super = (rows + 127) // 128
+        n_col_super = (n_blocks + 3) // 4
+        return n_row_super * n_col_super * 512
+
+    def alloc_fp4(self, rows: int, dim: int):
+        return (
+            torch.empty((rows, dim // 2), device="cuda", dtype=torch.uint8),
+            torch.empty((self.sfa_size_bytes(rows, dim),), device="cuda", dtype=torch.uint8),
+        )
+
+    def quantize_fp4_sfa_fp16(self, x, packed, sfa, is_sfb=False):
+        self._ops.quantize_fp4_sfa_fp16(x, packed, sfa, bool(is_sfb))
+
+    def quantize_fp4_sfa_bf16(self, x, packed, sfa, is_sfb=False):
+        self._ops.quantize_fp4_sfa_bf16(x, packed, sfa, bool(is_sfb))
+
+    def dequantize_fp4_sfa_fp16(self, packed, sfa, out, is_sfb=False):
+        self._ops.dequantize_fp4_sfa_fp16(packed, sfa, out, bool(is_sfb))
+
+    def nvfp4_gemm_bf16(self, a, b, sfa, sfb, out, alpha=1.0, variant=0):
+        self._ops.nvfp4_gemm_bf16(a, b, sfa, sfb, out, float(alpha), int(variant))
+
+    def nvfp4_gemm_bias_bf16(self, a, b, sfa, sfb, bias, out):
+        self._ops.nvfp4_gemm_bias_bf16(a, b, sfa, sfb, bias, out)
+
+    def nvfp4_gemm_bias_residual_bf16(
+        self, a, b, sfa, sfb, bias, residual, out
+    ):
+        self._ops.nvfp4_gemm_bias_residual_bf16(
+            a, b, sfa, sfb, bias, residual, out
+        )
+
+    def nvfp4_gemm_residual_bf16(self, a, b, sfa, sfb, residual, out, alpha=1.0):
+        self._ops.nvfp4_gemm_residual_bf16(
+            a, b, sfa, sfb, residual, out, float(alpha)
+        )
+
+    def nvfp4_gemm_bias_gelu_bf16(self, a, b, sfa, sfb, bias, out, alpha=1.0):
+        self._ops.nvfp4_gemm_bias_gelu_bf16(
+            a, b, sfa, sfb, bias, out, float(alpha)
+        )
+
+    def nvfp4_gemm_bias_gelu_nvfp4(
+        self, a, b, sfa, sfb, bias, out_packed, out_sfa, alpha=1.0
+    ):
+        self._ops.nvfp4_gemm_bias_gelu_nvfp4(
+            a, b, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
+        )
+
+    def nvfp4_gemm_streamk_bf16(self, a, b, sfa, sfb, out, alpha=1.0):
+        self._ops.nvfp4_gemm_streamk_bf16(a, b, sfa, sfb, out, float(alpha))
+
+    def nvfp4_gemm_streamk_bias_bf16(
+        self, a, b, sfa, sfb, bias, out, alpha=1.0
+    ):
+        self._ops.nvfp4_gemm_streamk_bias_bf16(
+            a, b, sfa, sfb, bias, out, float(alpha)
+        )
+
+
+class InstalledOps:
+    """Adapt the public return-value API to the in-place test interface."""
+
+    def __init__(self, module) -> None:
+        self._module = module
+
+    def sfa_size_bytes(self, rows: int, dim: int) -> int:
+        return int(self._module.sfa_size_bytes(rows, dim))
+
+    def alloc_fp4(self, rows: int, dim: int):
+        return (
+            torch.empty((rows, dim // 2), device="cuda", dtype=torch.uint8),
+            torch.empty(
+                (self._module.sfa_size_bytes(rows, dim),),
+                device="cuda",
+                dtype=torch.uint8,
+            ),
+        )
+
+    def quantize_fp4_sfa_fp16(self, x, packed, sfa, is_sfb=False):
+        self._module.quantize_fp4_sfa_fp16(
+            x, packed=packed, sfa=sfa, is_sfb=bool(is_sfb)
+        )
+
+    def quantize_fp4_sfa_bf16(self, x, packed, sfa, is_sfb=False):
+        self._module.quantize_fp4_sfa_bf16(
+            x, packed=packed, sfa=sfa, is_sfb=bool(is_sfb)
+        )
+
+    def dequantize_fp4_sfa_fp16(self, packed, sfa, out, is_sfb=False):
+        self._module.dequantize_fp4_sfa_fp16(
+            packed, sfa, out=out, is_sfb=bool(is_sfb)
+        )
+
+    def nvfp4_gemm_bf16(self, a, b, sfa, sfb, out, alpha=1.0, variant=0):
+        self._module.nvfp4_gemm_bf16(
+            a,
+            b,
+            sfa,
+            sfb,
+            alpha=float(alpha),
+            out=out,
+            variant=int(variant),
+        )
+
+    def nvfp4_gemm_bias_bf16(self, a, b, sfa, sfb, bias, out):
+        self._module.nvfp4_gemm_bias_bf16(
+            a, b, sfa, sfb, bias, out=out
+        )
+
+    def nvfp4_gemm_bias_residual_bf16(
+        self, a, b, sfa, sfb, bias, residual, out
+    ):
+        self._module.nvfp4_gemm_bias_residual_bf16(
+            a, b, sfa, sfb, bias, residual, out=out
+        )
+
+    def nvfp4_gemm_residual_bf16(self, a, b, sfa, sfb, residual, out, alpha=1.0):
+        self._module.nvfp4_gemm_residual_bf16(
+            a, b, sfa, sfb, residual, alpha=float(alpha), out=out
+        )
+
+    def nvfp4_gemm_bias_gelu_bf16(self, a, b, sfa, sfb, bias, out, alpha=1.0):
+        self._module.nvfp4_gemm_bias_gelu_bf16(
+            a, b, sfa, sfb, bias, alpha=float(alpha), out=out
+        )
+
+    def nvfp4_gemm_bias_gelu_nvfp4(
+        self, a, b, sfa, sfb, bias, out_packed, out_sfa, alpha=1.0
+    ):
+        self._module.nvfp4_gemm_bias_gelu_nvfp4(
+            a, b, sfa, sfb, bias, alpha=float(alpha),
+            out_packed=out_packed, out_sfa=out_sfa,
+        )
+
+    def nvfp4_gemm_streamk_bf16(self, a, b, sfa, sfb, out, alpha=1.0):
+        self._module.nvfp4_gemm_streamk_bf16(
+            a, b, sfa, sfb, alpha=float(alpha), out=out
+        )
+
+    def nvfp4_gemm_streamk_bias_bf16(
+        self, a, b, sfa, sfb, bias, out, alpha=1.0
+    ):
+        self._module.nvfp4_gemm_streamk_bias_bf16(
+            a, b, sfa, sfb, bias, alpha=float(alpha), out=out
+        )
+
+
+def _current_arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    if major >= 12:
+        return "12.0a"
+    return f"{major}.{minor}"
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    cutlass_include = Path(os.environ.get("FLASHRT_CUTLASS_INCLUDE", str(DEFAULT_CUTLASS_INCLUDE)))
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    if not cutlass_include.is_dir():
+        raise RuntimeError(f"missing CUTLASS include path: {cutlass_include}")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list())
+    namespace = "fp4_gemm_source_test"
+    capability = torch.cuda.get_device_capability(0)
+    if capability == (11, 0):
+        gemm_sources = [
+            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_nvfp4_w4a16_gemm_sm100.cu"),
+            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_fp4_gemm_bias_bf16_sm100.cu"),
+            str(PACKAGE / "csrc" / "quantize" / "quantize_fp4_sfa_bf16.cu"),
+            str(PACKAGE / "csrc" / "gemm" / "fp4" / "sm110_dispatch.cu"),
+        ]
+        source_define = "-DFLASHRT_FP4_GEMM_SOURCE_SM110_ONLY"
+    else:
+        gemm_sources = [
+            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_nvfp4_w4a16_gemm_sm120.cu"),
+            str(PACKAGE / "csrc" / "gemm" / "fp4" / "fp4_w4a4_mma_warpsplit_sm120.cu"),
+            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_nvfp4_gemm_bias_gelu_bf16out_sm120.cu"),
+            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_nvfp4_gemm_bias_gelu_fp4out_sm120.cu"),
+            str(PACKAGE / "csrc" / "gemm" / "fp4" / "cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cu"),
+        ]
+        source_define = None
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            *gemm_sources,
+            str(PACKAGE / "csrc" / "quantize" / "quantize_fp4_sfa.cu"),
+            str(PACKAGE / "csrc" / "dequantize_fp4_sfa.cu"),
+        ],
+        extra_include_paths=[
+            str(PACKAGE / "csrc"),
+            str(cutlass_include),
+            str(REGISTRATION_INCLUDE),
+        ],
+        extra_cflags=[flag for flag in ["-O3", "-DCUDA_KERNEL", source_define] if flag],
+        extra_cuda_cflags=[
+            "-O3",
+            "--expt-relaxed-constexpr",
+            "--expt-extended-lambda",
+            "-DCUDA_KERNEL",
+            "-DCUTLASS_ARCH_MMA_SM100_SUPPORTED=1",
+            *([source_define] if source_define else []),
+        ],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return InstalledOps(importlib.import_module("fp4_gemm"))
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def make_inputs(m: int, n: int, k: int, seed: int):
+    gen = torch.Generator(device="cuda")
+    gen.manual_seed(seed)
+    a = (torch.randn((m, k), device="cuda", generator=gen) * 0.25).to(torch.float16).contiguous()
+    b = (torch.randn((n, k), device="cuda", generator=gen) * 0.25).to(torch.float16).contiguous()
+    return a, b
+
+
+def metrics(got: torch.Tensor, expected: torch.Tensor) -> tuple[float, float, float, float]:
+    diff = (got.float() - expected.float()).abs().flatten()
+    return (
+        float(diff.max().item()),
+        float(diff.mean().item()),
+        float(torch.quantile(diff, 0.99).item()),
+        float(torch.nn.functional.cosine_similarity(got.float().flatten(), expected.float().flatten(), dim=0).item()),
+    )
+
+
+def check_bf16_threshold(max_abs: float, mean_abs: float, p99_abs: float, cosine: float) -> bool:
+    return max_abs <= 0.125 and mean_abs <= 0.005 and p99_abs <= 0.03125 and cosine >= 0.999
+
+
+def select_sm110_variant(shape: tuple[int, int, int]) -> int:
+    _m, n, k = shape
+    if n >= 4 * k:
+        return 1
+    if n == 3 * k:
+        return 2
+    return 0
+
+
+def prepare_quantized(ops: SourceOps, m: int, n: int, k: int):
+    a_fp16, b_fp16 = make_inputs(m, n, k, seed=7000 + m + n + k)
+    a_packed, sfa = ops.alloc_fp4(m, k)
+    b_packed, sfb = ops.alloc_fp4(n, k)
+    ops.quantize_fp4_sfa_fp16(a_fp16, a_packed, sfa, False)
+    ops.quantize_fp4_sfa_fp16(b_fp16, b_packed, sfb, True)
+    a_deq = torch.empty_like(a_fp16)
+    b_deq = torch.empty_like(b_fp16)
+    ops.dequantize_fp4_sfa_fp16(a_packed, sfa, a_deq, False)
+    ops.dequantize_fp4_sfa_fp16(b_packed, sfb, b_deq, True)
+    torch.cuda.synchronize()
+    expected = (a_deq.float() @ b_deq.float().T).to(torch.bfloat16)
+    return a_packed, b_packed, sfa, sfb, expected
+
+
+def prepare_quantized_full(ops: SourceOps, m: int, n: int, k: int):
+    a_fp16, b_fp16 = make_inputs(m, n, k, seed=9000 + m + n + k)
+    a_packed, sfa = ops.alloc_fp4(m, k)
+    b_packed, sfb = ops.alloc_fp4(n, k)
+    ops.quantize_fp4_sfa_fp16(a_fp16, a_packed, sfa, False)
+    ops.quantize_fp4_sfa_fp16(b_fp16, b_packed, sfb, True)
+    a_deq = torch.empty_like(a_fp16)
+    b_deq = torch.empty_like(b_fp16)
+    ops.dequantize_fp4_sfa_fp16(a_packed, sfa, a_deq, False)
+    ops.dequantize_fp4_sfa_fp16(b_packed, sfb, b_deq, True)
+    return a_packed, b_packed, sfa, sfb, a_deq, b_deq
+
+
+def run_case(ops: SourceOps, name: str, shape: tuple[int, int, int]) -> list[Metrics]:
+    m, n, k = shape
+    a_packed, b_packed, sfa, sfb, expected = prepare_quantized(ops, m, n, k)
+    results: list[Metrics] = []
+    variants = (-1, 0, 1, 2) if torch.cuda.get_device_capability(0) == (11, 0) else (0, 1, 2)
+    for variant in variants:
+        out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+        ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant)
+        torch.cuda.synchronize()
+        max_abs, mean_abs, p99_abs, cosine = metrics(out, expected)
+        results.append(
+            Metrics(
+                shape=name,
+                M=m,
+                N=n,
+                K=k,
+                workload="nvfp4_gemm_bf16",
+                variant=variant,
+                max_abs=max_abs,
+                mean_abs=mean_abs,
+                p99_abs=p99_abs,
+                cosine=cosine,
+                passed=check_bf16_threshold(max_abs, mean_abs, p99_abs, cosine),
+            )
+        )
+
+    return results
+
+
+def result_row(name, shape, workload, got, expected, *, fp4_output=False):
+    max_abs, mean_abs, p99_abs, cosine = metrics(got, expected)
+    if fp4_output:
+        mean_magnitude = float(expected.float().abs().mean().item())
+        rms = float(expected.float().square().mean().sqrt().item())
+        passed = (
+            cosine >= 0.9993
+            and mean_abs / max(mean_magnitude, 1e-12) <= 0.01
+            and p99_abs / max(rms, 1e-12) <= 0.15
+        )
+    else:
+        passed = (
+            cosine >= 0.999
+            and mean_abs <= 0.008
+            and p99_abs <= 0.0625
+        )
+    return Metrics(
+        shape=name,
+        M=shape[0],
+        N=shape[1],
+        K=shape[2],
+        workload=workload,
+        variant=None,
+        max_abs=max_abs,
+        mean_abs=mean_abs,
+        p99_abs=p99_abs,
+        cosine=cosine,
+        passed=passed,
+    )
+
+
+def run_epilogue_case(ops, name: str, shape: tuple[int, int, int]):
+    m, n, k = shape
+    a, b, sfa, sfb, a_deq, b_deq = prepare_quantized_full(ops, m, n, k)
+    matmul = a_deq.float() @ b_deq.float().T
+    bias = (torch.randn(n, device="cuda") * 0.02).to(torch.bfloat16)
+    residual = torch.randn((m, n), device="cuda", dtype=torch.bfloat16)
+    rows = []
+
+    out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+    ops.nvfp4_gemm_residual_bf16(a, b, sfa, sfb, residual, out)
+    expected = (matmul + residual.float()).to(torch.bfloat16)
+    rows.append(result_row(name, shape, "nvfp4_gemm_residual_bf16", out, expected))
+
+    ops.nvfp4_gemm_bias_gelu_bf16(a, b, sfa, sfb, bias, out)
+    expected_gelu = torch.nn.functional.gelu(
+        matmul + bias.float().view(1, -1), approximate="tanh"
+    ).to(torch.bfloat16)
+    rows.append(result_row(name, shape, "nvfp4_gemm_bias_gelu_bf16", out, expected_gelu))
+
+    out_packed, out_sfa = ops.alloc_fp4(m, n)
+    ops.nvfp4_gemm_bias_gelu_nvfp4(
+        a, b, sfa, sfb, bias, out_packed, out_sfa
+    )
+    out_deq = torch.empty((m, n), device="cuda", dtype=torch.float16)
+    ops.dequantize_fp4_sfa_fp16(out_packed, out_sfa, out_deq, False)
+    staged_packed, staged_sfa = ops.alloc_fp4(m, n)
+    ops.quantize_fp4_sfa_fp16(
+        expected_gelu.to(torch.float16), staged_packed, staged_sfa, False
+    )
+    staged_deq = torch.empty_like(out_deq)
+    ops.dequantize_fp4_sfa_fp16(
+        staged_packed, staged_sfa, staged_deq, False
+    )
+    rows.append(
+        result_row(
+            name, shape, "nvfp4_gemm_bias_gelu_nvfp4",
+            out_deq, staged_deq, fp4_output=True,
+        )
+    )
+
+    ops.nvfp4_gemm_streamk_bf16(a, b, sfa, sfb, out)
+    expected_linear = matmul.to(torch.bfloat16)
+    rows.append(result_row(name, shape, "nvfp4_gemm_streamk_bf16", out, expected_linear))
+
+    ops.nvfp4_gemm_streamk_bias_bf16(a, b, sfa, sfb, bias, out)
+    expected_bias = (matmul + bias.float().view(1, -1)).to(torch.bfloat16)
+    rows.append(result_row(name, shape, "nvfp4_gemm_streamk_bias_bf16", out, expected_bias))
+    return rows
+
+
+def run_sm110_epilogue_case(ops, name: str, shape: tuple[int, int, int]):
+    m, n, k = shape
+    a, b, sfa, sfb, a_deq, b_deq = prepare_quantized_full(ops, m, n, k)
+    matmul = a_deq.float() @ b_deq.float().T
+    bias = (torch.randn(n, device="cuda") * 0.02).to(torch.bfloat16)
+    residual = torch.randn((m, n), device="cuda", dtype=torch.bfloat16)
+    rows = []
+
+    out = torch.empty_like(residual)
+    ops.nvfp4_gemm_bias_bf16(a, b, sfa, sfb, bias, out)
+    expected_bias = (matmul + bias.float().view(1, -1)).to(torch.bfloat16)
+    rows.append(result_row(
+        name, shape, "nvfp4_gemm_bias_bf16", out, expected_bias
+    ))
+
+    before = residual.clone()
+    ops.nvfp4_gemm_bias_residual_bf16(
+        a, b, sfa, sfb, bias, residual, residual
+    )
+    expected_residual = (
+        matmul + bias.float().view(1, -1) + before.float()
+    ).to(torch.bfloat16)
+    rows.append(result_row(
+        name, shape, "nvfp4_gemm_bias_residual_bf16",
+        residual, expected_residual,
+    ))
+
+    out_packed, out_sfa = ops.alloc_fp4(m, n)
+    ops.nvfp4_gemm_bias_gelu_nvfp4(
+        a, b, sfa, sfb, bias, out_packed, out_sfa
+    )
+    out_deq = torch.empty((m, n), device="cuda", dtype=torch.float16)
+    ops.dequantize_fp4_sfa_fp16(out_packed, out_sfa, out_deq, False)
+    expected_gelu = torch.nn.functional.gelu(
+        matmul + bias.float().view(1, -1), approximate="tanh"
+    ).to(torch.bfloat16)
+    staged_packed, staged_sfa = ops.alloc_fp4(m, n)
+    ops.quantize_fp4_sfa_bf16(
+        expected_gelu, staged_packed, staged_sfa, False
+    )
+    staged_deq = torch.empty_like(out_deq)
+    ops.dequantize_fp4_sfa_fp16(
+        staged_packed, staged_sfa, staged_deq, False
+    )
+    rows.append(result_row(
+        name, shape, "nvfp4_gemm_bias_gelu_nvfp4",
+        out_deq, staged_deq, fp4_output=True,
+    ))
+    return rows
+
+
+def check_installed_compile(ops: InstalledOps) -> dict[str, object]:
+    a_packed, b_packed, sfa, sfb, _ = prepare_quantized(ops, 128, 128, 128)
+
+    def call(a, b, scale_a, scale_b):
+        return ops._module.nvfp4_gemm_bf16(a, b, scale_a, scale_b)
+
+    eager = call(a_packed, b_packed, sfa, sfb)
+    compiled = torch.compile(call, fullgraph=True)
+    got = compiled(a_packed, b_packed, sfa, sfb)
+    torch.cuda.synchronize()
+    max_abs = float((got.float() - eager.float()).abs().max().item())
+    passed = bool(
+        got.dtype == torch.bfloat16
+        and got.shape == eager.shape
+        and torch.equal(got, eager)
+    )
+    return {
+        "fullgraph": True,
+        "dtype": str(got.dtype),
+        "shape": list(got.shape),
+        "max_abs": max_abs,
+        "exact": bool(torch.equal(got, eager)),
+        "passed": passed,
+    }
+
+
+def check_bf16_quantizer(ops) -> dict[str, object]:
+    """Require the direct BF16 producer to preserve the established layout."""
+    cases = [
+        (1, 5120, False),
+        (1, 6144, False),
+        (1, 17408, False),
+        (16, 2048, False),
+        (128, 512, False),
+        (64, 1024, True),
+    ]
+    rows = []
+    for case_index, (m, k, is_sfb) in enumerate(cases):
+        torch.manual_seed(8100 + case_index)
+        x = (torch.randn((m, k), device="cuda") * 1.5).to(torch.bfloat16)
+        direct_packed, direct_sfa = ops.alloc_fp4(m, k)
+        compat_packed, compat_sfa = ops.alloc_fp4(m, k)
+        # CUTLASS SFA/SFB buffers contain alignment padding that producers do
+        # not write or consume. Zero it so a full-buffer equality check still
+        # proves every mapped scale byte lands at the same address.
+        direct_sfa.zero_()
+        compat_sfa.zero_()
+        ops.quantize_fp4_sfa_bf16(x, direct_packed, direct_sfa, is_sfb)
+        ops.quantize_fp4_sfa_fp16(
+            x.to(torch.float16), compat_packed, compat_sfa, is_sfb
+        )
+        torch.cuda.synchronize()
+        packed_exact = bool(torch.equal(direct_packed, compat_packed))
+        sfa_exact = bool(torch.equal(direct_sfa, compat_sfa))
+        direct_deq = torch.empty((m, k), device="cuda", dtype=torch.float16)
+        compat_deq = torch.empty_like(direct_deq)
+        ops.dequantize_fp4_sfa_fp16(
+            direct_packed, direct_sfa, direct_deq, is_sfb
+        )
+        ops.dequantize_fp4_sfa_fp16(
+            compat_packed, compat_sfa, compat_deq, is_sfb
+        )
+        torch.cuda.synchronize()
+        dequant_exact = bool(torch.equal(direct_deq, compat_deq))
+        rows.append(
+            {
+                "shape": [m, k],
+                "is_sfb": is_sfb,
+                "packed_exact": packed_exact,
+                "sfa_exact": sfa_exact,
+                "dequant_exact": dequant_exact,
+                "passed": packed_exact and sfa_exact and dequant_exact,
+            }
+        )
+    return {"rows": rows, "passed": all(row["passed"] for row in rows)}
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=sorted(MODES), default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+
+    results: list[Metrics] = []
+    selected_shapes = SM110_SHAPES if args.mode == "thor-models" else SHAPES
+    for name in MODES[args.mode]:
+        results.extend(run_case(ops, name, selected_shapes[name]))
+    capability = torch.cuda.get_device_capability(0)
+    if args.mode == "full" and capability != (11, 0):
+        for name, shape in EPILOGUE_SHAPES.items():
+            results.extend(run_epilogue_case(ops, name, shape))
+    if capability == (11, 0) and args.mode in {"full", "thor-models"}:
+        for name in (
+            "groot_n17_dit_qkv",
+            "groot_n17_dit_ffn_up",
+            "groot_n17_dit_ffn_down",
+        ):
+            results.extend(run_sm110_epilogue_case(
+                ops, name, SM110_SHAPES[name]
+            ))
+    compile_check = None
+    bf16_quantizer_check = check_bf16_quantizer(ops)
+    if args.backend == "installed" and args.mode == "full":
+        compile_check = check_installed_compile(ops)
+    passed = sum(1 for item in results if item.passed)
+    total = len(results)
+    if compile_check is not None:
+        total += 1
+        passed += int(bool(compile_check["passed"]))
+    total += 1
+    passed += int(bool(bf16_quantizer_check["passed"]))
+    payload = {
+        "backend": args.backend,
+        "mode": args.mode,
+        "device": torch.cuda.get_device_name(),
+        "torch": torch.__version__,
+        "passed": passed,
+        "total": total,
+        "results": [asdict(item) for item in results],
+        "compile_check": compile_check,
+        "bf16_quantizer_check": bf16_quantizer_check,
+    }
+    print(json.dumps(payload, indent=2))
+    if args.json_out:
+        out = Path(args.json_out)
+        out.parent.mkdir(parents=True, exist_ok=True)
+        out.write_text(json.dumps(payload, indent=2) + "\n")
+    return 0 if passed == total else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

torch-ext/fp4_gemm/__init__.py

@@ -0,0 +1,358 @@
+"""FlashRT FP4 GEMM kernels."""
+
+from __future__ import annotations
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def sfa_size_bytes(rows: int, dim: int) -> int:
+    if rows <= 0 or dim <= 0 or dim % 16 != 0:
+        raise ValueError("rows must be positive and dim must be positive/divisible by 16")
+    n_blocks = dim // 16
+    n_row_super = (rows + 127) // 128
+    n_col_super = (n_blocks + 3) // 4
+    return n_row_super * n_col_super * 512
+
+
+def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
+    return (
+        torch.empty((rows, dim // 2), device=device, dtype=torch.uint8),
+        torch.empty((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
+    )
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
+def _linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a4_gemv_warpsplit_bf16"))
+def _gemv_warpsplit_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, warps: int = 4, stages: int = 4) -> None:
+    if a_packed.shape[0] != 1:
+        raise RuntimeError("warp-split GEMV serves M=1 only")
+    if out.shape != (1, b_packed.shape[0]):
+        raise RuntimeError("out must have shape (1, N)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
+def _legacy_linear_fake(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    out: torch.Tensor,
+    alpha: float = 1.0,
+    variant: int = -1,
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_bf16"))
+def _bias_fake(a, b, sfa, sfb, bias, out) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_residual_bf16"))
+def _bias_residual_fake(a, b, sfa, sfb, bias, residual, out) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
+def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
+def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("dequantize_fp4_sfa_fp16"))
+def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor, is_sfb: bool = False) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_residual_bf16"))
+def _residual_fake(a, b, sfa, sfb, residual, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_bf16"))
+def _bias_gelu_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bias_gelu_nvfp4"))
+def _bias_gelu_nvfp4_fake(
+    a, b, sfa, sfb, bias, out_packed, out_sfa, alpha: float = 1.0
+) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bf16"))
+def _streamk_fake(a, b, sfa, sfb, out, alpha: float = 1.0) -> None:
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_streamk_bias_bf16"))
+def _streamk_bias_fake(a, b, sfa, sfb, bias, out, alpha: float = 1.0) -> None:
+    return None
+
+
+def quantize_fp4_sfa_fp16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_fp16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def quantize_fp4_sfa_bf16(
+    x: torch.Tensor,
+    packed: torch.Tensor | None = None,
+    sfa: torch.Tensor | None = None,
+    is_sfb: bool = False,
+):
+    """Quantize BF16 directly to packed E2M1 and CUTLASS SFA/SFB."""
+    if packed is None or sfa is None:
+        packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
+    ops.quantize_fp4_sfa_bf16(x, packed, sfa, bool(is_sfb))
+    return packed, sfa
+
+
+def dequantize_fp4_sfa_fp16(
+    packed: torch.Tensor,
+    sfa: torch.Tensor,
+    out: torch.Tensor | None = None,
+    is_sfb: bool = False,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((packed.shape[0], packed.shape[1] * 2), device=packed.device, dtype=torch.float16)
+    ops.dequantize_fp4_sfa_fp16(packed, sfa, out, bool(is_sfb))
+    return out
+
+
+def nvfp4_gemm_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
+    ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
+    return out
+
+
+def fp4_w4a4_gemv_warpsplit_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    *,
+    alpha: float = 1.0,
+    warps: int = 4,
+    stages: int = 4,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Warp-split-K NVFP4 W4A4 GEMV for the M=1 decode row (SM120).
+
+    Splits K across warps inside one block with a shared-memory reduce -
+    no cross-block intermediate, so it stays safe under CUDA-graph
+    replay - and fills the SMs the tiled GEMM underfills at long-K
+    small-M decode shapes. Same packed/scale layouts as the linear
+    entry points."""
+    if out is None:
+        out = torch.empty((1, b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
+    ops.fp4_w4a4_gemv_warpsplit_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(warps), int(stages))
+    return out
+
+
+def fp4_w4a16_linear_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+    variant: int = -1,
+) -> torch.Tensor:
+    """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
+    return nvfp4_gemm_bf16(
+        a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
+    )
+
+
+def nvfp4_gemm_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 NVFP4 GEMM with a fused per-column BF16 bias."""
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_bf16(a_packed, b_packed, sfa, sfb, bias, out)
+    return out
+
+
+def nvfp4_gemm_bias_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    *,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """SM110 NVFP4 GEMM with fused BF16 bias and residual add."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_bias_residual_bf16(
+        a_packed, b_packed, sfa, sfb, bias, residual, out
+    )
+    return out
+
+
+def nvfp4_gemm_residual_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    residual: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_gemm_residual_bf16(
+        a_packed, b_packed, sfa, sfb, residual, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_bias_gelu_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_bias_gelu_nvfp4(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out_packed: torch.Tensor | None = None,
+    out_sfa: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    m, n = a_packed.shape[0], b_packed.shape[0]
+    if out_packed is None:
+        out_packed = torch.empty((m, n // 2), device=a_packed.device, dtype=torch.uint8)
+    if out_sfa is None:
+        out_sfa = torch.empty((sfa_size_bytes(m, n),), device=a_packed.device, dtype=torch.uint8)
+    ops.nvfp4_gemm_bias_gelu_nvfp4(
+        a_packed, b_packed, sfa, sfb, bias, out_packed, out_sfa, float(alpha)
+    )
+    return out_packed, out_sfa
+
+
+def nvfp4_gemm_streamk_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bf16(
+        a_packed, b_packed, sfa, sfb, out, float(alpha)
+    )
+    return out
+
+
+def nvfp4_gemm_streamk_bias_bf16(
+    a_packed: torch.Tensor,
+    b_packed: torch.Tensor,
+    sfa: torch.Tensor,
+    sfb: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (a_packed.shape[0], b_packed.shape[0]),
+            device=a_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_gemm_streamk_bias_bf16(
+        a_packed, b_packed, sfa, sfb, bias, out, float(alpha)
+    )
+    return out
+
+
+__all__ = [
+    "dequantize_fp4_sfa_fp16",
+    "fp4_w4a16_linear_bf16",
+    "fp4_w4a4_gemv_warpsplit_bf16",
+    "nvfp4_gemm_bf16",
+    "nvfp4_gemm_bias_bf16",
+    "nvfp4_gemm_bias_gelu_bf16",
+    "nvfp4_gemm_bias_gelu_nvfp4",
+    "nvfp4_gemm_bias_residual_bf16",
+    "nvfp4_gemm_residual_bf16",
+    "nvfp4_gemm_streamk_bf16",
+    "nvfp4_gemm_streamk_bias_bf16",
+    "quantize_fp4_sfa_fp16",
+    "quantize_fp4_sfa_bf16",
+    "sfa_size_bytes",
+]

torch-ext/torch_binding.cpp

@@ -0,0 +1,606 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include <torch/all.h>
+#include <torch/library.h>
+
+#include <limits>
+
+#if defined(CUDA_KERNEL)
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#endif
+
+#include "dequantize_fp4_sfa.cuh"
+#if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY)
+#include "gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_bf16out_sm120.cuh"
+#include "gemm/fp4/cutlass_nvfp4_gemm_bias_gelu_fp4out_sm120.cuh"
+#include "gemm/fp4/cutlass_nvfp4_gemm_dn_streamk_bias_sm120.cuh"
+#include "gemm/fp4/cutlass_nvfp4_w4a16_gemm_sm120.cuh"
+#include "gemm/fp4/fp4_w4a4_mma_warpsplit_sm120.cuh"
+#endif
+#include "gemm/fp4/sm110_dispatch.cuh"
+#include "quantize/quantize_fp4_sfa.cuh"
+#include "registration.h"
+#include "torch_binding.h"
+
+flash_rt::hub::Sm110GemmDispatch flash_rt::hub::sm110_gemm_dispatch = nullptr;
+flash_rt::hub::Sm110GemmBiasDispatch
+    flash_rt::hub::sm110_gemm_bias_dispatch = nullptr;
+flash_rt::hub::Sm110GemmBiasResidualDispatch
+    flash_rt::hub::sm110_gemm_bias_residual_dispatch = nullptr;
+flash_rt::hub::Sm110GemmBiasGeluFp4Dispatch
+    flash_rt::hub::sm110_gemm_bias_gelu_fp4_dispatch = nullptr;
+flash_rt::hub::Sm110QuantizeBf16Dispatch
+    flash_rt::hub::sm110_quantize_bf16_dispatch = nullptr;
+
+namespace {
+
+void check_cuda_contiguous(torch::Tensor const& tensor, const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+}
+
+void check_uint8_cuda(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == torch::kUInt8,
+              name, " must have dtype torch.uint8");
+}
+
+void check_fp16_cuda(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype torch.float16");
+}
+
+void check_bf16_cuda(torch::Tensor const& tensor, const char* name) {
+  check_cuda_contiguous(tensor, name);
+  TORCH_CHECK(tensor.scalar_type() == torch::kBFloat16,
+              name, " must have dtype torch.bfloat16");
+}
+
+int checked_int(int64_t value, const char* name) {
+  TORCH_CHECK(value > 0 && value <= std::numeric_limits<int>::max(),
+              name, " must fit in positive int");
+  return static_cast<int>(value);
+}
+
+int64_t swizzled_bytes(int64_t rows, int64_t dim) {
+  TORCH_CHECK(rows > 0 && dim > 0 && dim % 16 == 0,
+              "rows must be positive and dim must be positive/divisible by 16");
+  const int64_t n_blocks = dim / 16;
+  const int64_t n_row_super = (rows + 127) / 128;
+  const int64_t n_col_super = (n_blocks + 3) / 4;
+  return n_row_super * n_col_super * 512;
+}
+
+void check_same_device(torch::Tensor const& a, torch::Tensor const& b,
+                       const char* a_name, const char* b_name) {
+  TORCH_CHECK(a.get_device() == b.get_device(),
+              a_name, " and ", b_name, " must be on the same CUDA device");
+}
+
+struct GemmShape {
+  int64_t m;
+  int64_t n;
+  int64_t k;
+};
+
+#if defined(CUDA_KERNEL)
+cudaDeviceProp const* current_device_properties(torch::Tensor const& anchor) {
+  return at::cuda::getDeviceProperties(anchor.get_device());
+}
+
+void require_sm120(torch::Tensor const& anchor, const char* operation) {
+  auto const* props = current_device_properties(anchor);
+  TORCH_CHECK(props->major == 12 && props->minor == 0,
+              operation, " is an SM120 fused epilogue; got SM",
+              props->major, props->minor,
+              ". On SM110 use nvfp4_gemm_bf16 with fp4-fused-ops producers.");
+}
+#endif
+
+GemmShape check_fp4_gemm_inputs(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb) {
+  check_uint8_cuda(a_packed, "a_packed");
+  check_uint8_cuda(b_packed, "b_packed");
+  check_uint8_cuda(sfa, "sfa");
+  check_uint8_cuda(sfb, "sfb");
+  TORCH_CHECK(a_packed.dim() == 2, "a_packed must have shape (M, K / 2)");
+  TORCH_CHECK(b_packed.dim() == 2, "b_packed must have shape (N, K / 2)");
+  const int64_t m = a_packed.size(0);
+  const int64_t n = b_packed.size(0);
+  const int64_t k_half = a_packed.size(1);
+  TORCH_CHECK(m > 0 && n > 0 && k_half > 0, "M, N, and K must be positive");
+  TORCH_CHECK(b_packed.size(1) == k_half,
+              "a_packed and b_packed must have the same K / 2 dimension");
+  const int64_t k = k_half * 2;
+  TORCH_CHECK(k % 16 == 0, "K must be divisible by 16");
+  TORCH_CHECK(sfa.numel() >= swizzled_bytes(m, k),
+              "sfa is too small for CUTLASS SFA layout");
+  TORCH_CHECK(sfb.numel() >= swizzled_bytes(n, k),
+              "sfb is too small for CUTLASS SFB layout");
+  check_same_device(a_packed, b_packed, "a_packed", "b_packed");
+  check_same_device(a_packed, sfa, "a_packed", "sfa");
+  check_same_device(a_packed, sfb, "a_packed", "sfb");
+  return {m, n, k};
+}
+
+int select_sm110_variant(GemmShape const& shape, int64_t requested) {
+  if (requested >= 0) return static_cast<int>(requested);
+  if (shape.n >= 4 * shape.k) return 1;
+  if (shape.n == 3 * shape.k) return 2;
+  return 0;
+}
+
+}  // namespace
+
+void fp4_w4a4_gemv_warpsplit_bf16(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb,
+    torch::Tensor& out,
+    double alpha,
+    int64_t warps,
+    int64_t stages) {
+  auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb);
+  check_bf16_cuda(out, "out");
+  TORCH_CHECK(shape.m == 1,
+              "warp-split GEMV serves the M=1 decode row only");
+  TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}),
+              "out must have shape (1, N)");
+  TORCH_CHECK(warps == 2 || warps == 4 || warps == 8,
+              "warps must be 2, 4 or 8");
+  TORCH_CHECK(stages == 3 || stages == 4 || stages == 6,
+              "stages must be 3, 4 or 6");
+  TORCH_CHECK(shape.n % 8 == 0, "N must be a multiple of 8");
+  TORCH_CHECK(shape.k % 64 == 0 && (shape.k / 64) % warps == 0,
+              "K must be a multiple of 64*warps");
+  check_same_device(a_packed, out, "a_packed", "out");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(a_packed.device());
+  auto const* props = current_device_properties(a_packed);
+  TORCH_CHECK(props->major == 12 && props->minor == 0,
+              "the warp-split GEMV is an SM120 kernel; got SM",
+              props->major, props->minor);
+#if defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY)
+  TORCH_CHECK(false, "SM120 FP4 GEMM source is not present in this build");
+#else
+  auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream();
+  const int rc = flash_rt::gemm::fp4_w4a4_mma_sm120_warpsplit_bf16out(
+      a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(),
+      checked_int(shape.n, "N"), checked_int(shape.k, "K"),
+      sfa.data_ptr(), sfb.data_ptr(), static_cast<float>(alpha),
+      static_cast<int>(warps), static_cast<int>(stages), stream);
+  TORCH_CHECK(rc == 0, "fp4_w4a4_gemv_warpsplit_bf16 failed with rc=", rc);
+#endif
+#else
+  TORCH_CHECK(false, "fp4-gemm was not built with CUDA support");
+#endif
+}
+
+void fp4_w4a16_linear_bf16(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb,
+    torch::Tensor& out,
+    double alpha,
+    int64_t variant) {
+  auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb);
+  check_bf16_cuda(out, "out");
+  TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}),
+              "out must have shape (M, N)");
+  TORCH_CHECK(variant >= -1 && variant <= 2,
+              "variant must be -1(auto), 0(default), 1(widen), or 2(pingpong)");
+  check_same_device(a_packed, out, "a_packed", "out");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(a_packed.device());
+  auto const* props = current_device_properties(a_packed);
+  TORCH_CHECK((props->major == 11 && props->minor == 0) ||
+                  (props->major == 12 && props->minor == 0),
+              "nvfp4_gemm_bf16 requires SM110 or SM120; got SM",
+              props->major, props->minor);
+  auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream();
+  if (props->major == 11) {
+    variant = select_sm110_variant(shape, variant);
+    TORCH_CHECK(flash_rt::hub::sm110_gemm_dispatch != nullptr,
+                "SM110 FP4 GEMM source is not present in this build");
+    flash_rt::hub::sm110_gemm_dispatch(
+        a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(),
+        checked_int(shape.m, "M"), checked_int(shape.n, "N"),
+        checked_int(shape.k, "K"), sfa.data_ptr(), sfb.data_ptr(),
+        static_cast<float>(alpha), variant, stream);
+  } else {
+#if defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY)
+    TORCH_CHECK(false, "SM120 FP4 GEMM source is not present in this build");
+#else
+    if (variant == 1) {
+      flash_rt::gemm::fp4_w4a16_gemm_sm120_bf16out_widen(
+          a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(),
+          checked_int(shape.m, "M"), checked_int(shape.n, "N"), checked_int(shape.k, "K"),
+          sfa.data_ptr(), sfb.data_ptr(), static_cast<float>(alpha), stream);
+    } else if (variant == 2) {
+      flash_rt::gemm::fp4_w4a16_gemm_sm120_bf16out_pingpong(
+          a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(),
+          checked_int(shape.m, "M"), checked_int(shape.n, "N"), checked_int(shape.k, "K"),
+          sfa.data_ptr(), sfb.data_ptr(), static_cast<float>(alpha), stream);
+    } else {
+      flash_rt::gemm::fp4_w4a16_gemm_sm120_bf16out(
+          a_packed.data_ptr(), b_packed.data_ptr(), out.data_ptr(),
+          checked_int(shape.m, "M"), checked_int(shape.n, "N"), checked_int(shape.k, "K"),
+          sfa.data_ptr(), sfb.data_ptr(), static_cast<float>(alpha), stream);
+    }
+#endif
+  }
+#endif
+}
+
+void nvfp4_gemm_bias_bf16(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb,
+    torch::Tensor const& bias,
+    torch::Tensor& out) {
+  auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb);
+  check_bf16_cuda(bias, "bias");
+  check_bf16_cuda(out, "out");
+  TORCH_CHECK(bias.dim() == 1 && bias.numel() == shape.n,
+              "bias must have shape (N,)");
+  TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}),
+              "out must have shape (M, N)");
+  check_same_device(a_packed, bias, "a_packed", "bias");
+  check_same_device(a_packed, out, "a_packed", "out");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(a_packed.device());
+  auto const* props = current_device_properties(a_packed);
+  TORCH_CHECK(props->major == 11 && props->minor == 0,
+              "nvfp4_gemm_bias_bf16 currently requires SM110; got SM",
+              props->major, props->minor);
+  TORCH_CHECK(flash_rt::hub::sm110_gemm_bias_dispatch != nullptr,
+              "SM110 fused-bias FP4 GEMM source is not present in this build");
+  auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream();
+  const int rc = flash_rt::hub::sm110_gemm_bias_dispatch(
+      a_packed.data_ptr(), sfa.data_ptr(), b_packed.data_ptr(), sfb.data_ptr(),
+      bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"),
+      checked_int(shape.n, "N"), checked_int(shape.k, "K"), stream);
+  TORCH_CHECK(rc == 0, "nvfp4_gemm_bias_bf16 failed with rc=", rc);
+#endif
+}
+
+void nvfp4_gemm_bias_residual_bf16(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb,
+    torch::Tensor const& bias,
+    torch::Tensor const& residual,
+    torch::Tensor& out) {
+  auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb);
+  check_bf16_cuda(bias, "bias");
+  check_bf16_cuda(residual, "residual");
+  check_bf16_cuda(out, "out");
+  TORCH_CHECK(bias.dim() == 1 && bias.numel() == shape.n,
+              "bias must have shape (N,)");
+  TORCH_CHECK(residual.sizes() == torch::IntArrayRef({shape.m, shape.n}),
+              "residual must have shape (M, N)");
+  TORCH_CHECK(out.sizes() == residual.sizes(), "out must match residual");
+  check_same_device(a_packed, bias, "a_packed", "bias");
+  check_same_device(a_packed, residual, "a_packed", "residual");
+  check_same_device(a_packed, out, "a_packed", "out");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(a_packed.device());
+  auto const* props = current_device_properties(a_packed);
+  TORCH_CHECK(props->major == 11 && props->minor == 0,
+              "nvfp4_gemm_bias_residual_bf16 currently requires SM110; got SM",
+              props->major, props->minor);
+  TORCH_CHECK(flash_rt::hub::sm110_gemm_bias_residual_dispatch != nullptr,
+              "SM110 bias-residual FP4 GEMM source is not present in this build");
+  auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream();
+  const int rc = flash_rt::hub::sm110_gemm_bias_residual_dispatch(
+      a_packed.data_ptr(), sfa.data_ptr(), b_packed.data_ptr(), sfb.data_ptr(),
+      bias.data_ptr(), residual.data_ptr(), out.data_ptr(),
+      checked_int(shape.m, "M"), checked_int(shape.n, "N"),
+      checked_int(shape.k, "K"), stream);
+  TORCH_CHECK(rc == 0, "nvfp4_gemm_bias_residual_bf16 failed with rc=", rc);
+#endif
+}
+
+void nvfp4_gemm_residual_bf16(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb,
+    torch::Tensor const& residual,
+    torch::Tensor& out,
+    double alpha) {
+  auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb);
+  check_bf16_cuda(residual, "residual");
+  check_bf16_cuda(out, "out");
+  TORCH_CHECK(residual.sizes() == torch::IntArrayRef({shape.m, shape.n}),
+              "residual must have shape (M, N)");
+  TORCH_CHECK(out.sizes() == residual.sizes(), "out must match residual");
+  check_same_device(a_packed, residual, "a_packed", "residual");
+  check_same_device(a_packed, out, "a_packed", "out");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(a_packed.device());
+  require_sm120(a_packed, "nvfp4_gemm_residual_bf16");
+  auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream();
+#if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY)
+  flash_rt::gemm::fp4_w4a16_gemm_residual_sm120_bf16out(
+      a_packed.data_ptr(), b_packed.data_ptr(), residual.data_ptr(),
+      out.data_ptr(), checked_int(shape.m, "M"), checked_int(shape.n, "N"),
+      checked_int(shape.k, "K"), sfa.data_ptr(), sfb.data_ptr(),
+      static_cast<float>(alpha), stream);
+#endif
+#endif
+}
+
+void nvfp4_gemm_bias_gelu_bf16(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb,
+    torch::Tensor const& bias,
+    torch::Tensor& out,
+    double alpha) {
+  auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb);
+  check_bf16_cuda(bias, "bias");
+  check_bf16_cuda(out, "out");
+  TORCH_CHECK(bias.dim() == 1 && bias.numel() == shape.n,
+              "bias must have shape (N,)");
+  TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}),
+              "out must have shape (M, N)");
+  check_same_device(a_packed, bias, "a_packed", "bias");
+  check_same_device(a_packed, out, "a_packed", "out");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(a_packed.device());
+  require_sm120(a_packed, "nvfp4_gemm_bias_gelu_bf16");
+  auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream();
+#if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY)
+  flash_rt::gemm::fp4_w4a16_gemm_bias_gelu_bf16out_sm120(
+      a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(),
+      bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"),
+      checked_int(shape.n, "N"), checked_int(shape.k, "K"),
+      static_cast<float>(alpha), stream);
+#endif
+#endif
+}
+
+void nvfp4_gemm_bias_gelu_nvfp4(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb,
+    torch::Tensor const& bias,
+    torch::Tensor& out_packed,
+    torch::Tensor& out_sfa,
+    double alpha) {
+  auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb);
+  check_bf16_cuda(bias, "bias");
+  check_uint8_cuda(out_packed, "out_packed");
+  check_uint8_cuda(out_sfa, "out_sfa");
+  TORCH_CHECK(bias.dim() == 1 && bias.numel() == shape.n,
+              "bias must have shape (N,)");
+  TORCH_CHECK(shape.n % 2 == 0 &&
+                  out_packed.sizes() ==
+                      torch::IntArrayRef({shape.m, shape.n / 2}),
+              "out_packed must have shape (M, N / 2)");
+  TORCH_CHECK(out_sfa.numel() >= swizzled_bytes(shape.m, shape.n),
+              "out_sfa is too small for output scale layout");
+  check_same_device(a_packed, bias, "a_packed", "bias");
+  check_same_device(a_packed, out_packed, "a_packed", "out_packed");
+  check_same_device(a_packed, out_sfa, "a_packed", "out_sfa");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(a_packed.device());
+  auto const* props = current_device_properties(a_packed);
+  auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream();
+  if (props->major == 11 && props->minor == 0) {
+    TORCH_CHECK(flash_rt::hub::sm110_gemm_bias_gelu_fp4_dispatch != nullptr,
+                "SM110 bias-GELU-FP4 GEMM source is not present in this build");
+    const int rc = flash_rt::hub::sm110_gemm_bias_gelu_fp4_dispatch(
+        a_packed.data_ptr(), sfa.data_ptr(), b_packed.data_ptr(), sfb.data_ptr(),
+        bias.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(),
+        checked_int(shape.m, "M"), checked_int(shape.n, "N"),
+        checked_int(shape.k, "K"), stream);
+    TORCH_CHECK(rc == 0, "nvfp4_gemm_bias_gelu_nvfp4 failed with rc=", rc);
+    return;
+  }
+  require_sm120(a_packed, "nvfp4_gemm_bias_gelu_nvfp4");
+#if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY)
+  flash_rt::gemm::fp4_w4a16_gemm_bias_gelu_fp4out_sm120(
+      a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(),
+      bias.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(),
+      checked_int(shape.m, "M"), checked_int(shape.n, "N"),
+      checked_int(shape.k, "K"), static_cast<float>(alpha), stream);
+#endif
+#endif
+}
+
+void nvfp4_gemm_streamk_bf16(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb,
+    torch::Tensor& out,
+    double alpha) {
+  auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb);
+  check_bf16_cuda(out, "out");
+  TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}),
+              "out must have shape (M, N)");
+  check_same_device(a_packed, out, "a_packed", "out");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(a_packed.device());
+  require_sm120(a_packed, "nvfp4_gemm_streamk_bf16");
+  auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream();
+#if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY)
+  flash_rt::gemm::fp4_w4a16_gemm_dn_streamk_bf16out_sm120(
+      a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(),
+      out.data_ptr(), checked_int(shape.m, "M"), checked_int(shape.n, "N"),
+      checked_int(shape.k, "K"), static_cast<float>(alpha), stream);
+#endif
+#endif
+}
+
+void nvfp4_gemm_streamk_bias_bf16(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb,
+    torch::Tensor const& bias,
+    torch::Tensor& out,
+    double alpha) {
+  auto shape = check_fp4_gemm_inputs(a_packed, b_packed, sfa, sfb);
+  check_bf16_cuda(bias, "bias");
+  check_bf16_cuda(out, "out");
+  TORCH_CHECK(bias.dim() == 1 && bias.numel() == shape.n,
+              "bias must have shape (N,)");
+  TORCH_CHECK(out.sizes() == torch::IntArrayRef({shape.m, shape.n}),
+              "out must have shape (M, N)");
+  check_same_device(a_packed, bias, "a_packed", "bias");
+  check_same_device(a_packed, out, "a_packed", "out");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(a_packed.device());
+  require_sm120(a_packed, "nvfp4_gemm_streamk_bias_bf16");
+  auto stream = at::cuda::getCurrentCUDAStream(a_packed.get_device()).stream();
+#if !defined(FLASHRT_FP4_GEMM_SOURCE_SM110_ONLY)
+  flash_rt::gemm::fp4_w4a16_gemm_dn_streamk_bias_bf16out_sm120(
+      a_packed.data_ptr(), b_packed.data_ptr(), sfa.data_ptr(), sfb.data_ptr(),
+      bias.data_ptr(), out.data_ptr(), checked_int(shape.m, "M"),
+      checked_int(shape.n, "N"), checked_int(shape.k, "K"),
+      static_cast<float>(alpha), stream);
+#endif
+#endif
+}
+
+void quantize_fp4_sfa_fp16(
+    torch::Tensor const& x,
+    torch::Tensor& packed,
+    torch::Tensor& sfa,
+    bool is_sfb) {
+  check_fp16_cuda(x, "x");
+  check_uint8_cuda(packed, "packed");
+  check_uint8_cuda(sfa, "sfa");
+  TORCH_CHECK(x.dim() == 2, "x must have shape (rows, dim)");
+  const int64_t rows = x.size(0);
+  const int64_t dim = x.size(1);
+  TORCH_CHECK(dim % 16 == 0, "x.shape[1] must be divisible by 16");
+  TORCH_CHECK(packed.sizes() == torch::IntArrayRef({rows, dim / 2}),
+              "packed must have shape (rows, dim / 2)");
+  TORCH_CHECK(sfa.numel() >= swizzled_bytes(rows, dim),
+              "sfa is too small for CUTLASS SFA/SFB layout");
+  check_same_device(x, packed, "x", "packed");
+  check_same_device(x, sfa, "x", "sfa");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  const int rc = flash_rt::fp4::quantize_fp4_dynamic_sfa_fp16(
+      x.data_ptr(), packed.data_ptr(), sfa.data_ptr(),
+      checked_int(rows, "rows"), checked_int(dim, "dim"), is_sfb, stream);
+  TORCH_CHECK(rc == 0, "quantize_fp4_dynamic_sfa_fp16 failed with rc=", rc);
+#endif
+}
+
+void quantize_fp4_sfa_bf16(
+    torch::Tensor const& x,
+    torch::Tensor& packed,
+    torch::Tensor& sfa,
+    bool is_sfb) {
+  check_bf16_cuda(x, "x");
+  check_uint8_cuda(packed, "packed");
+  check_uint8_cuda(sfa, "sfa");
+  TORCH_CHECK(x.dim() == 2, "x must have shape (rows, dim)");
+  const int64_t rows = x.size(0);
+  const int64_t dim = x.size(1);
+  TORCH_CHECK(dim % 16 == 0, "x.shape[1] must be divisible by 16");
+  TORCH_CHECK(packed.sizes() == torch::IntArrayRef({rows, dim / 2}),
+              "packed must have shape (rows, dim / 2)");
+  TORCH_CHECK(sfa.numel() >= swizzled_bytes(rows, dim),
+              "sfa is too small for CUTLASS SFA/SFB layout");
+  check_same_device(x, packed, "x", "packed");
+  check_same_device(x, sfa, "x", "sfa");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(x.device());
+  auto stream = at::cuda::getCurrentCUDAStream(x.get_device()).stream();
+  auto const* props = current_device_properties(x);
+  int rc = 0;
+  if (props->major == 11 && props->minor == 0) {
+    TORCH_CHECK(flash_rt::hub::sm110_quantize_bf16_dispatch != nullptr,
+                "SM110 vectorized BF16 FP4 quantizer is not present in this build");
+    rc = flash_rt::hub::sm110_quantize_bf16_dispatch(
+        x.data_ptr(), packed.data_ptr(), sfa.data_ptr(),
+        checked_int(rows, "rows"), checked_int(dim, "dim"), is_sfb, stream);
+  } else {
+    rc = flash_rt::fp4::quantize_fp4_dynamic_sfa_bf16(
+        x.data_ptr(), packed.data_ptr(), sfa.data_ptr(),
+        checked_int(rows, "rows"), checked_int(dim, "dim"), is_sfb, stream);
+  }
+  TORCH_CHECK(rc == 0, "quantize_fp4_dynamic_sfa_bf16 failed with rc=", rc);
+#endif
+}
+
+void dequantize_fp4_sfa_fp16(
+    torch::Tensor const& packed,
+    torch::Tensor const& sfa,
+    torch::Tensor& out,
+    bool is_sfb) {
+  check_uint8_cuda(packed, "packed");
+  check_uint8_cuda(sfa, "sfa");
+  check_fp16_cuda(out, "out");
+  TORCH_CHECK(out.dim() == 2, "out must have shape (rows, dim)");
+  const int64_t rows = out.size(0);
+  const int64_t dim = out.size(1);
+  TORCH_CHECK(dim % 16 == 0, "out.shape[1] must be divisible by 16");
+  TORCH_CHECK(packed.sizes() == torch::IntArrayRef({rows, dim / 2}),
+              "packed must have shape (rows, dim / 2)");
+  TORCH_CHECK(sfa.numel() >= swizzled_bytes(rows, dim),
+              "sfa is too small for CUTLASS SFA layout");
+  check_same_device(packed, sfa, "packed", "sfa");
+  check_same_device(packed, out, "packed", "out");
+#if defined(CUDA_KERNEL)
+  at::cuda::CUDAGuard device_guard(packed.device());
+  auto stream = at::cuda::getCurrentCUDAStream(packed.get_device()).stream();
+  flash_rt::fused_fp4::dequantize_fp4_sfa_fp16(
+      reinterpret_cast<const uint8_t*>(packed.data_ptr()),
+      reinterpret_cast<const uint8_t*>(sfa.data_ptr()),
+      reinterpret_cast<__half*>(out.data_ptr()),
+      checked_int(rows, "rows"), checked_int(dim, "dim"), is_sfb, stream);
+#endif
+}
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("nvfp4_gemm_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor! out, float alpha=1.0, int variant=-1) -> ()");
+  ops.def("fp4_w4a16_linear_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor! out, float alpha=1.0, int variant=-1) -> ()");
+  ops.def("fp4_w4a4_gemv_warpsplit_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor! out, float alpha=1.0, int warps=4, int stages=4) -> ()");
+  ops.def("nvfp4_gemm_bias_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor bias, Tensor! out) -> ()");
+  ops.def("nvfp4_gemm_bias_residual_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor bias, Tensor residual, Tensor! out) -> ()");
+  ops.def("nvfp4_gemm_residual_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor residual, Tensor! out, float alpha=1.0) -> ()");
+  ops.def("nvfp4_gemm_bias_gelu_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor bias, Tensor! out, float alpha=1.0) -> ()");
+  ops.def("nvfp4_gemm_bias_gelu_nvfp4(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor bias, Tensor! out_packed, Tensor! out_sfa, float alpha=1.0) -> ()");
+  ops.def("nvfp4_gemm_streamk_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor! out, float alpha=1.0) -> ()");
+  ops.def("nvfp4_gemm_streamk_bias_bf16(Tensor a_packed, Tensor b_packed, Tensor sfa, Tensor sfb, Tensor bias, Tensor! out, float alpha=1.0) -> ()");
+  ops.def("quantize_fp4_sfa_fp16(Tensor x, Tensor! packed, Tensor! sfa, bool is_sfb=False) -> ()");
+  ops.def("quantize_fp4_sfa_bf16(Tensor x, Tensor! packed, Tensor! sfa, bool is_sfb=False) -> ()");
+  ops.def("dequantize_fp4_sfa_fp16(Tensor packed, Tensor sfa, Tensor! out, bool is_sfb=False) -> ()");
+#if defined(CUDA_KERNEL)
+  ops.impl("nvfp4_gemm_bf16", torch::kCUDA, &fp4_w4a16_linear_bf16);
+  ops.impl("fp4_w4a16_linear_bf16", torch::kCUDA, &fp4_w4a16_linear_bf16);
+  ops.impl("fp4_w4a4_gemv_warpsplit_bf16", torch::kCUDA, &fp4_w4a4_gemv_warpsplit_bf16);
+  ops.impl("nvfp4_gemm_bias_bf16", torch::kCUDA, &nvfp4_gemm_bias_bf16);
+  ops.impl("nvfp4_gemm_bias_residual_bf16", torch::kCUDA, &nvfp4_gemm_bias_residual_bf16);
+  ops.impl("nvfp4_gemm_residual_bf16", torch::kCUDA, &nvfp4_gemm_residual_bf16);
+  ops.impl("nvfp4_gemm_bias_gelu_bf16", torch::kCUDA, &nvfp4_gemm_bias_gelu_bf16);
+  ops.impl("nvfp4_gemm_bias_gelu_nvfp4", torch::kCUDA, &nvfp4_gemm_bias_gelu_nvfp4);
+  ops.impl("nvfp4_gemm_streamk_bf16", torch::kCUDA, &nvfp4_gemm_streamk_bf16);
+  ops.impl("nvfp4_gemm_streamk_bias_bf16", torch::kCUDA, &nvfp4_gemm_streamk_bias_bf16);
+  ops.impl("quantize_fp4_sfa_fp16", torch::kCUDA, &quantize_fp4_sfa_fp16);
+  ops.impl("quantize_fp4_sfa_bf16", torch::kCUDA, &quantize_fp4_sfa_bf16);
+  ops.impl("dequantize_fp4_sfa_fp16", torch::kCUDA, &dequantize_fp4_sfa_fp16);
+#endif
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <torch/all.h>
+
+void fp4_w4a16_linear_bf16(
+    torch::Tensor const& a_packed,
+    torch::Tensor const& b_packed,
+    torch::Tensor const& sfa,
+    torch::Tensor const& sfb,
+    torch::Tensor& out,
+    double alpha,
+    int64_t variant);
+
+void quantize_fp4_sfa_fp16(
+    torch::Tensor const& x,
+    torch::Tensor& packed,
+    torch::Tensor& sfa,
+    bool is_sfb);
+
+void quantize_fp4_sfa_bf16(
+    torch::Tensor const& x,
+    torch::Tensor& packed,
+    torch::Tensor& sfa,
+    bool is_sfb);
+
+void dequantize_fp4_sfa_fp16(
+    torch::Tensor const& packed,
+    torch::Tensor const& sfa,
+    torch::Tensor& out,
+    bool is_sfb);