Commit 9c7d14f
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
- scripts
- tests
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
533 | 533 | | |
534 | 534 | | |
535 | 535 | | |
| 536 | + | |
| 537 | + | |
| 538 | + | |
| 539 | + | |
| 540 | + | |
| 541 | + | |
| 542 | + | |
| 543 | + | |
| 544 | + | |
| 545 | + | |
| 546 | + | |
| 547 | + | |
| 548 | + | |
| 549 | + | |
| 550 | + | |
| 551 | + | |
| 552 | + | |
| 553 | + | |
| 554 | + | |
| 555 | + | |
536 | 556 | | |
537 | 557 | | |
538 | 558 | | |
| |||
543 | 563 | | |
544 | 564 | | |
545 | 565 | | |
| 566 | + | |
546 | 567 | | |
547 | 568 | | |
548 | 569 | | |
549 | 570 | | |
550 | 571 | | |
551 | | - | |
552 | | - | |
| 572 | + | |
| 573 | + | |
553 | 574 | | |
554 | 575 | | |
555 | 576 | | |
| |||
564 | 585 | | |
565 | 586 | | |
566 | 587 | | |
| 588 | + | |
567 | 589 | | |
568 | 590 | | |
569 | 591 | | |
| |||
572 | 594 | | |
573 | 595 | | |
574 | 596 | | |
575 | | - | |
576 | | - | |
| 597 | + | |
| 598 | + | |
577 | 599 | | |
578 | 600 | | |
579 | 601 | | |
| |||
0 commit comments