mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-11-04 08:54:22 +08:00
qutaoandClaude Opus 4.8
c1ae801c61
T.copy_ue4m3_scale_tile and T.ue4m3_scale_tile_source_coords were NVFP4-specific helpers on the general T.* surface. Remove both and move the abstraction down a level: - The example now declares the packed scale source with its honest "tile rows" shape ([n_mn_blocks * n_k_blocks * block_MN, block_K // 64] uint32): the blockscaled_chunk_kmajor layout stores each (mn_block, k_block) tile contiguously, so with this view every tile is a plain rectangular slice and staging is a two-line T.Parallel loop with row-offset addressing only. The host obtains the view zero-copy via swizzle_blockscaled_chunk_kmajor_scale_words(...).reshape(-1, 4). - tilelang.quantize.nvfp4 gains blockscaled_chunk_kmajor_tile_source_coords, the hand-written offset formula for staging schedules that keep the 2-D [MN, K/64] view (the maint benchmark's producer warp group uses it). Compiler-internal staging was prototyped first (plain T.copy of the slice) and rejected for now: the raw-pointer bulk-TMA path this lowers to on some shapes never arrives on its mbarrier (deadlock, tracked upstream), and the synchronous fallback stalls the producer warps (-8.5% at 2048^3). The example comment documents this so the loop can become a plain T.copy once fixed. Validation on RTX 5090: 123 tests pass; example outputs bitwise-identical to the previous macro implementation on 5 shapes including M/N tail tiles (257x384x512, 130x128x256, 128x136x256); example perf at parity (100.1-100.4% over 2048^3/4096^3/8192^3/8192x8192x1536, interleaved A/B, generated CUDA identical modulo a loop variable name); maint WS benchmark kernel source byte-identical with --verify passing and perf unchanged (641.9 / 1310.5 / 1149.5 TF at 2048^3 / 8192^3 / 8192x8192x1536). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import tilelang.testing
|
|
|
|
|
|
def _load_sm120_example(monkeypatch):
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
example = repo_root / "examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.py"
|
|
monkeypatch.setattr(sys, "argv", [str(example)])
|
|
spec = importlib.util.spec_from_file_location("sm120_nvfp4_blockscaled_gemm_example", example)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def test_sm120_nvfp4_example_cli_keeps_internal_strategy_fixed(monkeypatch):
|
|
module = _load_sm120_example(monkeypatch)
|
|
args = module.parse_args()
|
|
|
|
assert not hasattr(args, "micro_pipeline")
|
|
assert not hasattr(args, "sf_layout")
|
|
assert not hasattr(args, "scale_storage_layout")
|
|
assert not hasattr(args, "manual_ws2")
|
|
|
|
assert module._SM120_SCALE_LAYOUT == "blockscaled_chunk_kmajor"
|
|
|
|
|
|
def test_sm120_nvfp4_example_source_has_no_legacy_strategy_flags():
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
source = (repo_root / "examples/gemm_sm120/sm120_nvfp4_blockscaled_gemm.py").read_text()
|
|
|
|
forbidden_flags = [
|
|
"--micro-pipeline",
|
|
"--manual-ws2",
|
|
"--manual-ws2-split",
|
|
"--manual-ws2-sf-load",
|
|
"--manual-ws2-sf-layout",
|
|
"--manual-ws2-ab-shared-storage",
|
|
"--manual-ws2-ab-copy-view",
|
|
"--scale-storage-layout",
|
|
"--sf-layout",
|
|
"--path",
|
|
"--run-cutlass",
|
|
"--cutlass-build-dir",
|
|
"--cutlass-binary",
|
|
"--rebuild-cutlass",
|
|
"--cmake",
|
|
"--nvcc",
|
|
]
|
|
for flag in forbidden_flags:
|
|
assert flag not in source
|
|
|
|
forbidden_harness_terms = [
|
|
"build_cutlass",
|
|
"subprocess",
|
|
"CMAKE_CUDA_COMPILER",
|
|
]
|
|
for term in forbidden_harness_terms:
|
|
assert term not in source
|
|
|
|
|
|
@tilelang.testing.requires_cuda_compute_version_eq(12, 0)
|
|
def test_sm120_nvfp4_example_kernel_handles_mn_tail_tiles(monkeypatch):
|
|
import pytest
|
|
|
|
torch = pytest.importorskip("torch")
|
|
|
|
from tilelang.quantize import swizzle_blockscaled_chunk_kmajor_scale_words
|
|
|
|
module = _load_sm120_example(monkeypatch)
|
|
for M, N, K in [(257, 384, 512), (130, 128, 256), (128, 136, 256)]:
|
|
kernel = module.sm120_nvfp4_blockscaled_gemm(M, N, K)
|
|
|
|
A = module._make_packed_fp4(M, K, seed=3)
|
|
B = module._make_packed_fp4(N, K, seed=4)
|
|
SFA_semantic = module._make_binary_scale_words(M, K, seed=5)
|
|
SFB_semantic = module._make_binary_scale_words(N, K, seed=6)
|
|
SFA = swizzle_blockscaled_chunk_kmajor_scale_words(SFA_semantic).reshape(-1, 4)
|
|
SFB = swizzle_blockscaled_chunk_kmajor_scale_words(SFB_semantic).reshape(-1, 4)
|
|
assert SFA.shape[0] % 128 == 0
|
|
assert SFB.shape[0] % 128 == 0
|
|
|
|
C = torch.empty((M, N), device="cuda", dtype=torch.bfloat16)
|
|
kernel(A, B, SFA, SFB, C)
|
|
torch.cuda.synchronize()
|
|
module._verify(A, B, SFA_semantic, SFB_semantic, C, torch.bfloat16)
|
|
|
|
|
|
def test_sm120_nvfp4_example_kernel_rejects_unsupported_tails(monkeypatch):
|
|
import pytest
|
|
|
|
module = _load_sm120_example(monkeypatch)
|
|
# simultaneous M and N tails hit a known copy-lowering boundary bug
|
|
with pytest.raises(ValueError, match="simultaneous M and N tail"):
|
|
module.sm120_nvfp4_blockscaled_gemm(257, 136, 512)
|
|
# bf16 output rows must stay 16-byte aligned
|
|
with pytest.raises(AssertionError, match="multiple of 8"):
|
|
module.sm120_nvfp4_blockscaled_gemm(128, 130, 256)
|