mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-11-01 08:23:24 +08:00
qutao
e84c318ef4
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.
47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
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",
|
|
]
|
|
for flag in forbidden_flags:
|
|
assert flag not in source
|