From a5164242d3dfd00b3cdcde6006c9875037cb2e8c Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Sun, 7 Jun 2026 17:37:30 -0700 Subject: [PATCH 1/6] test: fuzzer coverage from 9.18-9.24 fixed-bug mining MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Derived from a triage of the 134 fixed front-end bugs in cuDNN 9.18-9.24. - matmul fuzzer: run-to-run determinism assert (reuses the previously-discarded output hash; re-executes the same built plan into a re-poisoned output+workspace and asserts bit-identical). Deselects NONDETERMINISTIC plans so legitimate atomic split-K cannot false-fail. Env: MATMUL_DET_RERUNS / MATMUL_NUM_TESTS / MATMUL_FUZZ_SEED. - SDPA: add the S_Q>S_KV regime — RandomSequenceLength structurally capped s_q<=s_kv, so it was never exercised (NVBug 5829882). Clamped to s_q_max; wired into 9 suites. Env: MHAS_NUM_TESTS / MHAS_SEED_OFFSET. - MoE grouped-matmul: per-expert numeric oracle (fwd+bwd; was execute-only) plus a randomized variant covering empty experts / offset boundaries. - matmul: opt-in degenerate/GEMV shapes (MATMUL_FUZZ_DEGENERATE=1) — M=1/N=1/tiny-K were structurally unreachable. Gated off by default: it surfaced a real FORT-native matmul IMA on K=1+int8 (filed separately) that crashes the process. Co-Authored-By: Claude Opus 4.8 (1M context) --- test/python/sdpa/random_config.py | 19 ++- test/python/test_matmul_fuzzer.py | 81 ++++++++++-- test/python/test_mhas_v2.py | 24 ++-- test/python/test_moe_grouped_matmul.py | 163 +++++++++++++++++++++++++ 4 files changed, 265 insertions(+), 22 deletions(-) diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index 3de2100b0..adc661518 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -9,8 +9,14 @@ # fmt: off def generate_test_seeds(*, num_tests, rng_seed): + # Env overrides for focused/expanded sweeps without editing each suite: + # MHAS_NUM_TESTS - cap (or expand) the number of configs (e.g. small subset for sanitizer runs) + # MHAS_SEED_OFFSET - shift the geometry RNG seed to explore a fresh set of configs + import os + rng_seed = rng_seed + int(os.environ.get("MHAS_SEED_OFFSET", "0")) + n = int(os.environ.get("MHAS_NUM_TESTS", "0")) or num_tests rng = random.Random(rng_seed) - return [(i+1, num_tests, rng.randint(65536, 2147483647)) for i in range(num_tests)] + return [(i+1, n, rng.randint(65536, 2147483647)) for i in range(n)] def get_strides_from_indices(shape, indices=[0, 1, 2, 3], gaps=[0, 0, 0, 0], rng_geom=None): @@ -498,6 +504,7 @@ def __init__( s_kv_min = min(s_kv_min, s_kv_max) self.s_q_gen = RandomIntValue(min=s_q_min, max=s_q_max) self.s_kv_gen = RandomIntValue(min=s_kv_min, max=s_kv_max) + self._s_q_max = s_q_max # post-cap upper bound, used to clamp the s_q>s_kv branch self.distribution = RandomChoice(s_q_distribution) def __call__(self, rng): @@ -510,6 +517,16 @@ def __call__(self, rng): s_q = 1 elif distribution == "s_q=s_kv": s_q = s_kv + elif distribution == "s_q>s_kv": + # Query longer than key/value (cross-attention, chunked prefill). + # The default "s_q=random" branch below structurally caps s_q<=s_kv, so + # this regime is otherwise never exercised. Caught class: NVBug 5829882 + # (SDPA hang when S_Q > S_KV). + # Clamp to s_q_max so we don't overshoot the suite bound (esp. the SM80 + # OOM cap) and trigger spurious torch CUDA OOM instead of testing cuDNN. + if s_kv < self._s_q_max: + s_q = min(s_kv + rng.randint(1, max(1, s_kv)), self._s_q_max) + # else: s_kv already at the cap -> no room for s_q>s_kv within bounds; leave s_q. else: s_q = self.s_q_gen(rng) diff --git a/test/python/test_matmul_fuzzer.py b/test/python/test_matmul_fuzzer.py index 940ae1ae3..b7d7cdc30 100644 --- a/test/python/test_matmul_fuzzer.py +++ b/test/python/test_matmul_fuzzer.py @@ -326,6 +326,31 @@ def generate(self) -> MatmulConfig: N = self.random_dim() K = self.random_dim() + # Degenerate / GEMV shapes. random_dim() squares a sqrt-range int then rounds up + # to a multiple of 8, so M/N/K are always >= 8 and never 1 — the matrix-vector + # (M=1 / N=1) and tiny-K corners are structurally unreachable. Force them ~12% of + # the time. Caught class: NVBug 5866266 (force_jit core dump on a degenerate + # matmul) and 5990039 (no engine for IDENTITY/degenerate matmul). + # + # OPT-IN (default off): enabling this surfaced a real FORT-native matmul illegal + # memory access on K=1 with an int8 operand + large M (SM90, dev 9.30 + rel 9.24; + # root-caused to get_kernel_access_size() unclamped — bug filed separately). An IMA + # corrupts the CUDA context and kills the test process (it cannot be caught as a + # skip), so this axis is gated behind MATMUL_FUZZ_DEGENERATE=1 until that backend + # bug is fixed, after which the default can be flipped on. + if os.environ.get("MATMUL_FUZZ_DEGENERATE", "0") == "1" and self.rng.random() < 0.12: + kind = self.rng.choice(['m1', 'n1', 'tiny_k', 'm1_tiny_k', 'k1']) + if kind == 'm1': + M = 1 + elif kind == 'n1': + N = 1 + elif kind == 'tiny_k': + K = self.rng.choice([1, 2, 4, 8, 16]) + elif kind == 'm1_tiny_k': + M = 1; K = self.rng.choice([1, 8, 16]) + else: # k1 + K = 1 + # Data types - ensure compatible combinations if self.rng.random() < 0.8: # 80% of tests use same dtype for A and B (more stable) @@ -492,8 +517,19 @@ def compute_reference(config: MatmulConfig, A: torch.Tensor, B: torch.Tensor, bi def run_cudnn_matmul(config: MatmulConfig, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, - bias: Optional[torch.Tensor], cudnn_handle) -> Tuple[bool, str]: - """Run matmul using cuDNN and return success status and message.""" + bias: Optional[torch.Tensor], cudnn_handle, + det_reruns: Optional[int] = None) -> Tuple[bool, str]: + """Run matmul using cuDNN and return success status and message. + + If det_reruns > 0, re-execute the *same built plan* that many extra times into a + freshly garbage-poisoned output (and re-poisoned workspace) and assert the output + hash is identical run-to-run. This catches nondeterministic kernels (lost-update / + smem-pipeline races, uninitialized split-K workspace) without needing a golden ref + and without re-running heuristics. Caught class: NVBug 5897577 (sm120 matmul buffer + race), 6140341 (matmul-fuzzer numeric regression), 6217343-class. + """ + if det_reruns is None: + det_reruns = int(os.environ.get("MATMUL_DET_RERUNS", "2")) try: stream = torch.cuda.current_stream().cuda_stream cudnn.set_stream(handle=cudnn_handle, stream=stream) @@ -544,6 +580,15 @@ def run_cudnn_matmul(config: MatmulConfig, A: torch.Tensor, B: torch.Tensor, C: # Build and execute graph.validate() graph.build_operation_graph() + if det_reruns > 0: + # The determinism rerun-assert below must only judge plans cuDNN declares + # deterministic: a legitimate FP-atomic split-K plan carries the + # NONDETERMINISTIC numeric note and is *expected* to vary run-to-run, so + # deselect those here to avoid a spurious determinism failure. + try: + graph.deselect_numeric_notes([cudnn.numerical_note.NONDETERMINISTIC]) + except Exception: + pass graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) graph.check_support() graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) @@ -551,11 +596,15 @@ def run_cudnn_matmul(config: MatmulConfig, A: torch.Tensor, B: torch.Tensor, C: # Allocate workspace and fill with garbage to catch uninitialized memory bugs workspace_size = graph.get_workspace_size() workspace = torch.empty(workspace_size, device='cuda', dtype=torch.uint8) - if workspace_size > 0: - # Fill with random garbage + some NaN patterns to test proper workspace init - workspace.random_(0, 256) - nan_mask = torch.rand(workspace_size, device='cuda') < 0.1 - workspace[nan_mask] = 0xFF + + def poison_workspace(): + if workspace_size > 0: + # Fill with random garbage + some NaN patterns to test proper workspace init + workspace.random_(0, 256) + nan_mask = torch.rand(workspace_size, device='cuda') < 0.1 + workspace[nan_mask] = 0xFF + + poison_workspace() # Build variant pack variant_pack = {A_tensor: A, B_tensor: B, result: C} @@ -566,6 +615,20 @@ def run_cudnn_matmul(config: MatmulConfig, A: torch.Tensor, B: torch.Tensor, C: graph.execute(variant_pack, workspace, handle=cudnn_handle) torch.cuda.synchronize() + # Determinism check: re-execute the SAME plan into a re-poisoned output and + # workspace; the output hash must be bit-identical run-to-run. + if det_reruns > 0: + h0 = print_tensor_stats(C, tag=None) + for _ in range(det_reruns): + fill_with_garbage(C) # re-poison output: kernel must fully overwrite + poison_workspace() # re-poison workspace: catches uninit split-K + graph.execute(variant_pack, workspace, handle=cudnn_handle) + torch.cuda.synchronize() + hi = print_tensor_stats(C, tag=None) + if hi != h0: + return False, (f"NONDETERMINISTIC output: hash 0x{h0 >> 32:08X} != " + f"0x{hi >> 32:08X} across reruns (same plan, re-poisoned)") + return True, "success" except cudnn.cudnnGraphNotSupportedError as e: @@ -711,8 +774,8 @@ def get_test_params(request): # Fixed test list for default runs -DEFAULT_NUM_TESTS = 2048 -DEFAULT_SEED = 42 +DEFAULT_NUM_TESTS = int(os.environ.get("MATMUL_NUM_TESTS", "2048")) +DEFAULT_SEED = int(os.environ.get("MATMUL_FUZZ_SEED", "42")) TEST_PARAMS = tlist_with_configs(num_tests=DEFAULT_NUM_TESTS, rng_seed=DEFAULT_SEED) diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 191f205b1..67991588b 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -109,7 +109,7 @@ def test_sdpa_random_fwd_L0(env_info, test_no, request, cudnn_handle): # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=1, max=8, with_high_probability=[1,4]), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=256, d_v_min=1, d_v_max=256, head_dim_distribution={"d_qk=d_v":1, "d_qk=random":1}, with_high_probability=[(64,64), (128,128), (192,128), (256, 256)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -141,7 +141,7 @@ def test_sdpa_random_fwd_unified_L1(env_info, test_no, request, cudnn_handle): # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=1, max=8, with_high_probability=[1,4]), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=256, d_v_min=1, d_v_max=256, head_dim_distribution={"d_qk=d_v":1, "d_qk=random":1}, with_high_probability=[(64,64), (128,128), (192,128), (256, 256)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -188,7 +188,7 @@ def test_sdpa_random_bwd_L0(env_info, test_no, request, cudnn_handle): # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=8, max=16), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=192, d_v_min=1, d_v_max=128, head_dim_distribution={"d_qk=d_v":5, "d_qk=random":1}, with_high_probability=[(64,64), (128,128), (192,128), (256,256)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -369,7 +369,7 @@ def test_sdpa_random_fwd_ragged_L0(env_info, test_no, request, cudnn_handle): # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=1, max=8, with_high_probability=[1,4]), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=256, d_v_min=1, d_v_max=256, head_dim_distribution={"d_qk=d_v":1, "d_qk=random":1}, with_high_probability=[(64,64), (128,128), (192,128), (256, 256)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -401,7 +401,7 @@ def test_sdpa_random_fwd_ragged_unified_L1(env_info, test_no, request, cudnn_han # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=1, max=8, with_high_probability=[1,4]), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=256, d_v_min=1, d_v_max=256, head_dim_distribution={"d_qk=d_v":1, "d_qk=random":1}, with_high_probability=[(128,128), (192,128), (256, 256)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -443,7 +443,7 @@ def test_sdpa_random_fwd_ragged_offset_multiplier_unified_L1(env_info, test_no, # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=1, max=8, with_high_probability=[1,4]), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=256, d_v_min=1, d_v_max=256, head_dim_distribution={"d_qk=d_v":1, "d_qk=random":1}, with_high_probability=[(128,128), (192,128), (256, 256)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -477,7 +477,7 @@ def test_sdpa_random_bwd_ragged_L0(env_info, test_no, request, cudnn_handle): # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=8, max=16), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=192, d_v_min=1, d_v_max=128, head_dim_distribution={"d_qk=d_v":5, "d_qk=random":1}, with_high_probability=[(64,64), (128,128), (192,128), (256,256)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -514,7 +514,7 @@ def test_sdpa_fwd_paged_L0(env_info, test_no, request, cudnn_handle): # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=1, max=8, with_high_probability=[1,4]), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=64, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=64, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=128, d_v_min=1, d_v_max=128, head_dim_distribution={"d_qk=d_v":1, "d_qk=random":1}, with_high_probability=[(64,64), (128,128), (192,128)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -548,7 +548,7 @@ def test_sdpa_fwd_paged_unified_L0(env_info, test_no, request, cudnn_handle): # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=1, max=8, with_high_probability=[1,4]), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=64, s_kv_min=1, s_kv_max=512, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=64, s_kv_min=1, s_kv_max=512, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=128, d_v_min=1, d_v_max=128, head_dim_distribution={"d_qk=d_v":1, "d_qk=random":1}, with_high_probability=[(128,128), (192,128)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -583,7 +583,7 @@ def test_sdpa_random_fwd_unified_block_mask_L0(env_info, test_no, request, cudnn # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=1, max=8, with_high_probability=[1,4]), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=128, d_v_min=1, d_v_max=128, head_dim_distribution={"d_qk=d_v":1, "d_qk=random":1}, with_high_probability=[(128,128), (192,128)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -617,7 +617,7 @@ def test_sdpa_random_fwd_bias_L0(env_info, test_no, request, cudnn_handle): # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=1, max=8, with_high_probability=[1,4]), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=128, d_v_min=1, d_v_max=128, head_dim_distribution={"d_qk=d_v":1, "d_qk=random":1}, with_high_probability=[(64,64), (128,128), (192,128)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), @@ -650,7 +650,7 @@ def test_sdpa_random_bwd_bias_L0(env_info, test_no, request, cudnn_handle): # Create the randomization context within the test with RandomizationContext( batches=RandomBatchSize(min=8, max=16), - s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10}), + s_q_s_kv = RandomSequenceLength(s_q_min=1, s_q_max=4096, s_kv_min=1, s_kv_max=4096, s_q_distribution={"s_q=1":0, "s_q=s_kv":5, "s_q=random":10, "s_q>s_kv":3}), d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=192, d_v_min=1, d_v_max=128, head_dim_distribution={"d_qk=d_v":5, "d_qk=random":1}, with_high_probability=[(64,64), (128,128), (192,128), (256,256)]), head_count=RandomHeadGenerator(min=1, max=8, head_group_options=(1, 4, 1)), data_type=RandomChoice({torch.float16 : 1, torch.bfloat16 : 2}), diff --git a/test/python/test_moe_grouped_matmul.py b/test/python/test_moe_grouped_matmul.py index 65931399a..95adead97 100644 --- a/test/python/test_moe_grouped_matmul.py +++ b/test/python/test_moe_grouped_matmul.py @@ -27,6 +27,59 @@ def get_compute_capability() -> int: return major * 10 + minor +# --------------------------------------------------------------------------- +# Numeric oracle + layout helpers +# +# Layouts (from the graph tensor definitions below): +# token [1, T, H] row-major -> token[t, h] = data[t*H + h] +# weight [E, H, N] stride [H*N, 1, H] -> weight[e,h,n] = data[e*H*N + h + n*H] +# i.e. expert block is column-major [H,N] +# output [1, T, N] row-major -> output[t, n] = data[t*N + n] +# Expert e owns token rows [offset[e], offset[e+1]) with offset[E] := T. +# This turns the previously execute-only harness into a checked one, so silent +# wrong-result / grouped-offset / empty-expert defects (NVBug 6192149-class, +# 5921085 scatter OOB) are actually caught. +# --------------------------------------------------------------------------- + + +def _expert_weight_HN(weight_data, e, H, N): + """Reconstruct expert e's [H, N] weight matrix from the column-major flat block.""" + block = weight_data[e * H * N : (e + 1) * H * N] + return block.view(N, H).t().float() # data[h + n*H] -> M[h, n] + + +def moe_fwd_reference(token_data, weight_data, offsets, E, T, H, N): + tok = token_data.view(T, H).float() + out = torch.zeros(T, N, dtype=torch.float32, device=token_data.device) + bounds = list(offsets) + [T] + for e in range(E): + lo, hi = bounds[e], bounds[e + 1] + if hi > lo: + out[lo:hi] = tok[lo:hi] @ _expert_weight_HN(weight_data, e, H, N) + return out # [T, N] + + +def moe_bwd_reference(doutput_data, token_data, offsets, E, T, H, N): + """dweight[e] = token[e-rows]^T @ doutput[e-rows], returned as [E, H, N] (column-major flat).""" + tok = token_data.view(T, H).float() + do = doutput_data.view(T, N).float() + bounds = list(offsets) + [T] + dw = torch.zeros(E, H, N, dtype=torch.float32, device=token_data.device) + for e in range(E): + lo, hi = bounds[e], bounds[e + 1] + if hi > lo: + dw[e] = tok[lo:hi].t() @ do[lo:hi] # [H,N] + return dw # [E, H, N] + + +def _moe_tol(contract_dim): + # bf16 IO, fp32 accumulate; scale with sqrt(contraction length) like the matmul fuzzer. + import math + + s = max(1.0, math.sqrt(contract_dim / 128.0)) + return 2e-2 * s, 2e-2 * s + + @pytest.mark.skipif( cudnn.backend_version() < 91800, reason="moe_grouped_matmul requires cuDNN >= 9.18.0", @@ -146,6 +199,12 @@ def test_bf16_moe_grouped_matmul_fwd(cudnn_handle): workspace, handle=cudnn_handle, ) + torch.cuda.synchronize() + + # Numeric oracle (was previously execute-only). + ref = moe_fwd_reference(token_data, weight_data, first_token_offset_values, num_experts, token_num, hidden_size, weight_size) + rtol, atol = _moe_tol(hidden_size) + torch.testing.assert_close(output_data.view(token_num, weight_size).float(), ref, rtol=rtol, atol=atol) @pytest.mark.skipif( @@ -278,3 +337,107 @@ def test_bf16_moe_grouped_matmul_bwd(cudnn_handle): workspace, handle=cudnn_handle, ) + torch.cuda.synchronize() + + # Numeric oracle (was previously execute-only). dweight is [E,H,N] column-major flat: + # dweight[e,h,n] = data[e*H*N + h + n*H] == data.view(E,N,H)[e].t() + ref = moe_bwd_reference(doutput_data, token_data, first_token_offset_values, num_experts, token_num, hidden_size, weight_size) + dw_actual = dweight_data.view(num_experts, weight_size, hidden_size).transpose(1, 2).float() + rtol, atol = _moe_tol(token_num) # contraction is over tokens + torch.testing.assert_close(dw_actual, ref, rtol=rtol, atol=atol) + + +def _rand_offsets(E, T, rng): + """Non-decreasing first-token offsets, offset[0]=0; duplicates => empty experts.""" + starts = sorted(rng.randint(0, T) for _ in range(E)) + starts[0] = 0 + return starts + + +@pytest.mark.skipif( + cudnn.backend_version() < 91800, + reason="moe_grouped_matmul requires cuDNN >= 9.18.0", +) +@pytest.mark.L0 +@pytest.mark.parametrize("seed", list(range(16))) +def test_bf16_moe_grouped_matmul_fwd_randomized(cudnn_handle, seed): + """Randomized experts/tokens/offsets (incl. empty experts) + numeric oracle. + + The original harness used one fixed shape and never checked the result. This + exercises grouped-offset / empty-expert / token-boundary handling against a + PyTorch per-expert reference. Caught class: 6192149 (grouped MoE numerics), + 5921085 (scatter OOB on uneven offsets). + """ + import random as _random + + rng = _random.Random(seed) + + num_experts = rng.choice([2, 4, 8, 17, 36, 64]) + token_num = rng.choice([16, 64, 200, 555, 2000]) + hidden_size = rng.choice([64, 128, 256, 520]) + weight_size = rng.choice([64, 128, 248, 256]) + # Force at least one empty expert in ~half the configs. + offsets = _rand_offsets(num_experts, token_num, rng) + if seed % 2 == 0 and num_experts >= 2: + offsets[1] = 0 # expert 0 empty + + torch.manual_seed(seed) + + graph = cudnn.pygraph( + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + handle=cudnn_handle, + ) + tensor_token = graph.tensor( + name="token", + dim=[1, token_num, hidden_size], + stride=[token_num * hidden_size, hidden_size, 1], + data_type=cudnn.data_type.BFLOAT16, + ) + tensor_weight = graph.tensor( + name="weight", + dim=[num_experts, hidden_size, weight_size], + stride=[hidden_size * weight_size, 1, hidden_size], + data_type=cudnn.data_type.BFLOAT16, + ) + tensor_first_token_offset = graph.tensor( + name="first_token_offset", + dim=[num_experts, 1, 1], + stride=[1, 1, 1], + data_type=cudnn.data_type.INT32, + ) + tensor_output = graph.moe_grouped_matmul( + tensor_token, + tensor_weight, + tensor_first_token_offset, + mode=cudnn.moe_grouped_matmul_mode.NONE, + compute_data_type=cudnn.data_type.FLOAT, + name="moe_grouped_matmul", + ) + tensor_output.set_data_type(cudnn.data_type.BFLOAT16).set_output(True) + + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A]) + try: + graph.check_support() + except Exception as e: + pytest.skip(f"unsupported config: {e}") + graph.build_plans() + + token_data = torch.randn(token_num * hidden_size, dtype=torch.bfloat16, device="cuda") + weight_data = torch.randn(num_experts * hidden_size * weight_size, dtype=torch.bfloat16, device="cuda") + first_token_offset_data = torch.tensor(offsets, dtype=torch.int32, device="cuda") + output_data = torch.empty(token_num * weight_size, dtype=torch.bfloat16, device="cuda") + workspace = torch.empty(graph.get_workspace_size(), dtype=torch.uint8, device="cuda") + + graph.execute( + {tensor_token: token_data, tensor_weight: weight_data, tensor_first_token_offset: first_token_offset_data, tensor_output: output_data}, + workspace, + handle=cudnn_handle, + ) + torch.cuda.synchronize() + + ref = moe_fwd_reference(token_data, weight_data, offsets, num_experts, token_num, hidden_size, weight_size) + rtol, atol = _moe_tol(hidden_size) + torch.testing.assert_close(output_data.view(token_num, weight_size).float(), ref, rtol=rtol, atol=atol) From 588246bf97de65adfe96a632939a8103dc074dc0 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Sun, 28 Jun 2026 23:55:40 -0700 Subject: [PATCH 2/6] test(low-precision-matmul): use canonical block-reduced nvfp4 descale shape The fp4 matmul test passed a full-size descale (1,M,K)=(1,128,64) instead of the canonical F8_128x4 block-reduced (1,M,ceil(K/block) rounded to 4)=(1,128,4) (and B symmetrically). It only "passed" because scales were all 1.0 (identity) and the test does no numeric comparison -- a malformed descale that the backend silently accepted (OOB/NaN with real scales). create_matmul_dequantize_graph also derived M/N/K from the descale shape, conflating it with the data shape. Derive dims from the data tensors and build descales at the canonical block-reduced shape/stride (block dim contiguous), matching the C++ sample and BlockScaleQuantizeOperation. Now passes the new dequant shape guard. Co-Authored-By: Claude Opus 4.8 (1M context) --- test/python/test_low_precision_matmul.py | 32 +++++++++++++++--------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/test/python/test_low_precision_matmul.py b/test/python/test_low_precision_matmul.py index 0410009a9..4733efb2b 100644 --- a/test/python/test_low_precision_matmul.py +++ b/test/python/test_low_precision_matmul.py @@ -277,12 +277,11 @@ def create_matmul_dequantize_graph(cudnn_handle, A, B, A_descale, B_descale, BLO with cudnn.graph(cudnn_handle) as (g, _): - batch_size, M, N, K = ( - A_descale.shape[0], - A_descale.shape[1], - B_descale.shape[1], - A_descale.shape[2], - ) + # Derive dims from the data tensors (float4_e2m1fn_x2 packs two fp4 per byte). + batch_size = A.shape[0] + M = A.shape[1] + K = A.shape[2] * 2 + N = B.shape[2] * 2 A_cudnn_tensor = g.tensor( name="tensor_a", @@ -298,18 +297,23 @@ def create_matmul_dequantize_graph(cudnn_handle, A, B, A_descale, B_descale, BLO data_type=convert_to_cudnn_type(B.dtype), ) + # F8_128x4 descale: block dim (K) reduced by BLOCK_SIZE & rounded to 4, contiguous; other dim to 128. + k_scale = ((K + BLOCK_SIZE - 1) // BLOCK_SIZE + 3) // 4 * 4 + m_pad = (M + 127) // 128 * 128 + n_pad = (N + 127) // 128 * 128 + A_descale_tensor = g.tensor( name="block_descale_a", - dim=A_descale.shape, - stride=(M * K, K, 1), + dim=(batch_size, m_pad, k_scale), + stride=(m_pad * k_scale, k_scale, 1), data_type=convert_to_cudnn_type(A_descale.dtype), reordering_type=cudnn.tensor_reordering.F8_128x4, ) B_descale_tensor = g.tensor( name="block_descale_b", - dim=B_descale.shape, - stride=(N * K, 1, K), + dim=(batch_size, k_scale, n_pad), + stride=(k_scale * n_pad, 1, k_scale), data_type=convert_to_cudnn_type(B_descale.dtype), reordering_type=cudnn.tensor_reordering.F8_128x4, ) @@ -413,8 +417,12 @@ def test_low_precision_fp4_matmul(cudnn_handle): A = _bfloat16_to_float4_e2m1fn_x2(A_ref) B = _bfloat16_to_float4_e2m1fn_x2(B_ref) - A_descale = torch.full((batch_size, M, K), 1.0, dtype=torch.float8_e4m3fn, device="cuda") - B_descale = torch.full((batch_size, K, N), 1.0, device="cuda", dtype=torch.float8_e4m3fn) + # Canonical F8_128x4 block-reduced descale shapes (K reduced by BLOCK_SIZE & rounded to 4; other dim to 128). + k_scale = ((K + BLOCK_SIZE - 1) // BLOCK_SIZE + 3) // 4 * 4 + m_pad = (M + 127) // 128 * 128 + n_pad = (N + 127) // 128 * 128 + A_descale = torch.full((batch_size, m_pad, k_scale), 1.0, dtype=torch.float8_e4m3fn, device="cuda") + B_descale = torch.full((batch_size, k_scale, n_pad), 1.0, device="cuda", dtype=torch.float8_e4m3fn) g, uids = create_matmul_dequantize_graph(cudnn_handle, A, B, A_descale, B_descale, BLOCK_SIZE) From cb921b4eb0e7f2c02d57769a3a6c3d932251066d Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Mon, 29 Jun 2026 00:59:38 -0700 Subject: [PATCH 3/6] sample(sdpa-mxfp8): align fwd SF_V to d-contiguous (stride[3]==1) convention The fwd mxfp8 sample was the lone outlier declaring SF_V s_scale-contiguous (stride[2]==1); SF_Q/SF_K, the bwd sample, and test_mhas_v2 all use d-contiguous (stride[3]==1). The kernel reads block-scale factors via the F8_128x4 swizzle, so the declared inner stride is not load-bearing (verified: flipping it with fixed data is bit-identical) -- consistency/clarity fix, behavior unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- samples/cpp/sdpa/mxfp8_fwd.cpp | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/samples/cpp/sdpa/mxfp8_fwd.cpp b/samples/cpp/sdpa/mxfp8_fwd.cpp index 11ae7be34..f986fe97d 100644 --- a/samples/cpp/sdpa/mxfp8_fwd.cpp +++ b/samples/cpp/sdpa/mxfp8_fwd.cpp @@ -126,12 +126,11 @@ TEST_CASE("sdpa_mxfp8_fprop", "[graph][sdpa][mxfp8][forward]") { .set_data_type(fe::DataType_t::FP8_E8M0) .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - // Block scale tensor for V (FP8_E8M0 with F8_128x4 reordering) - // V scales the s (sequence) dimension since BMM2 (S @ V) contracts on s_kv - // The contracting dimension (s_scale) must be contiguous, so use COL_MAJOR-like strides - auto SF_V_dims = std::vector({b, h, s_scale_padded, d_padded}); - auto SF_V_strides = - std::vector({h * s_scale_padded * d_padded, s_scale_padded * d_padded, 1, s_scale_padded}); + // Block scale tensor for V (FP8_E8M0, F8_128x4). V is block-scaled along s; declared + // d-contiguous (stride[3]==1) uniformly with SF_Q/SF_K -- the kernel reads SF via the + // F8_128x4 swizzle, so the declared inner stride is not load-bearing. + auto SF_V_dims = std::vector({b, h, s_scale_padded, d_padded}); + auto SF_V_strides = std::vector({h * s_scale_padded * d_padded, s_scale_padded * d_padded, d_padded, 1}); auto SF_V = mha_graph.tensor(fe::graph::Tensor_attributes() .set_name("SF_V") From edbfba8442fdeb6285da7fa658cb2e1f3dbfcec5 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Mon, 29 Jun 2026 03:59:31 -0700 Subject: [PATCH 4/6] test(matmul-fuzzer): add MATMUL_FUZZ_UNALIGNED for FORT-native widening-cast corner Opt-in: emit non-mult-of-4 K/N so the bits_per_access<32 LDG+STS smem-staging path is reachable, where a widening-cast (int8/fp8->fp16/fp32) operand over-runs the staging buffer (silent wrong-result on unaligned K, IMA on unaligned N). Co-Authored-By: Claude Opus 4.8 (1M context) --- test/python/test_matmul_fuzzer.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/test/python/test_matmul_fuzzer.py b/test/python/test_matmul_fuzzer.py index b7d7cdc30..539d2aa93 100644 --- a/test/python/test_matmul_fuzzer.py +++ b/test/python/test_matmul_fuzzer.py @@ -351,6 +351,17 @@ def generate(self) -> MatmulConfig: else: # k1 K = 1 + # Unaligned leading dims (K/N not a multiple of 4). random_dim() rounds up to mult-of-8, + # so the bits_per_access<32 corner is normally unreachable -- that's where the FORT-native + # engine takes the LDG+STS smem-staging path, which over-runs for a widening cast (int8/fp8 + # -> fp16/fp32): silent wrong-results when K is unaligned, IMA when N is. Opt-in (IMA aborts). + if os.environ.get("MATMUL_FUZZ_UNALIGNED", "0") == "1" and self.rng.random() < 0.15: + unaligned = self.rng.choice([1, 2, 3, 5, 6, 7, 9, 13, 17, 30]) # bpa<32 for 8-bit load + if self.rng.random() < 0.5: + K = unaligned + else: + N = unaligned + # Data types - ensure compatible combinations if self.rng.random() < 0.8: # 80% of tests use same dtype for A and B (more stable) From e58f36c50ee25a02b24f4332fa04fe0c3ef7f782 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Mon, 29 Jun 2026 03:59:32 -0700 Subject: [PATCH 5/6] test: remove broken L2 mxfp8 SDPA test (home-grown swizzle reference) create_scale_factor_tensor_for_sdpa builds the F8_128x4 scale swizzle by hand inconsistently with the kernel, feeding mis-ordered scales -> fails numerically across cuDNN versions (incl. official 9.23.1.3). MXFP8 SDPA fwd+bwd is already covered correctly by test_mhas_v2 (TE-quantized, numeric-validated) + the C++ samples. Co-Authored-By: Claude Opus 4.8 (1M context) --- test/python/test_sdpa_mxfp8.py | 469 --------------------------------- 1 file changed, 469 deletions(-) delete mode 100644 test/python/test_sdpa_mxfp8.py diff --git a/test/python/test_sdpa_mxfp8.py b/test/python/test_sdpa_mxfp8.py deleted file mode 100644 index e925323b0..000000000 --- a/test/python/test_sdpa_mxfp8.py +++ /dev/null @@ -1,469 +0,0 @@ -""" -Test for MXFP8 (Microscaling FP8) Scaled Dot Product Attention. - -MXFP8 uses block-wise scale factors with E8M0 data type and F8_128x4 reordering. -Requirements: -- cuDNN 9.21.0 or later -- Blackwell GPU architecture or newer -""" - -import cudnn -import pytest -import torch -import math -from looseversion import LooseVersion - -# Try to import CUTLASS for scale factor conversion -try: - import cutlass.cute as cute - import cutlass - from cutlass.cute.runtime import from_dlpack - - @cute.jit - def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( - sf_ref_tensor: cute.Tensor, - sf_mma_tensor: cute.Tensor, - ): - """Convert scale factor tensor from MKL layout to mma specification M(32x4xrest_m)xK(4xrest_k)xL layout""" - sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) - sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) - for i in cutlass.range(cute.size(sf_ref_tensor)): - mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) - sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] - - HAS_CUTLASS = True -except Exception: - HAS_CUTLASS = False - cute = None - cvt_sf_MKL_to_M32x4xrm_K4xrk_L = None - - -def ceil_div(a: int, b: int) -> int: - return (a + b - 1) // b - - -def compute_mxfp8_scale_dims(b, h, s, d, block_size=32): - """ - Compute scale tensor dimensions for MXFP8. - - For Q/K: scale the d (hidden) dimension - For V: scale the s (sequence) dimension (BMM2 contracts on s) - - F8_128x4 reordering requires: - - Sequence dimension padded to multiple of 128 - - Scale dimension padded to multiple of 4 - """ - d_scale = ceil_div(d, block_size) - s_scale = ceil_div(s, block_size) - - s_padded = ceil_div(s, 128) * 128 - d_scale_padded = ceil_div(d_scale, 4) * 4 - s_scale_padded = ceil_div(s_scale, 4) * 4 - d_padded = ceil_div(d, 4) * 4 - - return { - "s_padded": s_padded, - "d_scale": d_scale, - "d_scale_padded": d_scale_padded, - "s_scale": s_scale, - "s_scale_padded": s_scale_padded, - "d_padded": d_padded, - } - - -def create_sf_layout_tensor_for_sdpa(l, mn, nk, sf_vec_size): - """Create scale factor tensor with F8_128x4 layout for SDPA.""" - sf_k = ceil_div(nk, sf_vec_size) - - atom_m = (32, 4) - atom_k = 4 - mma_shape = ( - l, - ceil_div(mn, atom_m[0] * atom_m[1]), - ceil_div(sf_k, atom_k), - atom_m[0], - atom_m[1], - atom_k, - ) - - mma_permute_order = (3, 4, 1, 5, 2, 0) - cute_f32_torch_tensor_cpu = torch.zeros(mma_shape, dtype=torch.float32).permute(mma_permute_order) - - return cute_f32_torch_tensor_cpu, sf_k - - -def create_scale_factor_tensor_for_sdpa(l, mn, k, sf_vec_size, dtype): - """ - Create scale factor tensor for SDPA with F8_128x4 reordering. - - Args: - l: batch dimension (b * h for SDPA) - mn: non-contracting dimension (s for Q/K, d for V) - k: contracting dimension to be scaled (d for Q/K, s for V) - sf_vec_size: block size (32 for MXFP8) - dtype: output dtype (torch.float8_e8m0fnu) - - Returns: - ref_tensor: reference tensor for computation [mn, sf_k, l] -> broadcast to [mn, k, l] - cute_tensor: F8_128x4 reordered tensor for cuDNN - """ - if not HAS_CUTLASS: - pytest.skip("CUTLASS is not installed; skipping MXFP8 tests.") - - cute_f32_torch_tensor_cpu, sf_k = create_sf_layout_tensor_for_sdpa(l, mn, k, sf_vec_size) - ref_shape = (l, mn, sf_k) - ref_permute_order = (1, 2, 0) - - # Create reference scale factors (small positive values for stability) - ref_f32_torch_tensor_cpu = torch.empty(ref_shape, dtype=torch.float32).uniform_(0.5, 2.0).permute(ref_permute_order).to(torch.int8).to(torch.float32) - - # Convert ref f32 tensor to cute f32 tensor with F8_128x4 layout - cvt_sf_MKL_to_M32x4xrm_K4xrk_L( - from_dlpack(ref_f32_torch_tensor_cpu), - from_dlpack(cute_f32_torch_tensor_cpu), - ) - - # Expand scale factors to match the original k dimension - # ref shape: [mn, sf_k, l] -> expand to [mn, sf_k * sf_vec_size, l] -> trim to [mn, k, l] - ref_expanded = ( - ref_f32_torch_tensor_cpu.permute(2, 0, 1).unsqueeze(-1).expand(l, mn, sf_k, sf_vec_size).reshape(l, mn, sf_k * sf_vec_size).permute(*ref_permute_order) - ) - ref_expanded = ref_expanded[:, :k, :] - - # Convert to E8M0 dtype - cute_torch_tensor = cute_f32_torch_tensor_cpu.to(torch.float8_e8m0fnu).cuda() - - return ref_expanded.cuda(), cute_torch_tensor - - -def create_fp8_tensor(shape, dtype=torch.float8_e4m3fn): - """Create FP8 tensor with random values.""" - # Generate in float32, clamp to FP8 range, convert - tensor_f32 = torch.randn(shape, dtype=torch.float32, device="cuda").clamp(-2.0, 2.0) - tensor_fp8 = tensor_f32.to(dtype) - return tensor_fp8, tensor_f32 - - -def compute_sdpa_ref(q_f32, k_f32, v_f32, sf_q_ref, sf_k_ref, sf_v_ref, attn_scale, use_causal_mask=False): - """ - Compute reference SDPA with MXFP8 dequantization. - - Args: - q_f32: Query tensor [B, H, S_q, D] in float32 - k_f32: Key tensor [B, H, S_kv, D] in float32 - v_f32: Value tensor [B, H, S_kv, D] in float32 - sf_q_ref: Q scale factors [S_q, D, B*H] expanded to [S_q, D, B*H] - sf_k_ref: K scale factors [S_kv, D, B*H] expanded to [S_kv, D, B*H] - sf_v_ref: V scale factors [D, S_kv, B*H] expanded to [D, S_kv, B*H] - attn_scale: attention scale factor - use_causal_mask: whether to apply causal masking - - Returns: - o_ref: Output tensor [B, H, S_q, D] in float32 - """ - b, h, s_q, d = q_f32.shape - _, _, s_kv, _ = k_f32.shape - - # Reshape for batch processing: [B, H, S, D] -> [B*H, S, D] - q = q_f32.reshape(b * h, s_q, d) - k = k_f32.reshape(b * h, s_kv, d) - v = v_f32.reshape(b * h, s_kv, d) - - # sf_q_ref: [S_q, D, B*H] -> [B*H, S_q, D] - sf_q = sf_q_ref.permute(2, 0, 1) - # sf_k_ref: [S_kv, D, B*H] -> [B*H, S_kv, D] - sf_k = sf_k_ref.permute(2, 0, 1) - # sf_v_ref: [D, S_kv, B*H] -> [B*H, S_kv, D] - sf_v = sf_v_ref.permute(2, 1, 0) - - # Dequantize Q and K (scale factors apply to D dimension) - q_dq = q * sf_q[:, :s_q, :d] - k_dq = k * sf_k[:, :s_kv, :d] - - # BMM1: S = Q @ K^T, shape [B*H, S_q, S_kv] - s = torch.bmm(q_dq, k_dq.transpose(1, 2)) * attn_scale - - # Apply causal mask if requested - if use_causal_mask: - mask = torch.triu(torch.ones(s_q, s_kv, device=s.device, dtype=torch.bool), diagonal=1) - s = s.masked_fill(mask.unsqueeze(0), float("-inf")) - - # Softmax - p = torch.softmax(s, dim=-1) - - # Dequantize V (scale factors apply to S_kv dimension) - v_dq = v * sf_v[:, :s_kv, :d] - - # BMM2: O = P @ V, shape [B*H, S_q, D] - o = torch.bmm(p, v_dq) - - # Reshape back to [B, H, S_q, D] - o_ref = o.reshape(b, h, s_q, d) - - return o_ref - - -@pytest.mark.L2 -def test_sdpa_mxfp8_with_reference(request, cudnn_handle): - """Test MXFP8 SDPA with reference computation.""" - if request.config.option.dryrun: - pytest.skip("dry run mode") - - cudnn_version = LooseVersion(cudnn.backend_version_string()) - if cudnn_version < "9.21.0": - pytest.skip("MXFP8 SDPA requires cuDNN 9.21.0 or higher") - - if torch.cuda.get_device_capability()[0] < 10: - pytest.skip("MXFP8 SDPA requires Blackwell or higher") - - if not HAS_CUTLASS: - pytest.skip("CUTLASS is not installed; skipping MXFP8 tests.") - - # Problem dimensions - b = 2 # batch size - h = 2 # number of heads - s = 512 # sequence length (both s_q and s_kv for simplicity) - d = 128 # hidden head dim - block_size = 32 - - attn_scale = 1.0 / math.sqrt(d) - - # Compute scale tensor dimensions - dims = compute_mxfp8_scale_dims(b, h, s, d, block_size) - - # Create FP8 input tensors (BHSD layout) - q_fp8, q_f32 = create_fp8_tensor((b, h, s, d)) - k_fp8, k_f32 = create_fp8_tensor((b, h, s, d)) - v_fp8, v_f32 = create_fp8_tensor((b, h, s, d)) - - # Create scale factor tensors - # SF_Q and SF_K: scale the D dimension, shape [S_padded, D_scale, B*H] for ref - # The cute tensor will have F8_128x4 layout - sf_q_ref, sf_q_cute = create_scale_factor_tensor_for_sdpa(b * h, dims["s_padded"], d, block_size, torch.float8_e8m0fnu) - sf_k_ref, sf_k_cute = create_scale_factor_tensor_for_sdpa(b * h, dims["s_padded"], d, block_size, torch.float8_e8m0fnu) - - # SF_V: scale the S dimension, shape [D_padded, S_scale, B*H] for ref - # Note: for V, the contracting dim is S, so we scale S - sf_v_ref, sf_v_cute = create_scale_factor_tensor_for_sdpa(b * h, dims["d_padded"], s, block_size, torch.float8_e8m0fnu) - - # Compute reference output - o_ref = compute_sdpa_ref(q_f32, k_f32, v_f32, sf_q_ref, sf_k_ref, sf_v_ref, attn_scale, use_causal_mask=True) - - # Build cuDNN graph - graph = cudnn.pygraph(io_data_type=cudnn.data_type.FP8_E4M3, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) - - # Q, K, V tensors - BHSD layout, contiguous - qkv_dims = (b, h, s, d) - qkv_strides = (h * s * d, s * d, d, 1) - - q = graph.tensor(dim=qkv_dims, stride=qkv_strides, data_type=cudnn.data_type.FP8_E4M3) - k = graph.tensor(dim=qkv_dims, stride=qkv_strides, data_type=cudnn.data_type.FP8_E4M3) - v = graph.tensor(dim=qkv_dims, stride=qkv_strides, data_type=cudnn.data_type.FP8_E4M3) - - # Block scale tensors for Q, K (FP8_E8M0 with F8_128x4 reordering) - # Shape: [B, H, S_padded, D_scale_padded], d_scale contiguous (stride[3]=1) - sf_qk_dims = (b, h, dims["s_padded"], dims["d_scale_padded"]) - sf_qk_strides = (h * dims["s_padded"] * dims["d_scale_padded"], dims["s_padded"] * dims["d_scale_padded"], dims["d_scale_padded"], 1) - - sf_q_tensor = graph.tensor(dim=sf_qk_dims, stride=sf_qk_strides, data_type=cudnn.data_type.FP8_E8M0, reordering_type=cudnn.tensor_reordering.F8_128x4) - sf_k_tensor = graph.tensor(dim=sf_qk_dims, stride=sf_qk_strides, data_type=cudnn.data_type.FP8_E8M0, reordering_type=cudnn.tensor_reordering.F8_128x4) - - # Block scale tensor for V (FP8_E8M0 with F8_128x4 reordering) - # Shape: [B, H, S_scale_padded, D_padded], s_scale contiguous (stride[2]=1) - sf_v_dims = (b, h, dims["s_scale_padded"], dims["d_padded"]) - sf_v_strides = (h * dims["s_scale_padded"] * dims["d_padded"], dims["s_scale_padded"] * dims["d_padded"], 1, dims["s_scale_padded"]) - - sf_v_tensor = graph.tensor(dim=sf_v_dims, stride=sf_v_strides, data_type=cudnn.data_type.FP8_E8M0, reordering_type=cudnn.tensor_reordering.F8_128x4) - - # Call MXFP8 SDPA - o, stats, amax_o = graph.sdpa_mxfp8( - q=q, - k=k, - v=v, - descale_q=sf_q_tensor, - descale_k=sf_k_tensor, - descale_v=sf_v_tensor, - attn_scale=attn_scale, - use_causal_mask=True, - generate_stats=True, - ) - - # Set output tensor properties - o_strides = (h * s * d, s * d, d, 1) - o.set_output(True).set_dim(qkv_dims).set_stride(o_strides).set_data_type(cudnn.data_type.BFLOAT16) - amax_o.set_output(True).set_dim((1, 1, 1, 1)).set_stride((1, 1, 1, 1)).set_data_type(cudnn.data_type.FLOAT) - stats.set_output(True).set_data_type(cudnn.data_type.FLOAT) - - # Validate and build the graph - try: - graph.validate() - graph.build_operation_graph() - graph.create_execution_plans([cudnn.heur_mode.A]) - graph.check_support() - graph.build_plans() - except cudnn.cudnnGraphNotSupportedError as e: - pytest.skip(f"MXFP8 SDPA not supported: {e}") - except Exception as e: - pytest.fail(f"Error building MXFP8 SDPA graph: {e}") - - # Allocate output tensors - o_gpu = torch.empty(b, h, s, d, dtype=torch.bfloat16, device="cuda") - amax_o_gpu = torch.zeros(1, 1, 1, 1, dtype=torch.float32, device="cuda") - stats_gpu = torch.empty(b, h, s, 1, dtype=torch.float32, device="cuda") - - # Reshape cute tensors for cuDNN (they need to be contiguous in the right layout) - # The cute tensor is already in F8_128x4 layout, reshape to [B, H, S_padded, D_scale_padded] - sf_q_cudnn = sf_q_cute.reshape(b, h, dims["s_padded"], dims["d_scale_padded"]).contiguous() - sf_k_cudnn = sf_k_cute.reshape(b, h, dims["s_padded"], dims["d_scale_padded"]).contiguous() - # For V, reshape to [B, H, S_scale_padded, D_padded] with s_scale contiguous - sf_v_cudnn = sf_v_cute.reshape(b, h, dims["s_scale_padded"], dims["d_padded"]).permute(0, 1, 3, 2).contiguous().permute(0, 1, 3, 2) - - # Build variant pack - variant_pack = { - q: q_fp8, - k: k_fp8, - v: v_fp8, - sf_q_tensor: sf_q_cudnn, - sf_k_tensor: sf_k_cudnn, - sf_v_tensor: sf_v_cudnn, - o: o_gpu, - amax_o: amax_o_gpu, - stats: stats_gpu, - } - - # Execute - workspace = torch.empty(graph.get_workspace_size(), dtype=torch.uint8, device="cuda") - graph.execute(variant_pack, workspace, handle=cudnn_handle) - torch.cuda.synchronize() - - # Compare results - o_gpu_f32 = o_gpu.float() - atol, rtol = 0.1, 0.2 # FP8 has limited precision - - torch.testing.assert_close(o_gpu_f32, o_ref, atol=atol, rtol=rtol) - - -@pytest.mark.L2 -def test_sdpa_mxfp8_graph_build(request): - """Test that MXFP8 SDPA graph can be built on supported configurations.""" - if request.config.option.dryrun: - pytest.skip("dry run mode") - - cudnn_version = LooseVersion(cudnn.backend_version_string()) - if cudnn_version < "9.21.0": - pytest.skip("MXFP8 SDPA requires cuDNN 9.21.0 or higher") - - if torch.cuda.get_device_capability()[0] < 10: - pytest.skip("MXFP8 SDPA requires Blackwell or higher") - - # Problem dimensions - b = 2 # batch size - h = 2 # number of heads - s = 512 # sequence length - d = 128 # hidden head dim - - attn_scale = 0.123 - - # Compute scale tensor dimensions - dims = compute_mxfp8_scale_dims(b, h, s, d) - - # Create graph - graph = cudnn.pygraph(io_data_type=cudnn.data_type.FP8_E4M3, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) - - # Q, K, V tensors (FP8_E4M3) - BHSD layout with packed QKV (bs3hd interleaved) - qkv_dims = (b, h, s, d) - qkv_strides = (s * 3 * h * d, d, 3 * h * d, 1) # bs3hd - o_strides = (s * h * d, d, h * d, 1) # bshd - - q = graph.tensor(dim=qkv_dims, stride=qkv_strides, data_type=cudnn.data_type.FP8_E4M3) - k = graph.tensor(dim=qkv_dims, stride=qkv_strides, data_type=cudnn.data_type.FP8_E4M3) - v = graph.tensor(dim=qkv_dims, stride=qkv_strides, data_type=cudnn.data_type.FP8_E4M3) - - # Block scale tensors for Q, K (FP8_E8M0 with F8_128x4 reordering) - # Q and K scale the d (hidden) dimension since BMM1 contracts on d - sf_qk_dims = (b, h, dims["s_padded"], dims["d_scale_padded"]) - sf_qk_strides = (h * dims["s_padded"] * dims["d_scale_padded"], dims["s_padded"] * dims["d_scale_padded"], dims["d_scale_padded"], 1) - - sf_q = graph.tensor(dim=sf_qk_dims, stride=sf_qk_strides, data_type=cudnn.data_type.FP8_E8M0, reordering_type=cudnn.tensor_reordering.F8_128x4) - - sf_k = graph.tensor(dim=sf_qk_dims, stride=sf_qk_strides, data_type=cudnn.data_type.FP8_E8M0, reordering_type=cudnn.tensor_reordering.F8_128x4) - - # Block scale tensor for V (FP8_E8M0 with F8_128x4 reordering) - # V scales the s (sequence) dimension since BMM2 (S @ V) contracts on s_kv - # The contracting dimension (s_scale) must be contiguous, so use COL_MAJOR-like strides - sf_v_dims = (b, h, dims["s_scale_padded"], dims["d_padded"]) - sf_v_strides = (h * dims["s_scale_padded"] * dims["d_padded"], dims["s_scale_padded"] * dims["d_padded"], 1, dims["s_scale_padded"]) # s_scale contiguous - - sf_v = graph.tensor(dim=sf_v_dims, stride=sf_v_strides, data_type=cudnn.data_type.FP8_E8M0, reordering_type=cudnn.tensor_reordering.F8_128x4) - - # Call MXFP8 SDPA - o, stats, amax_o = graph.sdpa_mxfp8( - q=q, k=k, v=v, descale_q=sf_q, descale_k=sf_k, descale_v=sf_v, attn_scale=attn_scale, use_causal_mask=True, generate_stats=True - ) - - # Set output tensor properties - o.set_output(True).set_dim(qkv_dims).set_stride(o_strides).set_data_type(cudnn.data_type.BFLOAT16) - amax_o.set_output(True).set_dim((1, 1, 1, 1)).set_stride((1, 1, 1, 1)).set_data_type(cudnn.data_type.FLOAT) - stats.set_output(True).set_data_type(cudnn.data_type.FLOAT) - - # Validate and build the graph - try: - graph.validate() - graph.build_operation_graph() - graph.create_execution_plans([cudnn.heur_mode.A]) - graph.check_support() - graph.build_plans() - except cudnn.cudnnGraphNotSupportedError as e: - pytest.skip(f"MXFP8 SDPA not supported: {e}") - except Exception as e: - pytest.fail(f"Error building MXFP8 SDPA graph: {e}") - - -@pytest.mark.L2 -def test_sdpa_mxfp8_unsupported_version(request): - """Test that MXFP8 SDPA returns GRAPH_NOT_SUPPORTED on old cuDNN versions.""" - if request.config.option.dryrun: - pytest.skip("dry run mode") - - cudnn_version = LooseVersion(cudnn.backend_version_string()) - if cudnn_version >= "9.21.0": - pytest.skip("This test is for older cuDNN versions") - - # Problem dimensions - b, h, s, d = 2, 2, 512, 128 - attn_scale = 0.123 - - dims = compute_mxfp8_scale_dims(b, h, s, d) - - graph = cudnn.pygraph(io_data_type=cudnn.data_type.FP8_E4M3, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) - - qkv_dims = (b, h, s, d) - qkv_strides = (s * 3 * h * d, d, 3 * h * d, 1) - - q = graph.tensor(dim=qkv_dims, stride=qkv_strides, data_type=cudnn.data_type.FP8_E4M3) - k = graph.tensor(dim=qkv_dims, stride=qkv_strides, data_type=cudnn.data_type.FP8_E4M3) - v = graph.tensor(dim=qkv_dims, stride=qkv_strides, data_type=cudnn.data_type.FP8_E4M3) - - sf_qk_dims = (b, h, dims["s_padded"], dims["d_scale_padded"]) - sf_qk_strides = (h * dims["s_padded"] * dims["d_scale_padded"], dims["s_padded"] * dims["d_scale_padded"], dims["d_scale_padded"], 1) - - sf_q = graph.tensor(dim=sf_qk_dims, stride=sf_qk_strides, data_type=cudnn.data_type.FP8_E8M0, reordering_type=cudnn.tensor_reordering.F8_128x4) - sf_k = graph.tensor(dim=sf_qk_dims, stride=sf_qk_strides, data_type=cudnn.data_type.FP8_E8M0, reordering_type=cudnn.tensor_reordering.F8_128x4) - - sf_v_dims = (b, h, dims["s_scale_padded"], dims["d_padded"]) - sf_v_strides = (h * dims["s_scale_padded"] * dims["d_padded"], dims["s_scale_padded"] * dims["d_padded"], 1, dims["s_scale_padded"]) - - sf_v = graph.tensor(dim=sf_v_dims, stride=sf_v_strides, data_type=cudnn.data_type.FP8_E8M0, reordering_type=cudnn.tensor_reordering.F8_128x4) - - o, stats, amax_o = graph.sdpa_mxfp8( - q=q, k=k, v=v, descale_q=sf_q, descale_k=sf_k, descale_v=sf_v, attn_scale=attn_scale, use_causal_mask=True, generate_stats=True - ) - - o_strides = (s * h * d, d, h * d, 1) - o.set_output(True).set_dim(qkv_dims).set_stride(o_strides).set_data_type(cudnn.data_type.BFLOAT16) - amax_o.set_output(True).set_dim((1, 1, 1, 1)).set_stride((1, 1, 1, 1)).set_data_type(cudnn.data_type.FLOAT) - stats.set_output(True).set_data_type(cudnn.data_type.FLOAT) - - # On older cuDNN, validate() should return GRAPH_NOT_SUPPORTED - with pytest.raises(cudnn.cudnnGraphNotSupportedError): - graph.validate() From 8ba1a9a7b8e7940467e075431af8916fd7031292 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Mon, 29 Jun 2026 03:59:32 -0700 Subject: [PATCH 6/6] fix: correct mislabeled IS_VIRTUAL tensor descriptor error message The IS_VIRTUAL SetAttribute failure reused the BYTE_ALIGNMENT error string. Co-Authored-By: Claude Opus 4.8 (1M context) --- include/cudnn_frontend_Tensor.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/cudnn_frontend_Tensor.h b/include/cudnn_frontend_Tensor.h index 9f1f7c13a..b6d5123e5 100644 --- a/include/cudnn_frontend_Tensor.h +++ b/include/cudnn_frontend_Tensor.h @@ -447,7 +447,7 @@ class TensorBuilder_v8 { set_error_and_throw_exception( &m_tensor, status, - "CUDNN_BACKEND_TENSOR_DESCRIPTOR: SetAttribute CUDNN_ATTR_TENSOR_BYTE_ALIGNMENT Failed"); + "CUDNN_BACKEND_TENSOR_DESCRIPTOR: SetAttribute CUDNN_ATTR_TENSOR_IS_VIRTUAL Failed"); return std::move(m_tensor); } }