diff --git a/python/cudnn/fla/gated_delta_rule.py b/python/cudnn/fla/gated_delta_rule.py index 11ace7f7a..4591ce8d3 100644 --- a/python/cudnn/fla/gated_delta_rule.py +++ b/python/cudnn/fla/gated_delta_rule.py @@ -94,7 +94,11 @@ def _to_native( beta = beta.to(q.dtype) if use_beta_sigmoid_in_kernel else beta.float() def thd(t): - return t.reshape(-1, *t.shape[2:]) + # FLA's fused QKV short-conv returns one compact [B,T,Q+K+V] + # allocation and splits it into strided q/k/v views. The native GDN + # kernels require compact THD inputs, so materialize only when needed; + # contiguous() is a no-op for the usual already-compact inputs. + return t.reshape(-1, *t.shape[2:]).contiguous() g2, beta2 = thd(g), thd(beta) if g2.shape[-1] != HO or beta2.shape[-1] != HO: diff --git a/test/python/linear_attention/test_fla_compat.py b/test/python/linear_attention/test_fla_compat.py index f78c02b70..d8bda57cc 100644 --- a/test/python/linear_attention/test_fla_compat.py +++ b/test/python/linear_attention/test_fla_compat.py @@ -188,6 +188,51 @@ def check(name, a, b, ref): check("d" + n, lv_fla[n].grad, lv_cud[n].grad, lv_ref[n].grad) +def test_parity_fused_layer_path_with_packed_qkv_views(): + """FLA's fused short-conv splits one packed output into non-compact Q/K/V + views. The shim must compact those views before entering native GDN while + preserving gradients back to the packed allocation.""" + B, T, H, HV, K, V = 1, 256, 16, 48, 128, 128 + widths = (H * K, H * K, HV * V) + gen = torch.Generator(device="cuda").manual_seed(4) + + packed = torch.randn(B, T, sum(widths), generator=gen, device="cuda", dtype=torch.bfloat16) + graw = torch.randn(B, T, HV, generator=gen, device="cuda", dtype=torch.bfloat16) + braw = torch.randn(B, T, HV, generator=gen, device="cuda", dtype=torch.bfloat16) + A_log = torch.log(torch.empty(HV, device="cuda").uniform_(0.1, 16, generator=gen)) + dt_bias = torch.randn(HV, generator=gen, device="cuda") + + def leaves(): + p = packed.detach().clone().requires_grad_(True) + q, k, v = p.split(widths, dim=-1) + lv = { + "packed": p, + "q": q.reshape(B, T, H, K), + "k": k.reshape(B, T, H, K), + "v": v.reshape(B, T, HV, V), + "graw": graw.detach().clone().requires_grad_(True), + "braw": braw.detach().clone().requires_grad_(True), + "A_log": A_log.detach().clone().requires_grad_(True), + "dt_bias": dt_bias.detach().clone().requires_grad_(True), + } + assert not lv["q"].is_contiguous() + assert not lv["k"].is_contiguous() + assert not lv["v"].is_contiguous() + return lv + + lv_fla, lv_cud = leaves(), leaves() + o_fla = _run_fused(chunk_gated_delta_rule, lv_fla) + o_cud = _run_fused(shim, lv_cud) + assert last_path() == "native", f"expected cuDNN native path, got {last_path()}" + + do = torch.randn_like(o_fla) + o_fla.backward(do) + o_cud.backward(do) + assert _relL2(o_cud, o_fla) <= C_SLACK * FLOOR + for n in ("packed", "graw", "braw", "A_log", "dt_bias"): + assert _relL2(lv_cud[n].grad, lv_fla[n].grad) <= C_SLACK * FLOOR, n + + kda_ops = pytest.importorskip("fla.ops.kda") chunk_kda = kda_ops.chunk_kda from cudnn.fla.kda import make_chunk_kda, last_path as kda_last_path