From 16e04d72e21544754d9e7c039af40c837c460236 Mon Sep 17 00:00:00 2001 From: Vedaanta Agarwalla Date: Thu, 13 Aug 2026 00:13:51 -0700 Subject: [PATCH] sdpa fp8 sm100: fix THD O-descriptor row stride (latent) The per-batch O-descriptor builder takes the DECLARED per-token element stride of packed O (QH * d_v). The fp8 kernel passed CFG.TILE_O - coincidentally correct only at QH == 1 with d_v == 128, otherwise landing every batch >= 1's descriptor base at 1/QH of the true offset inside earlier batches' O rows (region b receives batches {b, b+1}'s stores, the last region goes stale). Latent today - the engine declares thd=False and no API path drives fp8 THD - but real: reproduced and isolated on SM107 silicon via sentinel-V probes, with H=1xB=N and H=NxB=1 exact and H>1xB>1 corrupt. The f16 kernel has carried the same fix since its envelope work; this aligns the fp8 kernel before any THD enablement lands on top. Co-Authored-By: Claude Fable 5 --- python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py index 213eab081..6d7385517 100644 --- a/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py +++ b/python/cudnn/sdpa/fwd/kernels/prefill_d128_fp8_sm100.py @@ -1925,6 +1925,10 @@ def _tma_swz(byte_w: int): # THD: build the per-batch O descriptor array, then launch the exact # flat batch-outermost grid (n_thd_units host-computed); grid_x = units*CGA_M. # Works at cga1 (CGA_M=1) and cga2. + # Per-token element stride of packed O (QH * d_v) — NOT CFG.TILE_O, + # which is only coincidentally right at QH == 1 and otherwise lands + # every batch >= 1's descriptor base inside earlier batches' rows + # (the f16 kernel carries the same fix). _build_o_descs_kernel( o_tensor, tma_o_desc, @@ -1932,7 +1936,7 @@ def _tma_swz(byte_w: int): seq_kv_lens_tensor, cutlass.Int32(QH), cutlass.Int32(B), - cutlass.Int32(CFG.TILE_O), + cutlass.Int32(o_tensor.stride[1]), ).launch(grid=(1, 1, 1), block=(32, 1, 1), stream=stream) grid_shape = (n_thd_units * cutlass.Int32(CFG.CGA_M), cutlass.Int32(1), cutlass.Int32(1)) else: