diff --git a/fla/ops/gdn2/chunk_bwd.py b/fla/ops/gdn2/chunk_bwd.py index bf14650806..584429b601 100644 --- a/fla/ops/gdn2/chunk_bwd.py +++ b/fla/ops/gdn2/chunk_bwd.py @@ -60,7 +60,7 @@ for num_stages in [2, 3, 4] if not (IS_NVIDIA_HOPPER and BK == 32 and num_warps == 4) ], - key=['BT', 'STATE_V_FIRST'], + key=['BT', 'K', 'V', 'STATE_V_FIRST'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) diff --git a/tests/ops/test_gdn2.py b/tests/ops/test_gdn2.py index 0fcfd48263..16badb102f 100644 --- a/tests/ops/test_gdn2.py +++ b/tests/ops/test_gdn2.py @@ -472,3 +472,43 @@ def test_layer(num_heads, num_v_heads, use_short_conv): if p.requires_grad: assert p.grad is not None, f"{name}.grad is None" assert torch.isfinite(p.grad).all(), f"{name}.grad has non-finite values" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_chunk_bwd_autotune_key_covers_head_dims(): + """The fused WY/dqkg backward must autotune each (K, V) geometry separately. + + The kernel autotunes over BK/BV tile configs, whose relative performance + depends on the head dimensions K and V. Backwards through two different + (K, V) geometries at the same chunk size must therefore leave two distinct + autotuner cache entries; if K/V are missing from the autotune key, the + second geometry cache-hits on the first geometry's winner and runs a + config that was never measured for its shape. + """ + import fla.ops.gdn2.chunk_bwd as chunk_bwd_mod + + # The module symbol is Heuristics(Autotuner(JITFunction)); walk .fn down to + # the Autotuner, which owns the config cache. + tuner = chunk_bwd_mod.chunk_gdn2_bwd_kernel_wy_dqkg_fused + for _ in range(3): + if hasattr(tuner, 'cache'): + break + tuner = tuner.fn + assert hasattr(tuner, 'cache'), "could not locate the Autotuner on the fused kernel" + + saved = dict(tuner.cache) + tuner.cache.clear() + try: + for K, V in [(32, 64), (64, 128)]: + q, k, v, g, b, w, _, _ = _rand_inputs(1, 256, 2, K, V, torch.bfloat16) + for t in (q, k, v, g, b, w): + t.requires_grad_(True) + o, _ = chunk_gdn2(q=q, k=k, v=v, g=g, b=b, w=w, use_qk_l2norm_in_kernel=True) + o.sum().backward() + assert len(tuner.cache) == 2, ( + f"expected one autotune cache entry per (K, V) geometry, got {len(tuner.cache)}; " + f"the autotune key must include K and V" + ) + finally: + tuner.cache.clear() + tuner.cache.update(saved)