Skip to content

Commit 9c7d14f

Browse files
committed
feat(v4-path-d): generalize GDN Triton-frontend Path D to KDA-parity + runnable contract
GDN Path D (FLA Triton -> TileLang -> Metal) was fixed-prefill-only (Bx64xH1xK64/V32), test_only. Generalized to the full KDA-equivalent surface: - _compile_gdn_runtime_stages parametrized (total_tokens/num_sequences/num_chunks/h/hv/k/v/varlen/ scale/initial+final state/topology), mirroring _compile_kda_runtime_stages. - gdn_fwd_runtime_call reads shapes from inputs (HV%H==0, K<=256), custom scale, init/final state, packed varlen; added GDN scalar-gate chunk-local cumsum. - Backward = reuse Path B's fused recurrent Metal VJP (gdn_apply_path_d) - same (q,k,v,beta,g) signature, matches Path A fwd ~1e-8, runs M1+ (lowering FLA chunk_gated_delta_rule_bwd would hit the same M5 GEMM wall + need new OP_TABLE coverage). 5 finite grads. - Fixed a varlen ABI bug (kkt/recompute hardcoded cu_seqlens/chunk_indices len=1 -> TVM-FFI binding error); _gdn_varlen_meta_lengths. - New parity harness scripts/v4_gdn_path_d_parity.py (8 production-like shapes vs Path A FLA-naive). Contract: v4.gdn path_d test_only/blocked -> runnable_if_available/candidate, matching the already shipped KDA row (both compile+ABI green, numeric launch M5-gated). Default dispatch UNCHANGED: without the explicit unsafe-import env + local FLA/triton checkouts, Path D fail-closes to unavailable -> auto-pick stays path_c. Verified on M4 Max: all 8 harness shapes backend_available=True (lowering+compile+ABI green); backward finite. NUMERIC LAUNCH PARITY IS M5-GATED - the FLA chunk kernels lower to TileLang mpp::tensor_ops::matmul2d (Metal 4 / Apple M5); M4 launch fails 'undeclared mpp' (KDA Path D hits the identical wall). Capture green numeric parity on M5 via the harness. Tests 65 passed (failing launch-smoke tests pre-existing on pristine main, same M5 blocker); v3 untouched.
1 parent dfb4d03 commit 9c7d14f

5 files changed

Lines changed: 766 additions & 66 deletions

File tree

cppmega_v4/_tilelang/linear_attention_path_d_real.py

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -533,6 +533,26 @@ def _full_arg(extent: int) -> Tuple[int]:
533533
return (max(int(extent), 1),)
534534

535535

536+
def _gdn_varlen_meta_lengths(
537+
constexprs: Dict[str, Any],
538+
b: int,
539+
) -> Tuple[int, int]:
540+
"""Return ``(cu_seqlens_len, chunk_indices_len)`` honoring IS_VARLEN.
541+
542+
When packed varlen is active the kkt/recompute/chunk_o kernels receive the
543+
real cu_seqlens (length N+1) and chunk_indices (2 ints per chunk); under
544+
fixed-length they are unused single-element placeholders.
545+
"""
546+
547+
t = _runtime_t(constexprs)
548+
bt = int(constexprs["BT"])
549+
if bool(constexprs.get("IS_VARLEN", False)):
550+
n = _runtime_n(constexprs, b)
551+
nt = _runtime_num_chunks(constexprs, _num_chunks(t, bt))
552+
return _runtime_cu_len(constexprs, n + 1), nt * 2
553+
return 1, 1
554+
555+
536556
def _gdn_kkt_arg_buffer_shapes(
537557
constexprs: Dict[str, Any],
538558
grid: Optional[Tuple[int, ...]],
@@ -543,13 +563,14 @@ def _gdn_kkt_arg_buffer_shapes(
543563
k = int(constexprs["K"])
544564
bt = int(constexprs["BT"])
545565
b = _batch_from_grid(grid, hv)
566+
cu_len, chunk_indices_len = _gdn_varlen_meta_lengths(constexprs, b)
546567
return {
547568
0: _full_arg(b * t * h * k), # k
548569
1: _full_arg(b * t * hv), # g
549570
2: _full_arg(b * t * hv), # beta
550571
3: _full_arg(b * t * hv * bt), # A
551-
4: _full_arg(1), # cu_seqlens, unused when IS_VARLEN=False
552-
5: _full_arg(1), # chunk_indices, unused when IS_VARLEN=False
572+
4: _full_arg(cu_len), # cu_seqlens
573+
5: _full_arg(chunk_indices_len), # chunk_indices
553574
}
554575

555576

@@ -564,6 +585,7 @@ def _gdn_recompute_arg_buffer_shapes(
564585
v = int(constexprs["V"])
565586
bt = int(constexprs["BT"])
566587
b = _batch_from_grid(grid, hv)
588+
cu_len, chunk_indices_len = _gdn_varlen_meta_lengths(constexprs, b)
567589
return {
568590
0: _full_arg(b * t * h * k), # k
569591
1: _full_arg(b * t * hv * v), # v
@@ -572,8 +594,8 @@ def _gdn_recompute_arg_buffer_shapes(
572594
4: _full_arg(b * t * hv * v), # u
573595
5: _full_arg(b * t * hv * bt), # A
574596
6: _full_arg(b * t * hv), # g
575-
7: _full_arg(1), # cu_seqlens
576-
8: _full_arg(1), # chunk_indices
597+
7: _full_arg(cu_len), # cu_seqlens
598+
8: _full_arg(chunk_indices_len), # chunk_indices
577599
}
578600

579601

0 commit comments

Comments
 (0)