Support SM121 (NVIDIA GB10 / DGX Spark) via SM120 kernels

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
 

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

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(

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

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,

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,
     )