[CUDA] Add SM120 NVF4 block-scale MMA support#2364
Open
qqq-tao wants to merge 18 commits into
Open
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds SM120 NVF4 block-scale MMA support end-to-end: TL builtin intrinsic, PTX/mma template, CodeGen emission, 4-bit layout transforms and emitter wiring, CUTLASS CUDA reference kernel, TileLang correctness harness with packing/decoding utilities, and CUDA-gated tests. ChangesNVF4 Block-Scale MMA Support
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@maint/gemm/correctness_evaluation_nvf4_vs_cutlass.py`:
- Around line 253-256: The code sets cutlass_root using a hardcoded developer
path which is unsafe; update the logic in the block that defines repo and
cutlass_root (variables cutlass_root, repo and the use of
os.environ["CUTLASS_ROOT"]) to remove the hardcoded "/data/home/..." fallback
and instead: prefer the CUTLASS_ROOT environment variable if set, otherwise fall
back to repo / "3rdparty" / "cutlass"; ensure cutlass_root.exists() is checked
and raise a clear error or log if neither location exists so the failure is
explicit.
In `@tilelang/language/ast/ir.py`:
- Line 2140: The module's __all__ includes "ptx_mma_block_scale" but no symbol
by that name is defined or imported, which breaks wildcard imports; either
remove "ptx_mma_block_scale" from the __all__ list or add a proper
definition/import for ptx_mma_block_scale (e.g., define the function/class or
import it from its source) so the name is actually bound in this module; update
the __all__ entry near the existing list and ensure the symbol name matches
exactly the defined/imported identifier.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6a5bd78f-7a51-446a-bd8a-0078295fbd05
📒 Files selected for processing (16)
maint/gemm/correctness_evaluation_nvf4_vs_cutlass.pymaint/gemm/cutlass_nvf4_ref.cusrc/cuda/codegen/codegen_cuda.ccsrc/cuda/codegen/codegen_cuda.hsrc/op/builtin.ccsrc/op/builtin.hsrc/tl_templates/cuda/instruction/mma_block_scale.htesting/python/language/test_tilelang_language_nvf4_mma_block_scale.pytilelang/cuda/intrinsics/__init__.pytilelang/cuda/intrinsics/layout/mma_layout.pytilelang/cuda/intrinsics/macro/__init__.pytilelang/cuda/intrinsics/macro/mma_macro_generator.pytilelang/intrinsics/__init__.pytilelang/language/ast/ir.pytilelang/language/tir/ir.pytilelang/language/tir/op.py
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
tilelang/cuda/intrinsics/macro/mma_macro_generator.py (2)
1402-1407:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPreserve outer-dimension bases for plain scale buffers.
_scale_region_parts()treatsBufferdifferently from the A/B paths and drops every prefix dimension by returning([], 0, 0). The accesses at Lines 1478-1482 and Line 1510 then only index the trailing two axes, which breaks direct N-D scale buffers and sliced views. Reuse_legalize_to_buffer_region()here so scale buffers follow the same region contract as A/B.Suggested fix
`@staticmethod` def _scale_region_parts(scale_buf: Buffer | BufferRegion): - if isinstance(scale_buf, BufferRegion): - return scale_buf.buffer, [r.min for r in scale_buf.region[:-2]], scale_buf.region[-2].min, scale_buf.region[-1].min - if isinstance(scale_buf, Buffer): - return scale_buf, [], 0, 0 - raise ValueError(f"Unsupported scale buffer type: {type(scale_buf)}") + region = TensorCoreIntrinEmitter._legalize_to_buffer_region(scale_buf) + return ( + region.buffer, + [r.min for r in region.region[:-2]], + region.region[-2].min, + region.region[-1].min, + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/cuda/intrinsics/macro/mma_macro_generator.py` around lines 1402 - 1407, _scale_region_parts currently returns (buffer, [], 0, 0) for plain Buffer which drops outer-dimension bases and breaks N-D scale buffers; change it to call and reuse _legalize_to_buffer_region(scale_buf) so both Buffer and BufferRegion follow the same region contract as A/B. Update _scale_region_parts to accept either type, call _legalize_to_buffer_region when scale_buf is a Buffer to obtain the buffer and full region, and then extract the same tuples (buffer, [r.min for r in region[:-2]], region[-2].min, region[-1].min) as for BufferRegion; keep the existing ValueError for unsupported types. Ensure references to _scale_region_parts usage (the indexing code that expects preserved outer bases) continue to work unchanged.
1282-1333:⚠️ Potential issue | 🟠 Major | ⚡ Quick winValidate the fixed NVF4 contract in the constructor.
This emitter always lowers as
mxf4nvf4withk64,e2m1/e2m1, and the block-scale fragment layouts from this class, but it never rejects incompatiblea_dtype,b_dtype, oraccum_dtypeinputs. A caller can currently instantiate this public API with, for example,float16operands and still get NVF4 PTX emitted against mismatched fragment assumptions. Fail fast here instead of silently generating wrong code.Suggested guard
def __init__( self, a_dtype: str = T.float4_e2m1fn, b_dtype: str = T.float4_e2m1fn, accum_dtype: str = T.float32, @@ kind: str = "mxf4nvf4", scale_vec_size: int = 4, stype: str = "ue4m3", ): + if str(DataType(a_dtype)) != str(T.float4_e2m1fn) or str(DataType(b_dtype)) != str(T.float4_e2m1fn): + raise ValueError("SM120 block-scale MMA currently only supports float4_e2m1fn operands") + if str(DataType(accum_dtype)) != str(T.float32): + raise ValueError("SM120 block-scale MMA currently only supports float32 accumulation") self.block_scale_config = _get_block_scale_mma_config(kind, scale_vec_size, stype)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/cuda/intrinsics/macro/mma_macro_generator.py` around lines 1282 - 1333, The constructor currently forces block-scale config to NVF4 but does not validate the caller-provided dtypes, so callers can pass incompatible a_dtype/b_dtype/accum_dtype and silently generate wrong code; after calling _get_block_scale_mma_config(...) in __init__, add a guard that checks the resolved self.block_scale_config.kind (and/or its expected dtype descriptors from the config) against the incoming a_dtype, b_dtype, and accum_dtype parameters (e.g., ensure a_dtype and b_dtype match the NVF4 fragment dtypes such as T.float4_e2m1fn and accum_dtype matches the expected accumulator like T.float32 or whatever the config exposes), and raise a ValueError with a clear message if they mismatch; perform this validation before calling super().__init__ so invalid combinations fail fast.maint/gemm/correctness_evaluation_nvf4_vs_cutlass.py (1)
342-354:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winThe FP32 reference diagnostics are missing the UE4M3 block scales.
refis built from the decoded FP4 payloads only, sodiff_tl_refanddiff_cutlass_refare not checking the same computation as the block-scaled GEMMs. In the defaultNVF4_SCALE_MODE="varying"case, those numbers are inherently misleading. Either applysfa_logical/sfb_logicalper 16-K chunk when buildingref, or drop these prints until that scale-aware reference exists.Suggested minimal change
- ref = _decode_rowmajor_fp4(a, M, K) @ _decode_rowmajor_fp4(b, N, K).T - diff_tl_ref = (c_tl - ref).abs() - diff_cutlass_ref = (c_cutlass - ref).abs() print("scale_mode:", scale_mode) print("input_mode:", input_mode) print("max_abs_diff:", diff.max().item()) print("mean_abs_diff:", diff.mean().item()) print("max_abs_diff_transposed:", diff_t.max().item()) print("mean_abs_diff_transposed:", diff_t.mean().item()) - print("max_abs_diff_tilelang_ref:", diff_tl_ref.max().item()) - print("mean_abs_diff_tilelang_ref:", diff_tl_ref.mean().item()) - print("max_abs_diff_cutlass_ref:", diff_cutlass_ref.max().item()) - print("mean_abs_diff_cutlass_ref:", diff_cutlass_ref.mean().item())🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@maint/gemm/correctness_evaluation_nvf4_vs_cutlass.py` around lines 342 - 354, The FP32 reference `ref` is computed from raw decoded FP4 payloads so it doesn't include the UE4M3 block scales (`sfa_logical`/`sfb_logical`), making `diff_tl_ref` and `diff_cutlass_ref` invalid under NVF4_SCALE_MODE="varying"; fix by applying the per-block scales to the decoded tensors before matmul: after calling `_decode_rowmajor_fp4(a, M, K)` and `_decode_rowmajor_fp4(b, N, K)`, multiply each 16xK chunk of the decoded A by the corresponding entries in `sfa_logical` and each 16xK chunk of the decoded B by `sfb_logical` (or apply equivalent broadcasting per 16-K block) so `ref = scaled_decoded_a @ scaled_decoded_b.T` matches the block-scaled GEMM, otherwise remove the `diff_tl_ref`/`diff_cutlass_ref` prints until the scale-aware reference is implemented; reference symbols: `ref`, `_decode_rowmajor_fp4`, `sfa_logical`, `sfb_logical`, `diff_tl_ref`, `diff_cutlass_ref`.
♻️ Duplicate comments (1)
maint/gemm/correctness_evaluation_nvf4_vs_cutlass.py (1)
247-256:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFail fast when the CUTLASS root is missing.
This now avoids the developer-local fallback, but it still passes non-existent include roots into
load(). IfCUTLASS_ROOTis unset or wrong and3rdparty/cutlassis absent, the harness fails later with a compiler error instead of an explicit setup error.Suggested fix
repo = Path(__file__).resolve().parents[2] cutlass_root_env = os.environ.get("CUTLASS_ROOT") cutlass_root = Path(cutlass_root_env) if cutlass_root_env else repo / "3rdparty" / "cutlass" + if not cutlass_root.exists(): + raise RuntimeError( + f"CUTLASS not found at {cutlass_root}. " + "Set CUTLASS_ROOT or ensure 3rdparty/cutlass exists." + ) return load(🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@maint/gemm/correctness_evaluation_nvf4_vs_cutlass.py` around lines 247 - 256, Check that the computed cutlass_root (from cutlass_root_env or repo / "3rdparty" / "cutlass") actually exists before calling load; if it does not exist, raise a clear RuntimeError instructing the developer to set CUTLASS_ROOT or populate 3rdparty/cutlass, and only pass existing include paths (cutlass_root / "include" and cutlass_root / "tools" / "util" / "include") into the load(...) call instead of blindly passing non-existent paths; refer to the variables cutlass_root_env, cutlass_root, extra_include_paths and the load(...) call to locate where to add the existence check and error raise.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@maint/gemm/correctness_evaluation_nvf4_vs_cutlass.py`:
- Around line 342-354: The FP32 reference `ref` is computed from raw decoded FP4
payloads so it doesn't include the UE4M3 block scales
(`sfa_logical`/`sfb_logical`), making `diff_tl_ref` and `diff_cutlass_ref`
invalid under NVF4_SCALE_MODE="varying"; fix by applying the per-block scales to
the decoded tensors before matmul: after calling `_decode_rowmajor_fp4(a, M, K)`
and `_decode_rowmajor_fp4(b, N, K)`, multiply each 16xK chunk of the decoded A
by the corresponding entries in `sfa_logical` and each 16xK chunk of the decoded
B by `sfb_logical` (or apply equivalent broadcasting per 16-K block) so `ref =
scaled_decoded_a @ scaled_decoded_b.T` matches the block-scaled GEMM, otherwise
remove the `diff_tl_ref`/`diff_cutlass_ref` prints until the scale-aware
reference is implemented; reference symbols: `ref`, `_decode_rowmajor_fp4`,
`sfa_logical`, `sfb_logical`, `diff_tl_ref`, `diff_cutlass_ref`.
In `@tilelang/cuda/intrinsics/macro/mma_macro_generator.py`:
- Around line 1402-1407: _scale_region_parts currently returns (buffer, [], 0,
0) for plain Buffer which drops outer-dimension bases and breaks N-D scale
buffers; change it to call and reuse _legalize_to_buffer_region(scale_buf) so
both Buffer and BufferRegion follow the same region contract as A/B. Update
_scale_region_parts to accept either type, call _legalize_to_buffer_region when
scale_buf is a Buffer to obtain the buffer and full region, and then extract the
same tuples (buffer, [r.min for r in region[:-2]], region[-2].min,
region[-1].min) as for BufferRegion; keep the existing ValueError for
unsupported types. Ensure references to _scale_region_parts usage (the indexing
code that expects preserved outer bases) continue to work unchanged.
- Around line 1282-1333: The constructor currently forces block-scale config to
NVF4 but does not validate the caller-provided dtypes, so callers can pass
incompatible a_dtype/b_dtype/accum_dtype and silently generate wrong code; after
calling _get_block_scale_mma_config(...) in __init__, add a guard that checks
the resolved self.block_scale_config.kind (and/or its expected dtype descriptors
from the config) against the incoming a_dtype, b_dtype, and accum_dtype
parameters (e.g., ensure a_dtype and b_dtype match the NVF4 fragment dtypes such
as T.float4_e2m1fn and accum_dtype matches the expected accumulator like
T.float32 or whatever the config exposes), and raise a ValueError with a clear
message if they mismatch; perform this validation before calling
super().__init__ so invalid combinations fail fast.
---
Duplicate comments:
In `@maint/gemm/correctness_evaluation_nvf4_vs_cutlass.py`:
- Around line 247-256: Check that the computed cutlass_root (from
cutlass_root_env or repo / "3rdparty" / "cutlass") actually exists before
calling load; if it does not exist, raise a clear RuntimeError instructing the
developer to set CUTLASS_ROOT or populate 3rdparty/cutlass, and only pass
existing include paths (cutlass_root / "include" and cutlass_root / "tools" /
"util" / "include") into the load(...) call instead of blindly passing
non-existent paths; refer to the variables cutlass_root_env, cutlass_root,
extra_include_paths and the load(...) call to locate where to add the existence
check and error raise.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5bd71dcd-3f86-40b2-b9d8-6c494431985b
📒 Files selected for processing (6)
maint/gemm/correctness_evaluation_nvf4_vs_cutlass.pymaint/gemm/cutlass_nvf4_ref.cusrc/cuda/codegen/codegen_cuda.ccsrc/tl_templates/cuda/instruction/mma_block_scale.htilelang/cuda/intrinsics/macro/mma_macro_generator.pytilelang/language/ast/ir.py
🚧 Files skipped from review as they are similar to previous changes (3)
- src/cuda/codegen/codegen_cuda.cc
- src/tl_templates/cuda/instruction/mma_block_scale.h
- maint/gemm/cutlass_nvf4_ref.cu
Collaborator
|
Really appreciate your contribution! I think it's better to provide high-level API for blockscaled mma on sm120, like existing |
Contributor
Author
|
Thanks, that makes sense. This PR is currently a low-level SM120 NVF4 MMA emitter. It checks that the new mma.sync instruction works, and that the fragment layout and scale-factor lane mapping are correct. I agree that users should have a higher-level API, similar to the existing SM100 blockscaled GEMM API. The SM120 version should not copy the TCGEN05/TMEM behavior directly, but it can provide a similar user experience on top of the SM120 warp-level mma.sync path. For this PR, I can move the SM120 blockscaled MMA helper into gemm_sm120.h, so it sits with the SM120 GEMM code. Then we can add a higher-level tile op API as a follow-up, unless you prefer to include that in this PR. |
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/language/tir/op.py (1)
1079-1119:⚠️ Potential issue | 🟠 MajorFix
tl.ptx_mma_block_scaleintrinsic arity:set_num_inputs(21)doesn’t match CUDA lowering expectations.
src/op/builtin.cc:244-247registerstl.ptx_mma_block_scalewith.set_num_inputs(21).src/cuda/codegen/codegen_cuda.cc:2852lowers the intrinsic withICHECK_EQ(op->args.size(), 17U)and consumesop->args[0..16](accum_dtype/shape/layouts/k/vec_size/dtypes + A/B/C pointers & offsets +scale_a/scale_b).- The Python wrapper
language/tir/op.py:1079-1119passes exactly those 17call_intrinargs, so the registration count should be aligned (likely change.set_num_inputs(21)→17).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/language/tir/op.py` around lines 1079 - 1119, The registered intrinsic tl.ptx_mma_block_scale has a mismatched arity: the Python wrapper function ptx_mma_block_scale builds 17 call_intrin arguments and the CUDA lowering asserts ICHECK_EQ(..., 17), so update the intrinsic registration to use .set_num_inputs(17) (replace the current .set_num_inputs(21)) so the registration count matches the call_intrin arguments and the CUDA lowering expectations.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tilelang/language/tir/op.py`:
- Around line 1079-1119: The registered intrinsic tl.ptx_mma_block_scale has a
mismatched arity: the Python wrapper function ptx_mma_block_scale builds 17
call_intrin arguments and the CUDA lowering asserts ICHECK_EQ(..., 17), so
update the intrinsic registration to use .set_num_inputs(17) (replace the
current .set_num_inputs(21)) so the registration count matches the call_intrin
arguments and the CUDA lowering expectations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 82a9e9cc-83ba-44c4-acfa-79f6cc76f5fa
📒 Files selected for processing (5)
src/cuda/codegen/codegen_cuda.ccsrc/tl_templates/cuda/gemm_sm120.htesting/python/language/test_tilelang_language_nvf4_mma_block_scale.pytilelang/cuda/intrinsics/macro/mma_macro_generator.pytilelang/language/tir/op.py
💤 Files with no reviewable changes (1)
- tilelang/cuda/intrinsics/macro/mma_macro_generator.py
🚧 Files skipped from review as they are similar to previous changes (1)
- testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py
Collaborator
|
Nice work! May I ask whether we support TMA load for operands and SFs? If so, we may also illustrate that in our example. Also, I'm curious whether you've considered warp-specialization (either automatic version or handwritten) and its influence on the performance. |
Contributor
Author
|
Updated status: this PR now includes TMA loading for packed FP4 A/B and SFA/SFB, plus a warp-specialized SM120a NVFP4 performance path. Please see the latest compact status note for current validation and performance numbers: #2364 (comment) |
Contributor
Author
|
Current SM120 NVFP4 block-scaled GEMM status after the latest cleanup/update: This PR adds the first TileLang SM120a NVFP4 block-scaled GEMM path with:
Validation on Additional short-loop perf sanity check, using The remaining gap is not from missing outer-level TMA or warp specialization support. Our current understanding is that the next optimization stage needs deeper CUTLASS/CuTe-style consumer-side modeling inside TileLang lowering: A/B/SFA/SFB Durable PR status observed for this head: I am not pinning GitHub Actions state here because it is time-varying; please use the PR checks tab for the current Actions run state. |
Contributor
Author
|
@Rachmanino |
qqq-tao
force-pushed
the
nvf4-block-scale-sm120-pr-clean
branch
2 times, most recently
from
e736713 to
c3596f5
Compare
Rachmanino
reviewed
Jul 3, 2026
| from .macro.mma_macro_generator import ( # noqa: F401 | ||
| TensorCoreIntrinEmitter, | ||
| TensorCoreIntrinEmitterWithLadderTransform, | ||
| TensorCoreIntrinEmitterWithBlockScale, |
Collaborator
There was a problem hiding this comment.
Better to extend the original TensorCoreIntrinEmitter for mma to support blockscaling mode, which aligns with the impl of TCGEN05MMAIntricEmitter.
Collaborator
There was a problem hiding this comment.
Personally I think these 2 examples are a bit too complicated.
- I think we can make use of the auto ws provided by TileLang instead of handwritting ws
- For an illustrative example, i think non-persistent is sufficient
T.sm120_store_full_c_fragment_panel32_tma_bf16it is not very proper to expose such procedure as an API undertilelang.language, which usually only includes essential operators or intrinisics. A better alternative way it to consider how to lower to it automatically- For simplicity, kernel implementation, along with a simple benchmark result demonstration are sufficient for an example. We can remove stuff like finding nvcc/cmake/cutlass and so on, or consider moving them to the maint folder.
Collaborator
There was a problem hiding this comment.
If i understand correctly, multiple intrinsics exposed here that started with sm120_mma only differ in their micro pipeline strategies (i.e. how ldmatrix and atom mma are composed into a tiled blockscaled mma?). I think instead of exposing so many intrinsics, a unified sm120_mma_blockscaled API containing a parameter which controls the micro pipeline impl is cleaner and easier to maintain
| namespace { | ||
|
|
||
| bool EnableSM120CompactUnpackedFP4Shared() { | ||
| const char *value = std::getenv("TL_SM120_COMPACT_UNPACKED_FP4_SHARED"); |
Collaborator
There was a problem hiding this comment.
why we need TL_SM120_COMPACT_UNPACKED_FP4_SHARED and TL_SM120_MMA_RAW_UNPACKED_FP4_ACCESS_PTR
Contributor
Author
|
@Rachmanino Thank you for your review and valuable comments. I am fixing it. |
Remove the public micro_pipeline strategy knob from T.mma_gemm_blockscaled and route the packed-scale SM120 path internally through the package-pingpong lowering. Hide SM120 strategy helpers from the public TIR API and remove stale emitter experiments that were only used by the strategy sweep. Restore the optimized persistent SM120 NVFP4 example as the benchmark path, but keep the language surface clean by using sf_layout to select the packed scale contract and a regular T.copy epilogue instead of private C-fragment store APIs. Validation: pre-commit run --files on changed files; pytest testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py testing/python/language/test_tilelang_sm120_nvfp4_example_cli.py -q => 89 passed; 512^3 correctness passed; 8192^3 128x128x256 stages=2 timing => 0.7198 ms, 1527.57 TFLOPS.
Move the persistent CUTLASS comparison harness to maint/gemm and keep the public SM120 NVFP4 example focused on a non-persistent TileLang kernel using T.mma_gemm_blockscaled and normal T.copy output. Add a source-level guard so the public example does not grow CUTLASS build or subprocess benchmark plumbing again.
The SM120 fulltile package path now has a single production behavior: row-major shared A/B views and the packed-U4 ldmatrix.b16 operand fetch. Drop the TL_SM120_FULLTILE_* configuration macros (ROWMAJOR_VIEW, FP4_LDSM, FP4_REG_SHIFT, CUTE_ACCUM_*) together with the dead compact view helpers, and stop injecting -D flags from the example and the maint benchmark. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Merging the SM120 block-scale mode into TensorCoreIntrinEmitter changed mma()'s 4th/5th positional parameters to SFA_buf/SFB_buf, so the plain gemm lowering's positional mma(A_local, B_local, C_buf, ki) call bound ki to SFA_buf and raised 'Scale buffers require TensorCoreIntrinEmitter block-scale mode' for every non-blockscaled MMA GEMM (CI examples: attention_sink, deepseek_nsa/mhc/v32). Keep the base (A, B, C, k_inner) positional signature, make all scale-related parameters keyword-only, forward k_inner to the base implementation, and drop the unused blockscaled mma_atom override that shadowed the base method the same way. Add a signature-contract regression test. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
blockscaled_chunk_kmajor packed words are byte-identical to the CUTLASS/CuTeDSL canonical blocked SF layout (tile_atom_to_shape_SF, sf_vec_size=16, atom ((32,4),(16,4)):((16,4),(0,1)) tiled with order (2,1,3)); the two views bridge zero-copy via torch view casts. Document that on the packer and lock it with a cross-layout contract test. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Back-traced every symbol the example and the maint WS benchmark pull in, then walked transitive reachability from the production lowering (gemm_mma.py) and tests: - TensorCoreIntrinEmitter: drop 16 unreachable micro-pipeline strategy variants left over from bring-up (prefetched-scale orderings, owner / owner-wide / kpack scale copies, serpentine issue orders, selector probe) — the production path uses ldmatrix_a/b, mma, ldscale(_fragment), mma_full_b_atom_with_scale_fragments and the fulltile package backend. - gemm_sm120.h: drop two device helpers nothing emits (dual-issue same-B mma wrapper and the accum-inplace variant). - quantize/nvfp4.py: drop the one quantize kernel variant with no dispatch site (32x256 tma_io). No public surface change; language API stays T.mma_gemm_blockscaled + ptx_mma_block_scale. Validated with a fresh JIT cache: nvf4 language + example CLI + quantize layout + atom mma tests (122 passed), example and WS benchmark --verify both pass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The exact-bits shared storage sizing previously applied to every dtype with element_bits % 8 != 0 (int4/uint4/uint2/...), halving e.g. int4 shared allocations while SharedByteOffsetToLogicalIndexOffset kept the legacy bytes-per-element conversion for those dtypes — an inconsistent size/index pair for any non-FP4 sub-byte buffer in a merged arena. Restrict the new sizing to float4_e2m1fn scalar so it mirrors the byte-offset -> logical-index special case exactly: NVFP4 buffers get packed two-per-byte sizing and indexing, every other dtype keeps bit-exact upstream behavior. This pass now changes nothing outside the NVFP4 path. Validated after rebuilding libtilelang.so: nvf4 language + example CLI + quantize layout + access_ptr codegen + atom mma tests (129 passed), example and WS benchmark --verify pass with a fresh JIT cache. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
qqq-tao
force-pushed
the
nvf4-block-scale-sm120-pr-clean
branch
from
32cf205 to
be668c6
Compare
The example and the maint warp-specialized benchmark each carried a private macro with the blockscaled_chunk_kmajor tile addressing spelled out by hand. Move that knowledge into the library next to T.mma_gemm_blockscaled: - T.ue4m3_scale_tile_source_coords: the single source of the packed scale tile addressing (tiles are contiguous in this layout); words_per_kblock stays a parameter so future contracts (block_K != 256, MX scale packings) reuse it. - T.copy_ue4m3_scale_tile: cooperative whole-block staging macro with trace-time dtype/shape validation, used by the example. The scale operand stays visible in user kernels; only the layout math is hidden. The maint kernel keeps its producer-warp-group staging schedule and reuses the coords helper. Validated: nvf4 language + example CLI tests (91 passed incl. a new validation test), example and WS benchmark --verify pass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Production shapes rarely keep M (batch x seq) a multiple of 128. Lift the divisibility requirements to the validated envelope: - scale packers zero-pad rows to full 128-row tiles (the same ceil convention CUTLASS/cuBLAS/CuTeDSL use for the blocked SF tensor), so the packed source always has ceil(rows/128)*128 rows; - the example accepts an arbitrary M or an N tail: TMA loads zero-fill out-of-bounds rows and the C store is predicated. N keeps the 16-byte-aligned-rows rule (N % 8 == 0), the same contiguous-dim alignment CuTeDSL validates; - simultaneous M and N tails are rejected with a clear error: the fragment-to-global copy lowering emits the column guard bound from the wrong dimension in that case (upstream bug, tracked separately). Tests: tail-shape runtime correctness (257x384, 130x128, 128x136 with exact binary-scale references), rejection paths, and packer padding round-trip. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Two structural reductions, no behavior change: - Merge the three TMA-staged quantize kernel variants (16/32/64x256) into one _tilelang_nvfp4_blockscaled_quantize_kernel_tiled factory: their inner loops were line-identical, differing only in rows_per_cta, stage count and the store path (TMA store vs direct global stores), all of which are trace-time parameters now. Outputs verified bitwise-identical to the previous per-variant kernels for every dispatch configuration (quant bytes and scale source). - Move the 14 SM120 OMMA scale-selector contract oracles into the contract test file that is their only consumer; the quantize module stays focused on quantization and layout packing. 1334 -> 978 lines. Quantize suite: 32 passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
nvfp4.py: remove the two remaining bring-up era quantize kernels (the legacy no-TMA-load 16x256 schedule and the k64_per_cta=1 generic path, reachable only through tuning knobs nothing uses) and the use_tma_load / k64_per_cta parameters. The dispatcher now maps straight onto the single tiled kernel factory; outputs verified bitwise-identical to the previous code for every remaining configuration. 978 -> 710 lines (1334 before the slimming series). Example: drop the ones / constant-scale debug input modes and their CLI flags; the random-FP4 + random-binary-scale path is the only verification mode and checks exact equality. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
tilelang/quantize's identity is device-side conversion helpers plus host-side data-prep utilities (quantization.py's _tir_* converters, the lop3/mxfp intrin groups); a complete quantizer operator does not belong in the library. Move the TileLang BF16->NVFP4 kernel factory and its dispatcher to maint/gemm/tilelang_nvfp4_quantizer.py, next to the maint benchmark that consumes it. nvfp4.py now carries exactly the format contract: the _tl_* device helpers (FP4/UE4M3 encode/decode, scale-source coordinates) any user kernel can call, the blockscaled_chunk_kmajor packers/swizzles, and the torch reference quantizer — 524 lines (1334 at series start). The package re-export list shrinks to the four user-facing names; the alias wrappers are gone. Validated: quantize + example CLI suites (35 passed), maint benchmark and example --verify pass end-to-end. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Follow-up to the package export slimming: blockscaled_chunk_kmajor_word_offset is no longer re-exported at the tilelang.quantize package level. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[CUDA] Add SM120 NVFP4 block-scale MMA support
This PR adds the first TileLang GEMM path for SM120(a) (GeForce Blackwell,
e.g. RTX 5090) NVFP4 block-scaled tensor cores:
Design (shaped by review feedback — thanks @Rachmanino)
T.mma_gemm_blockscaled(A, B, C, SFA, SFB, ...),following the same scale-factor model as
T.tcgen05_gemm_blockscaled(users pass scale tensors,
k_start, K granularity; lowering derives theaddressing). The only new low-level intrinsic is
ptx_mma_block_scale,mirroring the existing
ptx_mma.TensorCoreIntrinEmittergains anis_blockscaledmode (aligned with how TCGEN05 integrates blockscaling);micro-pipeline selection happens inside the lowering.
examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.pyis a non-persistent tiled kernel using plain
T.Kernel/T.Pipelined/T.copy— no handwritten warp specialization, no special compile flags.A hand-written warp-specialized performance kernel (1 producer + 2
consumer warp groups) lives in
maint/gemm/benchmark_sm120_nvfp4_blockscaled_gemm.py.T.copy_ue4m3_scale_tile(SFA, SFA_shared, by, ko)call (the scale operandstays visible in the kernel), but the tile addressing of the packed layout
lives in one place; kernels with custom staging schedules (like the maint
warp-specialized kernel) reuse the lower-level
T.ue4m3_scale_tile_source_coords. Folding the staging into the op'slowering entirely is listed as follow-up work.
tilelang.quantizestays device-helper territory: the library gainsthe NVFP4 format contract only —
_tl_*device-side FP4/UE4M3encode/decode helpers usable from any kernel (mirroring the existing
_tir_*converters), theblockscaled_chunk_kmajorpackers/swizzles,and a torch reference quantizer. The tuned TileLang quantizer kernel
lives in
maint/gemm/tilelang_nvfp4_quantizer.pynext to the benchmarkthat consumes it.
Scale layout contract & interoperability
The scale operand uses the CUTLASS
BlockScaledBasicChunkK-major layout,exposed as
uint32 [MN, K/64](byte-identical to the canonical blocked SFlayout used by CUTLASS, CuTeDSL and cublasLt; the views convert zero-copy).
A contract test pins the packed bytes against an independently implemented
layout oracle. We validated interoperability by feeding the same device
bytes to this PR's kernels, CUTLASS (
CollectiveBuilder+KernelTmaWarpSpecializedPingpong), the CuTeDSLblackwell_geforceexample, and cuBLAS (
torch._scaled_mm-> cublasLt): all producebitwise-identical bf16 outputs across shapes, including adversarial
position-encoded inputs covering all FP4 codes (negatives included) and
100+ distinct scale values.
Supported instruction & future work
SM120
mmasupports several narrow-precision kinds (PTX ISA, warp-levelmmawith block scaling). This PR implements the NVFP4 variant only; theconfig registry is keyed by
(kind, scale_vec, scale type)so the otherscan be added incrementally:
kind::mxf4nvf44Xkind::mxf4nvf4/kind::mxf42Xkind::mxf8f6f41Xkind::f8f6f4Performance
Measured on GeForce RTX 5090 (driver 570.124.06, CUDA 12.8): CUDA-event
timing, 5 warmup + 50 iterations, median of 3 interleaved rounds at thermal
steady state (3 s preheat; round spread <= 1%), GPU held exclusively. All
implementations consume identical device bytes and produce bitwise-identical
outputs, so this is an apples-to-apples comparison. CUTLASS/CuTeDSL are shown
at their best tile config per shape; TileLang kernels use 128x128x256.
The simple example stays within ~5% of the best implementation at every
size; the warp-specialized maint kernel is fastest overall at 8192^3.
Core-pass changes are NVFP4-scoped
merge_shared_memory_allocations.cc: packed scalar NVFP4 shared buffersare sized bit-exactly (two elements per byte) and their merged-arena byte
offsets convert back to logical element indices with the matching x2 rule.
Both special cases are keyed on
float4_e2m1fnscalar only — every otherdtype (incl. int4/uint4) keeps bit-exact upstream behavior.
copy.cc: the TMA descriptor swizzle recast has an NVFP4-scalar-onlyspecial case; other dtypes unchanged.
Testing
testing/python/language/test_tilelang_language_nvf4_mma_block_scale.py— fragment/lane mapping vs CuTe, codegen, correctness (constant & varying
scales), API-surface contract tests.
testing/python/quantize/test_tilelang_quantize_nvfp4_scale_layout.py— scale-layout packers vs independent oracles (incl. the CuTeDSL blocked
layout), quantizer round-trip and error bounds.
testing/python/language/test_tilelang_sm120_nvfp4_example_cli.py— runs the example end-to-end on SM120.
exercised on local RTX 5090s (all tests above pass there); CI covers the
compile/lowering paths. Happy to help wire an SM120 runner if available.
Shape support
The packed scale source is zero-padded to full 128-row tiles by the
tilelang.quantizepackers, matching the ceil conventionCUTLASS/cuBLAS/CuTeDSL use for the blocked SF tensor.
Non-goals / follow-ups
block_K=128tiles, and a persistent tilescheduler for the wide-MN/short-K regime are follow-up work.
T.mma_gemm_blockscaledlowering(removing the explicit
T.copy_ue4m3_scale_tilecall) is follow-up work.since 2026-07-08, e.g.
fix_code) and unrelated to this PR.