diff --git a/python/cudnn/AGENTS.md b/python/cudnn/AGENTS.md index babedb9ff..e9b5cde90 100644 --- a/python/cudnn/AGENTS.md +++ b/python/cudnn/AGENTS.md @@ -218,6 +218,30 @@ python/cudnn/gemm/ Shared helpers (schedulers, metadata utils, e.g. `gemm/cutedsl/grouped/moe_*.py`) stay internal to the family package — never exported through `cudnn`. +## CuTeDSL kernel bodies + +**Do not factor code out of a `@cute.kernel` body into a plain Python helper.** +The DSL AST-transforms only the decorated function's own source: `for` becomes +an `ir_loop`, `if` becomes an `scf` region. A helper called from the kernel is +not transformed, so the ops it emits can land outside the enclosing region. + +Hoisting an 11-line block that ran correctly inline into a +`write_clamped_kv_descs(...)` helper — called from inside +`if nvvm.elect_sync() and tidx < 32:` — turned 212 passing forward tests into +31 failures (`Error building ...`, traceback through `ir_loop` → +`scf_execute_dynamic`). Unrolling the helper's own loop did not help; the +helper *call* was the problem. Duplicating the block across flavors is the +correct trade here. Factor only host-side code, or code you can mark +`@cute.jit`. + +Related: inside a kernel body, `for x in (a, b)` over a Python tuple is +rewritten into a dynamic `ir_loop` and cannot iterate heterogeneous objects +(e.g. `GridConstant[TensorMap]`). Unroll it, or use `cutlass.range_constexpr`. + +**Detector.** These break at `compile()`, not at import — `python -c "import ..."` +and `pytest --collect-only` both stay green. After any refactor of a kernel +body, run that flavor's own tests. + ## The APIBase contract (`api_base.py`) Every OSS kernel API extends `APIBase` and implements: diff --git a/python/cudnn/sdpa/fwd/api_dsl.py b/python/cudnn/sdpa/fwd/api_dsl.py index 390ebbd37..11dd93ab9 100644 --- a/python/cudnn/sdpa/fwd/api_dsl.py +++ b/python/cudnn/sdpa/fwd/api_dsl.py @@ -1333,13 +1333,12 @@ def scratch_workspace_bytes(self) -> int: # head-major (H, head_stride)); without one it compiles with # has_lse=False and no LSE buffer exists at all. No slq/slk # copies either: the metadata is built DEVICE-side by the setup - # kernel (issue #552). o_desc: 16 int64 per - # sequence + 16 spare, the per-sequence O TMA descriptors the - # builder kernel fills. - # o_desc: 16 int64 per sequence + the dead-unit pad slot; the - # FP8/MXFP8 flavors carry two more slots for the packed-total- - # clamped K/V runtime descriptors (see the kernels' THD closures). - o_desc_slots = b + (3 if self._fp8 else 1) + # kernel (issue #552). + # o_desc: 16 int64 per sequence + the dead-unit pad slot + two + # slots for the packed-total-clamped K/V runtime descriptors the + # setup kernel writes (see the kernels' THD closures). Every THD + # flavor carries those two now, not just FP8/MXFP8 (issue #624). + o_desc_slots = b + 3 return ws_align((4 * b + 4) * 4) + ws_align(o_desc_slots * 16 * 8) + (0 if self.has_sink else ws_align(qh * 4)) if self._fp8 and self.split_kv == 1: return 0 # dense FP8/MXFP8: no per-execute scratch (dummies are cached one-time) @@ -1675,9 +1674,10 @@ def _thd_pack(self, q_buf, k_buf, v_buf, o_buf, sinks, seq_kv_lens, seq_q_lens, # before any consumer read, so stale bytes never survive — a fill # here is a wasted kernel launch on the execute hot path (Rule 1). with _torch_stream_context(current_stream, dev): - # +2 slots on the FP8/MXFP8 flavors: the packed-total-clamped K/V - # runtime descriptors the setup kernel writes after the pad slot. - o_desc_slots = b + (3 if self._fp8 else 1) + # +2 past the pad slot: the packed-total-clamped K/V runtime + # descriptors the setup kernel writes. Every THD flavor carries + # them now, not just FP8/MXFP8 (issue #624). + o_desc_slots = b + 3 o_desc = carver.take(o_desc_slots * 16, torch.int64) if carver is not None else torch.empty(o_desc_slots * 16, dtype=torch.int64, device=dev) # The PLAN-TIME envelope grid — dead units exit by kernel contract. # PERSISTENT THD grid: cap the launch at what the device can hold diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py index 64263f32f..78a70ff7f 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_f16_sm100.py @@ -601,6 +601,7 @@ def _kernel( n_q_supers=n_q_supers, n_qh=n_qh, n_batch=n_batch, + o_desc_words=o_desc_words, qh_per_kh=qh_per_kh, is_leader=is_leader, cta_in_pair=cta_in_pair, @@ -672,6 +673,7 @@ def _tmaldg_warp_group( n_q_supers, n_qh, n_batch, + o_desc_words, qh_per_kh, is_leader, cta_in_pair, @@ -693,8 +695,19 @@ def _tmaldg_warp_group( mb_q_reload = bars.mb_q_o_alias if cutlass.const_expr(IS_QO_ALIAS) else bars.mb_q_empty tma_q = GmemTileTma(tma_q_desc) - tma_k = GmemTileTma(tma_k_desc) - tma_v = GmemTileTma(tma_v_desc) + if cutlass.const_expr(CFG.THD_VARLEN): + # THD: K/V ride the setup kernel's packed-total-clamped runtime + # descriptors (o_desc_words slots n_batch+1 / n_batch+2), so the last + # sequence's tile-tail lands as exact zeros instead of reading the + # buffer's capacity tail (issue #624). Same closure shape as the dense + # GmemTileTma, so every load site below stays branch-free. + _k_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(1)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + _v_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(2)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + tma_k = lambda *coords: tma_slice_runtime_desc(_k_rt_ptr, *coords) # noqa: E731 + tma_v = lambda *coords: tma_slice_runtime_desc(_v_rt_ptr, *coords) # noqa: E731 + else: + tma_k = GmemTileTma(tma_k_desc) + tma_v = GmemTileTma(tma_v_desc) q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, @@ -2199,6 +2212,8 @@ def _tma_swz(byte_w: int): _build_thd_meta_o_descs_kernel( o_tensor, tma_o_desc, + tma_k_desc, + tma_v_desc, o_desc_words, seq_kv_lens_tensor, thd_q_lens_tensor, @@ -2412,9 +2427,10 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if CFG.SEQ_Q_LENS_PRESENT else None ) - # Per-batch O TMA-descriptor array (16 int64 = 128 B each) + 1 pad slot; - # dummy 1-elem when THD off (kernel never reads it). - _odesc_len = (b * _TENSOR_MAP_QWORDS + _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1 + # Per-batch O TMA-descriptor array (16 int64 = 128 B each) + 1 pad slot + # + 2 slots for the packed-total-clamped K/V runtime descriptors the setup + # kernel writes (issue #624); dummy 1-elem when THD off (never read). + _odesc_len = (b * _TENSOR_MAP_QWORDS + 3 * _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1 fake_o_desc = cute.runtime.make_fake_compact_tensor( cutlass.Int64, (_odesc_len,), diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py index 47125817d..db5f49d06 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d192_d128_f16_sm100.py @@ -634,6 +634,7 @@ def _kernel( n_q_supers=n_q_supers, n_qh=n_qh, n_batch=n_batch, + o_desc_words=o_desc_words, qh_per_kh=qh_per_kh, is_leader=is_leader, cta_in_pair=cta_in_pair, @@ -685,6 +686,7 @@ def _tmaldg_warp_group( n_q_supers, n_qh, n_batch, + o_desc_words, qh_per_kh, is_leader, cta_in_pair, @@ -706,8 +708,19 @@ def _tmaldg_warp_group( mb_q_reload = bars.mb_q_o_alias if cutlass.const_expr(IS_QO_ALIAS) else bars.mb_q_empty tma_q = GmemTileTma(tma_q_desc) - tma_k = GmemTileTma(tma_k_desc) - tma_v = GmemTileTma(tma_v_desc) + if cutlass.const_expr(CFG.THD_VARLEN): + # THD: K/V ride the setup kernel's packed-total-clamped runtime + # descriptors (o_desc_words slots n_batch+1 / n_batch+2), so the last + # sequence's tile-tail lands as exact zeros instead of reading the + # buffer's capacity tail (issue #624). Same closure shape as the dense + # GmemTileTma, so every load site below stays branch-free. + _k_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(1)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + _v_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(2)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + tma_k = lambda *coords: tma_slice_runtime_desc(_k_rt_ptr, *coords) # noqa: E731 + tma_v = lambda *coords: tma_slice_runtime_desc(_v_rt_ptr, *coords) # noqa: E731 + else: + tma_k = GmemTileTma(tma_k_desc) + tma_v = GmemTileTma(tma_v_desc) q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, @@ -2266,6 +2279,8 @@ def _tma_swz(byte_w: int): _build_thd_meta_o_descs_kernel( o_tensor, tma_o_desc, + tma_k_desc, + tma_v_desc, o_desc_words, seq_kv_lens_tensor, thd_q_lens_tensor, @@ -2471,7 +2486,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): ) # Per-batch O TMA-descriptor array (16 int64 = 128 B each) + 1 pad slot; # dummy 1-elem when THD off (kernel never reads it). - _odesc_len = (b * _TENSOR_MAP_QWORDS + _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1 + # +2 slots beyond the pad: the packed-total-clamped K/V runtime + # descriptors the setup kernel writes (issue #624). + _odesc_len = (b * _TENSOR_MAP_QWORDS + 3 * _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1 fake_o_desc = cute.runtime.make_fake_compact_tensor( cutlass.Int64, (_odesc_len,), diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py index b6b491c52..dd02511b8 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d256_f16_sm100.py @@ -447,6 +447,7 @@ def _kernel( n_q_supers=n_q_supers, n_qh=n_qh, n_batch=n_batch, + o_desc_words=o_desc_words, qh_per_kh=qh_per_kh, is_leader=is_leader, cta_in_pair=cta_in_pair, @@ -493,6 +494,7 @@ def _tmaldg_warp_group( n_q_supers, n_qh, n_batch, + o_desc_words, qh_per_kh, is_leader, cta_in_pair, @@ -502,8 +504,19 @@ def _tmaldg_warp_group( kv_state = PipelineState.start(phase=1) tma_q = GmemTileTma(tma_q_desc) - tma_k = GmemTileTma(tma_k_desc) - tma_v = GmemTileTma(tma_v_desc) + if cutlass.const_expr(CFG.THD_VARLEN): + # THD: K/V ride the setup kernel's packed-total-clamped runtime + # descriptors (o_desc_words slots n_batch+1 / n_batch+2), so the last + # sequence's tile-tail lands as exact zeros instead of reading the + # buffer's capacity tail (issue #624). Same closure shape as the dense + # GmemTileTma, so every load site below stays branch-free. + _k_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(1)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + _v_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(2)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + tma_k = lambda *coords: tma_slice_runtime_desc(_k_rt_ptr, *coords) # noqa: E731 + tma_v = lambda *coords: tma_slice_runtime_desc(_v_rt_ptr, *coords) # noqa: E731 + else: + tma_k = GmemTileTma(tma_k_desc) + tma_v = GmemTileTma(tma_v_desc) q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, @@ -1770,6 +1783,8 @@ def _tma_swz(byte_w: int): _build_thd_meta_o_descs_kernel( o_tensor, tma_o_desc, + tma_k_desc, + tma_v_desc, o_desc_words, seq_kv_lens_tensor, thd_q_lens_tensor, @@ -1960,7 +1975,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if CFG.SEQ_Q_LENS_PRESENT else None ) - _odesc_len = (b * _TENSOR_MAP_QWORDS + _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1 + # +2 slots beyond the pad: the packed-total-clamped K/V runtime + # descriptors the setup kernel writes (issue #624). + _odesc_len = (b * _TENSOR_MAP_QWORDS + 3 * _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1 fake_o_desc = cute.runtime.make_fake_compact_tensor( cutlass.Int64, (_odesc_len,), diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py index 15bcf30f1..1df375e06 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d512_f16_sm100.py @@ -625,6 +625,7 @@ def _kernel( n_q_supers=n_q_supers, n_qh=n_qh, n_batch=n_batch, + o_desc_words=o_desc_words, qh_per_kh=qh_per_kh, is_leader=is_leader, cta_in_pair=cta_in_pair, @@ -1644,6 +1645,7 @@ def _tmaldg_warp_group( n_q_supers, n_qh, n_batch, + o_desc_words, qh_per_kh, is_leader, cta_in_pair, @@ -1655,8 +1657,19 @@ def _tmaldg_warp_group( o_empty_for_v_state = PipelineState.start(phase=1) tma_q = GmemTileTma(tma_q_desc) - tma_k = GmemTileTma(tma_k_desc) - tma_v = GmemTileTma(tma_v_desc) + if cutlass.const_expr(CFG.THD_VARLEN): + # THD: K/V ride the setup kernel's packed-total-clamped runtime + # descriptors (o_desc_words slots n_batch+1 / n_batch+2), so the last + # sequence's tile-tail lands as exact zeros instead of reading the + # buffer's capacity tail (issue #624). Same closure shape as the dense + # GmemTileTma, so every load site below stays branch-free. + _k_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(1)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + _v_rt_ptr = (o_desc_words.iterator.raw_ptr() + (n_batch + cutlass.Int32(2)) * cutlass.Int32(TENSOR_MAP_QWORDS)).tospace(cutlass.AddressSpace.generic) + tma_k = lambda *coords: tma_slice_runtime_desc(_k_rt_ptr, *coords) # noqa: E731 + tma_v = lambda *coords: tma_slice_runtime_desc(_v_rt_ptr, *coords) # noqa: E731 + else: + tma_k = GmemTileTma(tma_k_desc) + tma_v = GmemTileTma(tma_v_desc) q_super_idx, head_idx, batch_idx, split_idx = _decode_initial_split( sched.bidx_init, @@ -1975,6 +1988,8 @@ def _tma_swz(byte_w: int): _build_thd_meta_o_descs_kernel( o_tensor, tma_o_desc, + tma_k_desc, + tma_v_desc, o_desc_words, seq_kv_lens_tensor, thd_q_lens_tensor, @@ -2161,7 +2176,9 @@ def _fake_bshd(shape, stride, dtype=STORAGE_DTYPE, bpe=CFG.BPE): if CFG.SEQ_Q_LENS_PRESENT else None ) - _odesc_len = (b * _TENSOR_MAP_QWORDS + _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1 + # +2 slots beyond the pad: the packed-total-clamped K/V runtime + # descriptors the setup kernel writes (issue #624). + _odesc_len = (b * _TENSOR_MAP_QWORDS + 3 * _TENSOR_MAP_QWORDS) if CFG.THD_VARLEN else 1 fake_o_desc = cute.runtime.make_fake_compact_tensor( cutlass.Int64, (_odesc_len,), diff --git a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py index 556c58369..79651a888 100644 --- a/python/cudnn/sdpa/fwd/kernels/thd_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/thd_sm100.py @@ -231,8 +231,12 @@ def build_thd_meta_o_kv_descs_kernel( n_batch: cutlass.Int32, o_row_stride: cutlass.Int32, ) -> None: - """``build_thd_meta_o_descs_kernel`` + packed-total-clamped K/V descriptors - (the FP8/MXFP8 THD flavors). + """THD setup for the FP8/MXFP8 flavors: metadata, per-batch O descriptors + and the packed-total-clamped K/V descriptors. + + Same body as ``build_thd_meta_o_descs_kernel`` minus the persistent + scheduler's live-unit total and claim counter, which these flavors do not + launch with. Both kernels clamp K/V (issue #624). The K/V loads tile in TILE_N rows, so the LAST sequence's tile steps past the packed KV total into the buffer's capacity tail — caller-owned bytes @@ -305,6 +309,8 @@ def build_thd_meta_o_kv_descs_kernel( def build_thd_meta_o_descs_kernel( o_tensor: cute.Tensor, base_o_desc: cutlass.GridConstant[tmap.TensorMap], + base_k_desc: cutlass.GridConstant[tmap.TensorMap], + base_v_desc: cutlass.GridConstant[tmap.TensorMap], o_desc_words: cute.Tensor, meta_t: cute.Tensor, q_lens_t: cute.Tensor, @@ -316,7 +322,9 @@ def build_thd_meta_o_descs_kernel( cga_tile_m: cutlass.Int32, n_clusters: cutlass.Int32, ) -> None: - """Per-execute THD setup, one elected thread (issue #552, D2H removal): + """Per-execute THD setup for the f16/bf16 flavors, one elected thread — + ``build_thd_meta_o_kv_descs_kernel`` plus the persistent scheduler's + live-unit total and claim counter (issue #552, D2H removal): build the [seq_kv_lens(B) | cu_seqlens_q(B+1) | cu_seqlens_k(B+1)] metadata buffer DEVICE-side from the caller's length tensors — ``(B,)`` per-batch lengths (serial cumsum; B is small) or the ``(B+1,)`` cu prefix-sum form @@ -364,6 +372,27 @@ def build_thd_meta_o_descs_kernel( new_value=s_i, ord=2, ) + # Packed-total-clamped K/V runtime descriptors (issue #624). K/V load + # in TILE_N rows, so the LAST sequence's tile steps past the packed KV + # total into the buffer's capacity tail — caller-owned bytes that may + # never have been written. Masked S columns are NaN-safe (the mask is + # a select), but BMM2's P·V is not: 0 · NaN == NaN wipes every valid + # row of the tile. Patching the seq extent (GLOBAL_DIM ord=2) to + # cu_k[B] makes those rows TMA-OOB, so they land as EXACT ZEROS + # without touching memory — no fill kernel, and nothing written into + # the caller's buffer. Mirrors build_thd_meta_o_kv_descs_kernel, which + # the FP8/MXFP8 flavors have used for this since they were written. + t_kv = cutlass.Int32(meta[cutlass.Int32(3) * n_batch + cutlass.Int32(1)]) # cu_k[B] + k_dptr = desc_base + (n_batch + cutlass.Int32(1)) * cutlass.Int32(TENSOR_MAP_QWORDS) + k_src = Pointer(base_k_desc.get_ptr(), dtype=cutlass.Int64) + for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS): + (k_dptr + i).store((k_src + i).load()) + nvvm.tensormap_replace(nvvm.TensormapField.GLOBAL_DIM, k_dptr, new_value=t_kv, ord=2) + v_dptr = desc_base + (n_batch + cutlass.Int32(2)) * cutlass.Int32(TENSOR_MAP_QWORDS) + v_src = Pointer(base_v_desc.get_ptr(), dtype=cutlass.Int64) + for i in cutlass.range_constexpr(TENSOR_MAP_QWORDS): + (v_dptr + i).store((v_src + i).load()) + nvvm.tensormap_replace(nvvm.TensormapField.GLOBAL_DIM, v_dptr, new_value=t_kv, ord=2) nvvm.fence_proxy_release( nvvm.MemScope.GPU, from_proxy=nvvm.Proxy.GENERIC, diff --git a/test/AGENTS.md b/test/AGENTS.md index ff216bfca..c3dfd9d16 100644 --- a/test/AGENTS.md +++ b/test/AGENTS.md @@ -49,3 +49,43 @@ pytest fe_api/gemm/ # OSS kernel tests - **A regression test must be seen RED.** Before trusting one, run it against the unfixed code — restore the old line, confirm it fails, restore the fix. `test_dsl_sm100_thd_interleaved_kv_views` and `test_varlen_backward_does_not_sync` were both checked this way, and both were genuinely red beforehand; a test written for a bug and never seen to fail is asserting an unknown. - **Seed before you allocate.** `torch.manual_seed()` after constructing the inputs seeds nothing that matters. Two runs meant to be compared then differ by data, and the assertion fails (or worse, passes) for a reason unrelated to what is under test — if two runs must be comparable, build the inputs once and reuse them. - **When you remove a fallback, invert its counter assertion — do not delete it.** Tests that asserted `calls["bwd_cpp"]` incremented had to become "`calls["bwd"]` increments **and** `bwd_cpp` does not", so a silent regression to the old path fails the suite instead of passing it. + +### Confirm you are testing the code you edited + +`pip install -e .` does **not** put the package on `sys.path`. It installs a +`sys.meta_path` finder (`__editable___nvidia_cudnn_frontend_*_finder.py`) whose +`MAPPING` hard-codes an absolute path to the checkout it was installed from. +Meta-path finders run *before* `sys.path`, so **`PYTHONPATH` cannot shadow it** — +if you edit a different clone or a git worktree, your changes are silently not +under test. Symptoms are indistinguishable from a real result: a probe that +should change the output leaves it bit-identical, and edits appear to do nothing. + +Check first, every time you work outside the installed checkout: + +```bash +python -c "import cudnn; print(cudnn.__file__)" # must be YOUR tree +``` + +`conftest.py` prints the same path in its banner (`cuDNN Frontend Path:`) — read +it rather than assuming. To point the editable install at another tree for one +run, patch the finder's `MAPPING` from a `sitecustomize.py` on `PYTHONPATH` +(`site` imports it after processing `.pth` files, so the finder already exists): + +```python +# sitecustomize.py -- the finder module name embeds the installed version, so +# discover it rather than hard-coding it (it changes when __version__ bumps). +import importlib, pkgutil + +name = next( + m.name + for m in pkgutil.iter_modules() + if m.name.startswith("__editable___nvidia_cudnn_frontend_") and m.name.endswith("_finder") +) +importlib.import_module(name).MAPPING["cudnn"] = "/path/to/your/worktree/python/cudnn" +``` + +The same trap hides *inside* a run: `python/cudnn/frost/template_loader.py` +loads kernel templates by absolute path via `spec_from_file_location`, so the +template that serves a config may come from elsewhere too. To find out which +template a test actually compiles, log `path` at the top of `load_template` — +do not infer it from `_pick_flavor` by reading the source. diff --git a/test/python/sdpa/fp16.py b/test/python/sdpa/fp16.py index 2d2db5e07..e53be63b2 100644 --- a/test/python/sdpa/fp16.py +++ b/test/python/sdpa/fp16.py @@ -169,6 +169,18 @@ def allocate_tensors(cfg, rng_data_gen, perf=False): allocs[TensorUid.dK] = alloc_tensor((max_t_kv, cfg.h_k, cfg.d_qk), cfg.data_type, strides=k_strides) allocs[TensorUid.dV] = alloc_tensor((max_t_kv, cfg.h_v, cfg.d_v), cfg.data_type, strides=v_strides) allocs[TensorUid.dO] = alloc_tensor((max_t_q, cfg.h_q, cfg.d_v), cfg.data_type, strides=o_strides, rng=rng_data_gen, mean=0.0, std=0.1, sparse_int=si) + # NaN-poison the capacity tail (tokens past the last ragged offset). + # No engine may ever read those rows; finite random data there hides + # capacity-vs-live-token binding bugs because the padding mask turns + # them into exact zeros (0 x finite = 0), while real recycled device + # memory holds NaN bit patterns (0 x NaN = NaN) — GitHub issue #624. + # This makes f16/bf16 consistent with the fp8 harness, whose + # convert_uniform_to_packed always NaN-fills the tail. + total_t_q, total_t_kv = sum(cfg.seq_len_q), sum(cfg.seq_len_kv) + for uid, total in ((TensorUid.q, total_t_q), (TensorUid.k, total_t_kv), (TensorUid.v, total_t_kv)): + allocs[uid][0][total:] = float("nan") + if cfg.is_train: + allocs[TensorUid.dO][0][total_t_q:] = float("nan") else: allocs[TensorUid.q] = alloc_tensor(cfg.shape_q, cfg.data_type, strides=cfg.stride_q, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si) allocs[TensorUid.k] = alloc_tensor(cfg.shape_k, cfg.data_type, strides=cfg.stride_k, rng=rng_data_gen, mean=-0.5, std=1.0, sparse_int=si) @@ -733,8 +745,8 @@ def _to_seq_len_ref(seq_len_gpu, cu_seq_len_gpu): if cfg.is_ragged and cfg.is_train: dO_ref = convert_packed_to_uniform(dO_ref, seq_len_q_ref, cfg.s_q) - max_t_q = max(64, ((seq_len_q_ref.sum().item() + 63) // 64) * 64) if cfg.is_ragged else None - max_t_kv = max(64, ((seq_len_kv_ref.sum().item() + 63) // 64) * 64) if cfg.is_ragged else None + max_t_q = packed_token_capacity(seq_len_q_ref.tolist()) if cfg.is_ragged else None + max_t_kv = packed_token_capacity(seq_len_kv_ref.tolist()) if cfg.is_ragged else None attn_scale = 0.125 diff --git a/test/python/sdpa/fp8.py b/test/python/sdpa/fp8.py index bcc16584d..c051dcc09 100644 --- a/test/python/sdpa/fp8.py +++ b/test/python/sdpa/fp8.py @@ -27,6 +27,7 @@ profile_execution, note_frost_routing, ) +from .random_config import packed_token_capacity # fmt: off @@ -359,8 +360,11 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle): if is_ragged: seq_len_q_gpu = torch.tensor(seq_len_q_list, dtype=torch.int32, device="cuda").view(-1) seq_len_kv_gpu = torch.tensor(seq_len_kv_list, dtype=torch.int32, device="cuda").view(-1) - max_t_q = max(64, ((seq_len_q_gpu.sum().item() + 63) // 64) * 64) - max_t_kv = max(64, ((seq_len_kv_gpu.sum().item() + 63) // 64) * 64) + # Guaranteed capacity tail (> total tokens); convert_uniform_to_packed + # NaN-fills it, so engines that read past the last ragged offset fail + # deterministically (GitHub #624). + max_t_q = packed_token_capacity(seq_len_q_list) + max_t_kv = packed_token_capacity(seq_len_kv_list) # With the ragged offset multiplier, offsets are stored in coarser units # (divided by the per-tensor multiplier; always divides evenly) and the diff --git a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py index 8ad146950..2db5f5b91 100644 --- a/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py +++ b/test/python/sdpa/frost/test_sdpa_fwd_dsl_sm100.py @@ -958,6 +958,108 @@ def _dense_buf(packed): torch.testing.assert_close(o_out, o_ref, atol=5e-2, rtol=3e-2) +# Issue #624: a THD caller binds K/V at BUFFER CAPACITY, not at the packed +# total, and the descriptor extent is the host-derived capacity. Rows in +# [total, capacity) are therefore inside the extent and get loaded. They are +# masked, so P == 0 -- but 0 * NaN == NaN, so an UNINITIALIZED capacity tail +# wipes whole tiles of O. The setup kernel patches the K/V descriptor extent to +# the packed total cu_k[B] device-side so those rows are TMA-OOB and land as +# exact zeros instead. Total 397 is deliberately not a multiple of TILE_N, so +# the last sequence's tail tile steps past it. +@pytest.mark.L0 +# (d_qk, d_v) rather than a single d: d192/d128 is its own kernel flavor, and +# a single-d parametrization cannot reach it. +@pytest.mark.parametrize("d_qk,d_v", [(128, 128), (192, 128), (256, 256), (512, 512)], ids=["llama", "dsv3", "qwen", "dsv4"]) +@pytest.mark.parametrize("dtype", _DTYPES, ids=_DTYPE_IDS) +@torch_fork_set_rng(seed=0) +def test_dsl_sm100_thd_nan_capacity_tail(dtype, d_qk, d_v): + """THD with a NaN-poisoned capacity tail: O must be finite and correct.""" + _require_dsl() + import cudnn + + dev = "cuda" + H = 8 + seq_lens = [200, 150, 47] + B = len(seq_lens) + S_max = max(seq_lens) + T = sum(seq_lens) + cu = [0] + for s_i in seq_lens: + cu.append(cu[-1] + s_i) + scale = 1.0 / math.sqrt(d_qk) + + q_pk = torch.randn(T, H, d_qk, device=dev, dtype=dtype) + k_pk = torch.randn(T, H, d_qk, device=dev, dtype=dtype) + v_pk = torch.randn(T, H, d_v, device=dev, dtype=dtype) + + def _stride(dd): + return (S_max * H * dd, dd, H * dd, 1) + + def _poisoned_buf(packed, dd): + """Capacity-sized storage; only [0, T) is live, the tail is NaN. + + Returns the strided view only -- it keeps the storage alive, so there + is no need to bind the flat tensor as well.""" + stor = torch.full((B * S_max * H * dd,), float("nan"), device=dev, dtype=dtype) + stor[: T * H * dd] = packed.reshape(-1) + return stor.as_strided((B, H, S_max, dd), _stride(dd)) + + q_gpu = _poisoned_buf(q_pk, d_qk) + k_gpu = _poisoned_buf(k_pk, d_qk) + v_gpu = _poisoned_buf(v_pk, d_v) + o_stor = torch.zeros(B * S_max * H * d_v, device=dev, dtype=dtype) + o_gpu = o_stor.as_strided((B, H, S_max, d_v), _stride(d_v)) + + slq = torch.tensor(seq_lens, dtype=torch.int32, device=dev).view(B, 1, 1, 1) + slk = slq.clone() + cu_t = torch.tensor(cu, dtype=torch.int64, device=dev) + ro_qk = (cu_t * H * d_qk).view(B + 1, 1, 1, 1) + ro_v = (cu_t * H * d_v).view(B + 1, 1, 1, 1) + + io = cudnn.data_type.HALF if dtype == torch.float16 else cudnn.data_type.BFLOAT16 + g = cudnn.pygraph(io_data_type=io, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) + tq = g.tensor(dim=[B, H, S_max, d_qk], stride=list(_stride(d_qk)), data_type=io, name="q") + tk = g.tensor(dim=[B, H, S_max, d_qk], stride=list(_stride(d_qk)), data_type=io, name="k") + tv = g.tensor(dim=[B, H, S_max, d_v], stride=list(_stride(d_v)), data_type=io, name="v") + sq = g.tensor_like(slq) + skv = g.tensor_like(slk) + qro = g.tensor_like(ro_qk) + kro = g.tensor_like(ro_qk) + vro = g.tensor_like(ro_v) + oro = g.tensor_like(ro_v) + tq.set_ragged_offset(qro) + tk.set_ragged_offset(kro) + tv.set_ragged_offset(vro) + o, _ = g.sdpa( + name="sdpa", q=tq, k=tk, v=tv, generate_stats=False, attn_scale=scale, use_causal_mask=True, use_padding_mask=True, seq_len_q=sq, seq_len_kv=skv + ) + o.set_output(True).set_dim([B, H, S_max, d_v]).set_stride(list(_stride(d_v))) + o.set_ragged_offset(oro) + + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A]) + _select_engine(g, engine_name()) + g.check_support() + g.build_plans() + vp = {tq: q_gpu, tk: k_gpu, tv: v_gpu, o: o_gpu, sq: slq, skv: slk, qro: ro_qk, kro: ro_qk, vro: ro_v, oro: ro_v} + g.execute(vp, torch.empty(max(g.get_workspace_size(), 1), device=dev, dtype=torch.uint8)) + torch.cuda.synchronize() + + o_out = o_stor[: T * H * d_v].reshape(T, H, d_v) + assert not torch.isnan(o_out).any(), f"{int(torch.isnan(o_out).sum())} NaNs in O -- the capacity tail reached BMM2" + + o_ref = torch.zeros(T, H, d_v, device=dev, dtype=dtype) + for b in range(B): + lo, hi = cu[b], cu[b + 1] + qb = q_pk[lo:hi].permute(1, 0, 2).unsqueeze(0) + kb = k_pk[lo:hi].permute(1, 0, 2).unsqueeze(0) + vb = v_pk[lo:hi].permute(1, 0, 2).unsqueeze(0) + ob = _ref_sdpa_full(qb, kb, vb, scale=scale, is_causal=True) + o_ref[lo:hi] = ob.squeeze(0).permute(1, 0, 2) + torch.testing.assert_close(o_out, o_ref, atol=5e-2, rtol=3e-2) + + @pytest.mark.L0 @pytest.mark.parametrize("d", _FLAVORS, ids=_FLAVOR_IDS) @pytest.mark.parametrize("dtype", _DTYPES, ids=_DTYPE_IDS) @@ -1706,9 +1808,11 @@ def test_dsl_sm100_thd_declared_total_bounds_capacity_tail(): FINITE. A caller that over-allocates and leaves the tail unwritten therefore poisons whole tiles through ``0 * NaN`` (issue #624). - Declaring the total clamps the TMA extent to the packed total, putting that - tail out of reach. Same graph, same data, only the tail fill differs: with - the declaration a NaN tail must be inert.""" + Two mechanisms now put that tail out of reach: declaring the total clamps + the host-derived extent, and the THD setup kernel clamps the K/V extent to + ``cu_k[B]`` device-side regardless. Same graph, same data, only the tail + fill differs -- a NaN tail must be inert either way, and the two paths must + agree bit for bit.""" _require_dsl() import cudnn @@ -1782,13 +1886,22 @@ def _run(tail_value, declare_total): ) assert torch.equal(declared_nan, declared_zero), "O must not depend on the capacity tail once the total is declared" - # Control: the tail is genuinely reachable without the declaration, so the - # assertions above are testing the clamp rather than a benign shape. + # The undeclared path is NaN-safe too, and deliberately so: the THD setup + # kernel now patches the K/V descriptor extent to cu_k[B] device-side on + # every flavor, which bounds the tail without any host declaration. Before + # that clamp this same call poisoned ~50% of O, and this assertion was + # inverted -- it asserted the NaN to prove the tail was reachable. undeclared_nan = _run(float("nan"), declare_total=False) - assert torch.isnan(undeclared_nan).any(), ( - "expected the undeclared path to read the capacity tail (issue #624); if this no longer " - "holds the extent is exact by other means and this test needs rethinking" + undeclared_zero = _run(0.0, declare_total=False) + assert not torch.isnan(undeclared_nan).any(), ( + f"the device-side K/V extent clamp must bound the capacity tail with or without a declared " + f"total: {int(torch.isnan(undeclared_nan).sum())} NaNs in O" ) + assert torch.equal(undeclared_nan, undeclared_zero), "O must not depend on the capacity tail" + # Declaring the total is now a host-side tightening (capacity, grid and + # workspace sizing) rather than the only thing standing between the kernel + # and the tail -- so both paths must agree bit for bit. + assert torch.equal(declared_nan, undeclared_nan), "declaring the packed total must not change the result" @pytest.mark.L0 diff --git a/test/python/sdpa/mxfp8.py b/test/python/sdpa/mxfp8.py index 5c44dd06b..3ad0c4871 100644 --- a/test/python/sdpa/mxfp8.py +++ b/test/python/sdpa/mxfp8.py @@ -31,6 +31,14 @@ quantize_to_mxfp8, ) # noqa: F401 (re-exported; mxfp8_ref imports it from here) +# NOTE: this harness is dense-full only. The sdpa_mxfp8 python API exposes no +# seq_len/padding arguments (and the MXFP8 engines defer THD/varlen), so +# cfg.is_padding / cfg.seq_len_q / cfg.seq_len_kv are intentionally ignored +# here. When padding/THD support lands in the API, wire it through the shared +# packed_token_capacity / convert_uniform_to_packed helpers so the ragged +# capacity tails come NaN-poisoned (see GitHub issue #624 for why that +# poisoning is load-bearing). + # fmt: off def ceil_div(a: int, b: int) -> int: diff --git a/test/python/sdpa/random_config.py b/test/python/sdpa/random_config.py index d4f638fb3..65659cff1 100644 --- a/test/python/sdpa/random_config.py +++ b/test/python/sdpa/random_config.py @@ -42,8 +42,13 @@ def get_strides_from_indices(shape, indices=[0, 1, 2, 3], gaps=[0, 0, 0, 0], rng def packed_token_capacity(seq_lens): - """Token capacity of a packed (ragged) buffer, rounded up to a multiple of 64.""" - return max(64, ((sum(seq_lens) + 63) // 64) * 64) + """Token capacity of a packed (ragged) buffer: total tokens rounded up to the + next multiple of 64, always strictly greater than the total. The surplus + guarantees every ragged buffer has a capacity tail past the last ragged + offset — the harnesses NaN-poison that tail so an engine that reads it + (e.g. binding K/V views to capacity instead of live token counts, + GitHub #624) fails deterministically instead of only on recycled memory.""" + return (sum(seq_lens) // 64 + 1) * 64 def get_strides_from_layout(shape, layout, gaps=[0, 0, 0, 0], rng_geom=None): diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 9e746ae39..2b7c9c3f3 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -700,6 +700,10 @@ def test_sdpa_fp8_fwd_L0(env_info, test_no, request, cudnn_handle): output_type=RandomChoice({torch.float8_e4m3fn: 1, torch.float8_e5m2: 1, torch.float16: 2}), with_sliding_mask=SlidingWindowMaskGenerator(causal=10, left_window_only=5, right_window_only=5, band_around_diag=10, no_mask=10), diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}), + # KNOWN GAP: a dense "padded" draw currently runs as full — exec_sdpa_fp8 + # binds seq_len tensors only for the paged and ragged paths, so the + # padding mask is never applied here (only the ragged fp8 suites below + # exercise real padding). is_ragged_or_padded_or_full=RandomChoice({"ragged": 0, "padded": 1, "full": 1}), with_sink_token=RandomChoice({True : 1, False : 2}), ) as randomization_ctx: @@ -949,7 +953,13 @@ def test_sdpa_mxfp8_fwd_L0(env_info, test_no, request, cudnn_handle): output_type=RandomChoice({torch.float16: 2, torch.bfloat16: 1}), # FP16 more often for tighter tolerance testing with_sliding_mask=SlidingWindowMaskGenerator(causal=10, left_window_only=5, right_window_only=5, band_around_diag=10, no_mask=10), diag_align=RandomChoice({cudnn.diagonal_alignment.TOP_LEFT : 1, cudnn.diagonal_alignment.BOTTOM_RIGHT : 1}), - is_ragged_or_padded_or_full=RandomChoice({"ragged": 0, "padded": 1, "full": 3}), + # full-only: the sdpa_mxfp8 API has no seq_len/padding arguments, so a + # "padded" draw would silently run dense-full (exec_sdpa_mxfp8 never + # reads seq_len_q/kv) and inflate padded coverage. When the API grows + # seq-len support, re-add padded/ragged draws — the shared + # packed_token_capacity / convert_uniform_to_packed helpers then give + # the NaN-poisoned capacity tails that catch the GitHub #624 class. + is_ragged_or_padded_or_full=RandomChoice({"full": 1}), with_sink_token=RandomChoice({True : 1, False : 2}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed)