diff --git a/README.md b/README.md index efa2375..bfe85b7 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Key features: Requirements: -- SM90, SM100, SM103 or SM120 +- SM90, SM100, SM103, SM120 or SM121 - CUDA 12.8 or above - PyTorch 2.8 or above diff --git a/flash_qla/ops/gated_delta_rule/chunk/__init__.py b/flash_qla/ops/gated_delta_rule/chunk/__init__.py index 698e18f..051a7d1 100644 --- a/flash_qla/ops/gated_delta_rule/chunk/__init__.py +++ b/flash_qla/ops/gated_delta_rule/chunk/__init__.py @@ -17,14 +17,14 @@ from .blackwell import get_warmup_chunks, get_warmup_chunks_bidi, correct_initial_states, correct_terminal_states from .blackwell.cp_bwd import fused_gdr_dh_ws as fused_gdr_dh CHUNK_SIZE = 64 -elif tilelang.contrib.nvcc.get_target_compute_version() == "12.0": +elif tilelang.contrib.nvcc.get_target_compute_version() in ["12.0", "12.1"]: from .blackwell_sm120 import fused_gdr_fwd, fused_gdr_h, kkt_solve from .blackwell_sm120 import get_warmup_chunks, get_warmup_chunks_bidi, correct_initial_states, correct_terminal_states fused_gdr_bwd = None fused_gdr_dh = None CHUNK_SIZE = 32 else: - raise ValueError(f"FlashQLA now support sm90, sm100, sm103 and sm120 only. Found compute version: {tilelang.contrib.nvcc.get_target_compute_version()}") + raise ValueError(f"FlashQLA now support sm90, sm100, sm103, sm120 and sm121 only. Found compute version: {tilelang.contrib.nvcc.get_target_compute_version()}") from .cp_context import intra_card_cp_preprocess, intra_card_cp_preprocess_bwd, _calc_cp_seqs, _create_cu_seqlens from flash_qla.utils import input_guard diff --git a/flash_qla/ops/gated_delta_rule/chunk/cp_context.py b/flash_qla/ops/gated_delta_rule/chunk/cp_context.py index 5a265e6..645332f 100644 --- a/flash_qla/ops/gated_delta_rule/chunk/cp_context.py +++ b/flash_qla/ops/gated_delta_rule/chunk/cp_context.py @@ -21,12 +21,12 @@ from .blackwell import get_warmup_chunks, get_warmup_chunks_bidi, fused_gdr_h, correct_initial_states, correct_terminal_states from .blackwell.cp_bwd import fused_gdr_dh_ws as fused_gdr_dh ARCH = "SM103" -elif tilelang.contrib.nvcc.get_target_compute_version() == "12.0": +elif tilelang.contrib.nvcc.get_target_compute_version() in ["12.0", "12.1"]: from .blackwell_sm120 import get_warmup_chunks, get_warmup_chunks_bidi, fused_gdr_h, correct_initial_states, correct_terminal_states fused_gdr_dh = None ARCH = "SM120" else: - raise ValueError(f"FlashQLA now support sm90, sm100 and sm103 only. Found compute version: {tilelang.contrib.nvcc.get_target_compute_version()}") + raise ValueError(f"FlashQLA now support sm90, sm100, sm103, sm120 and sm121 only. Found compute version: {tilelang.contrib.nvcc.get_target_compute_version()}") MULTI_PROCESSOR_COUNT = torch.cuda.get_device_properties().multi_processor_count @@ -123,7 +123,7 @@ def _calc_cp_seqs( Be * H <= 32 and max(num_chunks) >= 192 ) else: - raise ValueError(f"FlashQLA now support sm90, sm100 and sm103 only. Found compute version: {tilelang.contrib.nvcc.get_target_compute_version()}") + raise ValueError(f"FlashQLA now support sm90, sm100, sm103, sm120 and sm121 only. Found compute version: {tilelang.contrib.nvcc.get_target_compute_version()}") if use_cp: cp_cu_seqlens = torch.tensor( diff --git a/tests/conftest.py b/tests/conftest.py index 92f84d9..a8b7dfc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,7 +19,7 @@ ARCH = "SM100" elif _cv == "10.3": ARCH = "SM103" - elif _cv == "12.0": + elif _cv in ["12.0", "12.1"]: ARCH = "SM120" except Exception: pass diff --git a/tests/test_gdr_unit.py b/tests/test_gdr_unit.py index 0607e0a..bdead57 100644 --- a/tests/test_gdr_unit.py +++ b/tests/test_gdr_unit.py @@ -24,7 +24,7 @@ DETERMINISM_ITERS = 1000 HEAD_DIM_K = 128 HEAD_DIM_V = 128 -CHUNK_SIZE = 32 if tilelang.contrib.nvcc.get_target_compute_version() == "12.0" else 64 +CHUNK_SIZE = 32 if tilelang.contrib.nvcc.get_target_compute_version() in ["12.0", "12.1"] else 64 REF_DTYPE = torch.float64 DATA_DTYPE = torch.bfloat16 DEVICE = "cuda" @@ -781,6 +781,8 @@ def test_bwd_cp_cache( @pytest.mark.parametrize("bwd_cp", [False, True], ids=["bwd_no_cp", "bwd_cp"]) def test_mixed_cp_control(state_v_first, fwd_cp, bwd_cp): """Mixing auto_cp=True/False between fwd and bwd should still match reference.""" + if tilelang.contrib.nvcc.get_target_compute_version() in ["12.0", "12.1"]: + pytest.skip("Backward is disabled on SM120/SM121") B, T, Hk, Hv = 1, 32768, 4, 4 ( q, k, v, g, beta, do, diff --git a/tests/test_sm120_backward_disabled.py b/tests/test_sm120_backward_disabled.py index 76a4a5b..b4877a7 100644 --- a/tests/test_sm120_backward_disabled.py +++ b/tests/test_sm120_backward_disabled.py @@ -219,8 +219,8 @@ def test_cp_preprocess_bwd_raises_not_implemented(): @pytest.mark.parametrize("state_v_first", [False, True], ids=["kv", "vk"]) def test_bwd_works_on_non_sm120(state_v_first): """Sanity check: backward must *not* raise on supported architectures.""" - if tilelang.contrib.nvcc.get_target_compute_version() == "12.0": - pytest.skip("Backward is disabled on SM120 — this test only runs on SM90/SM100/SM103") + if tilelang.contrib.nvcc.get_target_compute_version() in ["12.0", "12.1"]: + pytest.skip("Backward is disabled on SM120/SM121 — this test only runs on SM90/SM100/SM103") q, k, v, g, beta, do, h0, dht, scale = _make_inputs( use_h0=True, state_v_first=state_v_first, )