[DPLR] A Tilelang implementation of DPLR - #1069
Conversation
Register a TileLang backend on chunk_dplr_delta_rule (which also covers
chunk_rwkv7) via the generic dispatch system. The backend wraps the
custom-op TileLang DPLR pipeline (intra A-stage, WY repr, fused h+o
forward, streaming backward) behind a verifier that falls back to the
default Triton implementation for unsupported configurations: CP,
non-fp16/bf16 dtypes, K != V, head dims outside {64, 128}, and
chunk_size 64 off Hopper (the intra backward exceeds sm_120's 99KB
shared-memory cap there).
Also fixes three latent kernel bugs surfaced by the new route-parity
tests:
- WY inversion raced on a fragment accumulator spread across threads;
the j-reduction now runs serially per lane.
- fp16 forwards overflowed the gate-centered exp2 operands (inf/NaN);
q-side GEMM operands are kept fp32 for fp16 inputs.
- K=128 at BT=16 failed to compile (vectorize-planner issue in the
ThreadSync pass) and the fused h+o kernel used a warp count no
16-row GEMM can partition; the A-stage now compiles a non-vectorized
variant for BT <= 16 and the h+o config caps threads there.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
# Conflicts: # tests/ops/test_backends.py
The fused h+o state pass parallelizes only over (V/BV, N, H) and walks chunks serially, so it loses to the split Triton path on small grids (up to 10x slower at B2 H8 on PRO 6000) while winning at the design envelope (1.73x fwd+bwd at B8 H32 D64 cs32). The verifier now rejects calls whose state-pass grid underfills half the GPU, falling back to Triton. Route-parity tests move to shapes above the crossover. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Verified each rejection branch directly against the kernels (all raise
hard errors or silently miscompute when bypassed): cp_context is
silently dropped by the backend method (no boundary-state exchange in
the TileLang path), fp32 has no launchable stream-backward schedule,
mixed gk dtypes fail kernel dtype checks, K != V is a stream-backward
constraint, and head dims outside {64, 128} hit explicit wrappers. All
branches stay. Also raise TypeError on unexpected kwargs in the
TileLang wrapper so a direct cp_context call fails loudly instead of
being swallowed, and gate chunk_size 64 on the device's shared-memory
opt-in cap instead of Hopper-only (allows A100-class parts where the
BT=64 intra backward fits).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
FLA's own tests and many training setups keep gk in fp32 while the activations are fp16/bf16. The from_gk A-stage and the DERIVE_GE intra backward now take gk in its native dtype (gate math was already fp32 everywhere), and the backward's recycled dgk buffer is only reused when it matches gk's dtype. The verifier now requires only k/v/a/b to match q.dtype; gk may be fp16/bf16/fp32. Verified exact parity vs Triton (0.0000) for bf16/fp16, K=64/128, chunk 16/32, rect and varlen. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The TileLang DPLR chunk path now handles cp_context using FLA's boundary-state scheme: after the WY stage, each rank exchanges its compact local state contribution [c|M] via the existing chunk_gated_delta_rule_fwd_h_pre_process and folds it into the corrected initial state; the backward symmetrically folds the following ranks' dh via chunk_gated_delta_rule_bwd_dhu_pre_process before the streaming dh scan. The FLACPContext is threaded through a ContextVar (custom ops cannot take the object) and persisted on the autograd ctx so the exchange runs in the backward as well. Verified against the naive DPLR recurrence on 2 GPUs (route assertion included) and against the Triton CP path on identical inputs (ratios <= 0.006). The harness in tests/context_parallel/test_cp_dplr.py gains an op_chunk_size and a route-assertion switch for this. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Measured ~0.5x vs Triton (the non-vectorized A-stage and 2-warp h+o path at BT=16 do not pay off at K=128); fall back to Triton there. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The in-CTA cumsum and the gating loop now read gk from a shared tile loaded once, and gi is written with one batched coalesced store instead of a global load+store per serial step (wind_rwkv-style register/shared scan). chunk_dplr_fwd_intra_from_gk at B8 T4096 H32 D64 cs32 bf16: 1.85ms -> 1.52ms (-18%). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Adds a TileLang backend for the generalized delta rule DPLR “chunk” pipeline and wires it into FLA’s backend dispatch system, aiming to improve performance via additional fusion while preserving correctness by falling back to the existing Triton implementation when unsupported.
Changes:
- Register
chunk_dplr_delta_ruleunder backend dispatch (generalized_delta_rule.dplr) and add a TileLang backend implementation with verifier-based routing. - Implement TileLang kernels/stages for DPLR forward/backward (A-stage, WY, fused H/O forward, streaming backward, cumsum/layout utilities).
- Add unit tests for verifier behavior, backend availability gating, route parity, and CP routing.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/ops/test_dplr_tilelang.py | New tests covering verifier acceptance/rejection and TileLang-vs-Triton parity routing. |
| tests/ops/test_backends.py | Extends existing backend availability tests to include the DPLR TileLang backend module. |
| tests/context_parallel/test_cp_dplr.py | Adds a CP test case that asserts the TileLang DPLR route is taken. |
| fla/ops/generalized_delta_rule/dplr/chunk.py | Hooks chunk_dplr_delta_rule into the backend dispatch decorator for DPLR. |
| fla/ops/generalized_delta_rule/dplr/backends/init.py | Introduces a DPLR backend registry and registers the TileLang backend. |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/init.py | Adds DPLRTileLangBackend with availability checks and a shape/perf verifier. |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/utils.py | Implements shared chunk layout utilities for rect/varlen execution. |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/cumsum.py | Implements chunk-local cumsum (PyTorch fast path for rect, TileLang for varlen). |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/wy_fast_fwd.py | TileLang forward for WY inversion and W/U products. |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/wy_fast_bwd.py | TileLang backward for WY path (dA_ab, dA_ak, dv, dag). |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_A_fwd.py | TileLang intra-chunk A-stage forward (incl. from-gk specialization for K=64). |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_A_bwd.py | TileLang intra-chunk backward producing dq/dk/da/db/dgk (with fused q-side recompute). |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_h_fwd.py | TileLang chunk recurrent-state forward producing per-chunk state and v_new. |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_ho_fwd.py | TileLang fused H/O forward kernel (optionally with saved context buffers). |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_stream_bwd.py | TileLang streaming backward kernel(s) with schedule selection based on SMEM caps. |
| fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk.py | Main TileLang DPLR chunk op implementation using torch.library.custom_op + autograd wiring. |
Comments suppressed due to low confidence (1)
fla/ops/generalized_delta_rule/dplr/backends/tilelang/init.py:84
- The
chunk_size == 64verifier path callstorch.cuda.get_device_properties(...)via_smem_optin_bytes(...)before checkingq.is_cuda. If a caller runs the verifier on CPU tensors withchunk_size=64, this will raise instead of cleanly returning a fallback-to-Triton rejection.
if _smem_optin_bytes(q.device.index or 0) < smem_need:
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if q.dtype not in (torch.float16, torch.bfloat16): | ||
| return False, f"TileLang backend does not support dtype {q.dtype}; fall back to Triton" |
| return backend_module.RWKV6TileLangBackend | ||
| return backend_module.KDATileLangBackend | ||
| if backend_module is kda_tilelang_backend: | ||
| return backend_module.KDATileLangBackend | ||
| return backend_module.DPLRTileLangBackend |
| if assert_tilelang: | ||
| assert route_spy, "TileLang backend route was not taken" |
| ok, reason = DPLRTileLangBackend().chunk_dplr_delta_rule_verifier( | ||
| *_verifier_inputs(K=K, dtype=dtype), chunk_size=chunk_size, cp_context=SimpleNamespace(), | ||
| ) |
Interior chunks now load their kg/w/qg/v/u/A tiles with bulk T.copy (vectorized cp.async) instead of scalar predicated elementwise loops; tail chunks keep the masked path. chunk_dplr_fwd_ho at B8 T4096 H32 D64 cs32 bf16: 1.42ms -> 1.03ms (-27%). The chunk loop is switched to T.Pipelined so num_stages pipelining is available; it is a no-op at the current default (0), and stages 2/3 measured no further gain at the design shape, so the default stays 0. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The three TL_DISABLE_DATA_RACE_CHECK: True entries in chunk_A_fwd.py are unnecessary: with the checker enabled, the intra and from_gk kernels compile warning-free at BT 16/32/64, K 64/128, bf16/fp16 (the WY race that motivated silencing was fixed properly in prepare_wy_repr_fwd). Keeping the warning-only checker on matches the rest of the package and preserves the safety net for future edits. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 16 out of 16 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (7)
tests/ops/test_dplr_tilelang.py:53
- This test asserts the verifier accepts a CUDA configuration, but it will run (and fail) on CPU-only environments because
_verifier_inputs()usesfla.utils.devicewhich can be'cpu'. Add a CUDA guard.
def test_chunk_verifier_accepts_fp32_gk():
tests/ops/test_dplr_tilelang.py:61
- This test expects the verifier to accept a CUDA-only path, but it can run on CPU-only CI where
_verifier_inputs()usesfla.utils.device='cpu', causing the verifier to reject. Gate it on CUDA availability.
def test_chunk_verifier_accepts_chunk64_on_large_smem_device(monkeypatch):
tests/ops/test_dplr_tilelang.py:80
test_chunk_verifier_rejectsdepends on CUDA-specific rejection reasons (dtype mismatch, SMEM caps, etc.). On CPU-only environments the verifier will instead reject with the generic CUDA-only message, breaking the expectedreason in whyassertions. Add a CUDA guard.
def test_chunk_verifier_rejects(monkeypatch, case: str, reason: str):
tests/context_parallel/test_cp_dplr.py:276
- The test patches
DPLRTileLangBackend.chunk_dplr_delta_rulebut never restores it. Even though this runs in a spawned worker, restoring avoids leaking the patch to later code paths in the same process (e.g., if additional assertions/logging are added or the worker is reused).
if assert_tilelang:
assert route_spy, "TileLang backend route was not taken"
tests/ops/test_dplr_tilelang.py:45
- These verifier-acceptance tests will fail on CPU-only CI because
fla.utils.devicecan be'cpu', making the verifier returnFalse(q.is_cudais required). Gate these tests on CUDA availability (or otherwise ensure inputs are CUDA tensors).
This issue also appears in the following locations of the same file:
- line 53
- line 61
- line 80
@pytest.mark.parametrize(('K', 'dtype', 'chunk_size'), [(64, torch.bfloat16, 32), (128, torch.bfloat16, 32), (64, torch.float16, 16)])
tests/context_parallel/test_cp_dplr.py:480
- This test unconditionally deletes
FLA_TILELANGin thefinallyblock. If the variable was already set in the environment (e.g. a user/CI opt-in), this permanently removes the prior value for the remainder of the test session. Save and restore the previous value instead.
os.environ['FLA_TILELANG'] = '1'
try:
run_cp_test_with_spawn(
world_size=2,
test_name="CP2_TileLang",
T=1024, H=128, D=64,
lengths=[400, 624],
dtype=torch.bfloat16,
op_chunk_size=32,
assert_tilelang=True,
)
finally:
del os.environ['FLA_TILELANG']
fla/ops/generalized_delta_rule/dplr/backends/tilelang/init.py:104
- The small-grid heuristic ignores
cp_context: whencp_contextis provided,chunk_dplr_delta_rulecallers typically passcu_seqlens=Noneand rely oncp_context.cu_seqlensinstead. Usingq.shape[0]in that case underestimatesn_seqsand can incorrectly reject the TileLang route for CP workloads.
bv = 64 if v.shape[-1] <= 64 else 32
n_seqs = len(cu_seqlens) - 1 if cu_seqlens is not None else q.shape[0]
grid = n_seqs * q.shape[2] * ((v.shape[-1] + bv - 1) // bv)
build_varlen_chunk_layout was rebuilt on every call (~0.9ms of host time in small op launches, and torch.diff's int32 path stalls the stream). It is now memoized by cu_seqlens identity (bounded deque of 4, refs held so ids cannot recycle — same contract as fla's tensor_cache), and lengths are computed with a slice-sub instead. Varlen fwd+bwd vs Triton at hidden4096 D64 cs32: 1.00x -> 1.44x. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The tvm_ffi adapter resolves dynamic output shapes on every launch with str(tirx.Var) comparisons, which lower to hundreds of TIR _script calls per launch on multi-output kernels (~2 ms for the 9-output from_gk A-stage). Small and varlen workloads are dominated by this host cost: measured 0.16x-0.7x vs Triton at <=8K tokens while the GPU kernels themselves were already faster. Drop out_idx from the remaining forward factories and pass preallocated outputs instead, matching the pattern the backward kernels already used. Forward host time per call drops 2.91 ms -> 0.27 ms (varlen h1024); varlen fwd+bwd speedups move from 0.16-1.06x to 0.73-1.97x on a contended box.
Every stage rebuilt the same rect layout per call (~6 small CUDA ops each, ~30 extra launches per forward). Key it by (B, T, chunk_size, device) like the varlen layout cache.
The verifier ignored cp_context, so a CP call with initial_state, output_final_state, or no cu_seqlens was accepted and then raised inside the backend; reject these and fall back to Triton instead. The grid guard also read the sequence count from the packed batch dim (always 1 under CP) — take it from the context's cu_seqlens. Review follow-ups: make _backend_cls exhaustive in test_backends.py, restore the route spy in the CP worker, and test CP verifier acceptance with a real context instead of a bare SimpleNamespace.
zhiyuan1i
left a comment
zhiyuan1i
left a comment
There was a problem hiding this comment.
Thanks — impressive work. We checked the generated CUDA and the headline claims hold: the fused fwd h+o single sweep (no h/v_new global round-trip) and the never-materialized bwd dh are real, and the stage mapping onto the Triton pipeline is complete. A few items to address before merge:
- Verifier routing must require
safe_gate=True and chunk_size == 64(inline). The TileLang A-stage always uses the mid-chunk-centered tensorcore scheme, which is only valid while(chunk_size/2) * max|gk|stays in fp32 range. Defaultsafe_gate=Falsecallers assert nothing about gate range (Triton'ssub_intrais correct for arbitrary gates — those calls must stay on Triton), and the documented safe range [-5, 0) only covers BT≤32: at BT=64 the centered half-range is 160 log2 and overflows fp32. RWKV7 satisfies both conditions (boundedw, explicitchunk_size=64). Longer term, consider GDN2's explicitlower_boundcontract. - Please add route-parity tests for
initial_state=None+output_final_state=False— every parity case currently passes h0 and requires the final state, but RWKV7 training hits exactly those uncovered branches. A checkpoint-elision test and a torch.compile/opcheck smoke test would also be worthwhile, since both are headline features. - Robustness (details inline):
initial_stateneeds a contiguity guard; exceptions from the impl (JIT failure, schedule selection) escape instead of falling back, contrary to the PR description; and the (K=128, chunk64) smem gate does not match the stream-bwd requirement — on A100 this passes the verifier and then raises. A100 is a target platform for us. - Benchmark evidence: per the repo's MR-readiness conventions, kernel PRs need before/after numbers on the same hardware, dense + varlen. If Hopper access is the blocker, we can run them on our H100 — happy to help.
- Style: please move the raw
torch.cuda.get_device_capability/propertiesqueries into smallfla.utilshelpers, per the repo guidelines.
| output_final_state: bool = False, | ||
| cu_seqlens: torch.LongTensor | None = None, | ||
| cu_seqlens_cpu: torch.LongTensor | None = None, | ||
| safe_gate: bool = False, |
There was a problem hiding this comment.
safe_gate is accepted here but never checked. Please route only when safe_gate=True and chunk_size == 64: the kernels' mid-chunk-centered tensorcore scheme overflows fp32 when (chunk_size/2) * max|gk| ≳ 127 log2. safe_gate=False callers make no gate-range assertion — Triton's sub_intra stays correct for arbitrary gates, so those calls must not be rerouted. And safe_gate=True alone is not sufficient for BT=64: the documented [-5, 0) range gives a 160-log2 half-range at BT=64 (it only licenses BT≤32). RWKV7 satisfies both conditions (w architecturally clamped to (-0.61, 0), explicit chunk_size=64).
| if chunk_size == 64: | ||
| # the intra backward at BT=64 needs 131200B (K=64) or 148480B | ||
| # (K=128) of shared memory per block; smaller caps cannot run it | ||
| smem_need = 131200 if k.shape[-1] <= 64 else 148480 |
There was a problem hiding this comment.
This gate is derived from the intra-bwd requirement, but the stream-bwd schedules need more: at (K=128, BT=64) high=295424B and low=167936B, both above A100's 166912B optin — so on A100 the call passes this check and then _select_stream_bwd_schedule raises at backward time instead of falling back. A100 is a target platform for us; please derive the gate from the schedule's own arithmetic (or reject the combo here).
| scale_f = float(q.shape[-1] ** -0.5 if scale is None else scale) | ||
| n_seqs = len(cu_seqlens) - 1 if cu_seqlens is not None else q.shape[0] | ||
| if initial_state is not None: | ||
| h0 = initial_state |
There was a problem hiding this comment.
initial_state needs a contiguity guard: a non-contiguous fp32 h0 would be read with wrong strides by the kernels (the Triton path is protected by input_guard, which this backend bypasses). Please normalize here or check contiguity in the verifier.
| ): | ||
| if not torch.cuda.is_available(): | ||
| return [{"BV": 16, "threads": 64}] | ||
| major = torch.cuda.get_device_capability()[0] |
There was a problem hiding this comment.
This queries the current device rather than q.device — on multi-GPU with a different current device the configs come from the wrong GPU while the Triton path (via input_guard's device ctx) works. Same pattern in _ho_fragment_merge_flags and _chunk_h_fwd_configs. Also consider moving these raw torch.cuda queries into fla.utils helpers, per repo guidelines.
| dgk_out_f = dgk_out.reshape(N_tokens, H, K).contiguous() | ||
| # Temporary ABI shim: the fused kernel specialization ignores dAqk/dAqb, | ||
| # but the shared JIT signature still requires same-shape bf16 inputs. | ||
| dummy_dA = q.new_empty((N_tokens, H, BT)) |
There was a problem hiding this comment.
Nit: two full-size [N_tokens, H, BT] allocations that the FUSE_QSIDE_DA kernel never reads (~128MB dead memory at T=8192/H=64/BT=64). Consider 0-elem tensors or splitting the kernel signature.
| return wy_fast_bwd_tl | ||
|
|
||
|
|
||
| def chunk_dplr_bwd_wy( |
There was a problem hiding this comment.
Nit: dead code — only the _into variant is used.
| dw_f = dw.reshape(N_tokens, H, K).contiguous() | ||
| du_f = du.reshape(N_tokens, H, V).contiguous() | ||
| dv0_f = dv0.reshape(N_tokens, H, V).contiguous() | ||
| dA_ab_f = dA_ab_out.reshape(N_tokens, H, BT).contiguous() |
There was a problem hiding this comment.
Nit: reshape(...).contiguous() on an out-param silently writes into a temporary copy if a caller ever passes a non-contiguous tensor — the kernel's writes would be dropped. Please assert out.is_contiguous() here and in the other _into wrappers.
| ) | ||
|
|
||
|
|
||
| def test_cp2_tilelang_route(): |
There was a problem hiding this comment.
Nit: this should skip when DPLRTileLangBackend.is_available() is False — on a 2-GPU machine without tilelang/nvcc it fails instead of skipping (the requires_tilelang_route pattern in test_dplr_tilelang.py covers this). Also please restore FLA_TILELANG to its previous value rather than del.
| token = _DPLR_CP_CONTEXT.set(cp_context) | ||
| try: | ||
| if disable_recompute: | ||
| if not is_compiling and checkpoint_phase == _CHECKPOINT_PHASE_FORWARD or ( |
There was a problem hiding this comment.
Nit: please parenthesize — A and B or (C and D and E) happens to parse correctly here but is fragile.
- verifier: route safe_gate inputs, gate chunk_size=64 by the schedule-derived smem footprint, and log fallbacks - chunk wrapper: elide the checkpoint context when recompute is disabled, pass CP metadata through a ContextVar side channel, and keep h0 contiguous - per-device kernel configs via cached get_device_capability / get_device_smem_optin helpers that respect device_index - tests: extend TileLang DPLR parity coverage and the CP constraint checks
…for chunk_size=64 - stream bwd: new "mid" schedule selected between high and low_v2 aliases K/V in shared and stages the output tile so K=128/BT=64 fits the cc90 opt-in smem budget (stream kernel 17.0 -> 2.44 ms on H800; cs64 fwd+bwd 0.37-0.55x -> 0.54-0.81x vs Triton) - bwd intra, wu_fwd and wy_fast_bwd: 256 threads and bulk T.copy tile loads on cc90 with BT>=64 only; cc120 and BT<=32 keep the previous scalar path, where the bulk path measured slower
The K=128 low-smem stream backward runs one serial chunk scan per (seq, head) block with no V split, and its 97KB shared footprint leaves no room to prefetch, so below half the SM count the serial chain is exposed: measured 0.44-0.63x vs Triton at varlen D128 with N*H=48 on the PRO 6000. The verifier now rejects low_v2 below the same half-SM fill the forward grid guard already uses, and no longer relies on the cs64-only schedule check (acceptance implies launchability at every chunk size). The low schedule also moves to 256 threads at K=V=128, halving the persistent (K, V) fp32 dh fragment to 64 registers per thread (stream kernel 2.81 -> 2.64 ms).
Compare the routed op against the per-token fp32 recurrence (dplr_recurrence) for output, final state and all input gradients at the same 0.007/0.008 ratios the Triton path is held to in test_dplr_delta, covering rect and varlen, bf16/fp16, chunk 16/32/64 and the safe-gate stress regimes (saturated lower bound at cs32, RWKV7-clamped at cs64). Also adds the verifier reject case for the underfilled low-smem stream backward grid guard.
…unks The low_v2 schedule loaded every input tile and q-side slice with scalar predicated element accesses, which cap at ~1.5TB/s. Mirror the high/mid kernel: interior (full-tile) chunks bulk-load qg/bg/w/kg/do/A via vectorized T.copy, v_new and v are staged per chunk like do so the q-side slices stay in shared memory, and exactly-divisible h slices bulk-load. Boundary chunks and ragged Vs keep the scalar predicated path. Gate: tests/ops/test_dplr_tilelang.py 48 passed, 6 skipped. Bench (PRO 6000, cs32 D128 bf16, fwd+bwd vs Triton, dr=0/dr=1): varlen h1024 N6 0.62/0.47 -> 0.96/0.78, rect h1024 0.71/0.57 -> 0.98/0.83, varlen h4096 N3 0.88/0.73 -> 1.19/1.04, varlen h2560 N5 2.26/1.50 -> 3.10/2.18, rect h2560 1.06/0.86 -> 1.24/1.03.
…roups The low_v2 schedule reduced the kg*dk and bg*db dgk terms as two serial over-BT chains with only K lanes active. Port the high/mid schedule's (4, K) lane-group split so all threads participate, and account for the new staging tile in the shared-memory formula. Gate: smoke parity (forced native vs Triton) max-rel ~1e-2 bf16. Bench (PRO 6000, cs32 D128 bf16, fwd+bwd vs Triton, dr=0/dr=1): varlen h1024 N6 0.96/0.78 -> 1.05/0.87, rect h1024 0.98/0.83 -> 1.04/0.89, varlen h4096 N3 1.19/1.04 -> 1.25/1.11, varlen h2560 N5 3.10/2.18 -> 3.28/2.34, rect h2560 1.24/1.03 -> 1.27/1.06.
…eLang backend Addresses the remaining PR review style item: raw torch.cuda.get_device_capability/get_device_properties/current_device queries now go through fla.utils helpers (get_device_capability, get_device_smem_optin), matching the verifier. The dead current_device fallbacks in the ho/h config pickers are dropped (all call sites pass kg.device.index); the stream-bwd schedule selection keeps the device name fetch only on its error path. Gates: tests/ops/test_dplr_tilelang.py + tests/ops/test_backends.py 64 passed/6 skipped; tests/context_parallel/test_cp_dplr.py 5 passed/2 skipped.
The four fp32 (BT,BT) dA staging tiles become two: dA1/dA2 serve dAqk/dAqb for the q-side GEMM group, then are overwritten with dAak/dAab (strict-lower mask applied after the overwrite) for the a-side group. Per-accumulator GEMM order is unchanged, so results are bit-identical to the four-tile layout; the footprint drops by 2*(BT,BT) fp32 tiles, taking the BT=64/BK=32 config from 131200B to 98432B — under cc120's 101376B smem optin, one of the two stages that blocked native chunk_size=64 on 99KB devices. Evidence: bit-exact vs HEAD at BT=32 rect+varlen with boundary chunks (max diff 0.0); BT=32 perf flat (rect h2560 fwd+bwd: D64 8.659 vs 8.695ms, D128 15.723 vs 15.728ms vs HEAD). Gates: tests/ops/test_dplr_tilelang.py 48 passed/6 skipped.
The gating loop re-reads gk from global memory (L2-resident after the serial scan) instead of gk_shared, ending the tile's lifetime at the scan so the allocator packs it against the fp32 operand tiles. Values are identical, so results are bit-exact; the BT=64 footprint drops from 98816B (bf16 gk) / 107008B (fp32 gk) to 98304B for both — under cc120's 101376B optin. Together with the A-backward staging aliasing this makes every chunk_size=64 D64 stage launchable on 99KB devices; native cs64-D64 ran end-to-end for the first time there, with parity vs the fp32 naive recurrence inside the gate envelope (o 0.0037 vs 0.007 threshold, worst dgk 0.0064 vs 0.008). Evidence: bit-exact vs HEAD at BT=32 rect+varlen (max diff 0.0); cs64 acceptance arithmetic stays unchanged until native cs64 wins a bench. Gates: tests/ops/test_dplr_tilelang.py 48 passed/6 skipped.
With the fused A-backward down to 98432B and the from_gk A-forward at 98304B, every BT=64 D64 stage fits a 101376B optin (the stream backward selects low_v2 at 82944B), so the verifier's binding constant drops 131200 -> 98432 and cc120-class devices accept cs64 at K=64. K=128 cs64 stays rejected there — every stream schedule is >=167936B — and the D128 control rows measure exactly 1.00x, confirming the fallback costs nothing. Native cs64-D64 wins every measured cell vs Triton cs64 (safe_gate, rwkv7 gates, bf16): routed sweep rect h2560/h4096 dr=0 fwd 1.76/1.85x, fwd+bwd 1.70/1.71x, dr=1 1.37-1.39x fwd and 1.20-1.21x fwd+bwd; varlen N2/N3 1.50-1.73x fwd and 1.25-1.65x fwd+bwd. Native parity vs the fp32 naive recurrence sits inside the gate envelope on all outputs (o 0.0037 vs 0.007, worst dgk 0.0064 vs 0.008). The reject-case chunk64_small_smem_k64 becomes an explicit accept test; the K=128 small-smem reject stays, now documented as stream-bound. Gates: tests/ops/test_dplr_tilelang.py 52 passed/2 skipped (the four cs64 cases now run natively); tests/ops/test_backends.py 16 passed; tests/context_parallel/test_cp_dplr.py 5 passed/2 skipped; tests/ops/test_dplr_delta.py 28 passed + 1 pre-existing D100-cs64 suite flake (rejected by head dim before the cs64 arithmetic; isolated reruns pass).
The 256-thread + vectorized-T.copy configs for wu_fwd and wy_fast_bwd were gated cc90-only at BT>=64, but the kernel-level win reproduces on cc120 (wu_fwd BT64 2.12x at D64 / 1.67x at D128, wy_fast_bwd BT64 1.47x), and bulk-vs-scalar is bit-exact at BT64 rect+varlen including boundary chunks. With cs64 now verifier-accepted on 99KB devices this was the largest cs64 deficit: wu_fwd ran 3.53ms/iter vs Triton's 0.88ms at dr=1. End-to-end cs64-D64 vs Triton: dr=1 fwd+bwd 1.18-1.23x -> 1.30-1.38x, dr=0 fwd+bwd 1.59-1.69x -> 1.79-1.90x, dr=0 fwd up to 2.16x. BT=32 keeps the scalar path (bulk measured flat to slightly slower there). Gates: tests/ops/test_dplr_tilelang.py 52 passed/2 skipped (the native cs64 cases exercise the bulk path).
…ckward The ctx forward already computes qg/kg/ag/bg/w/A_ab_inv/A_ak/A_qk/A_qb; return them as extra op outputs so the disable_recompute backward reads them from ctx instead of recomputing the from_gk/intra A-stage, prepare_wy_repr, and wu. Only the cumsum gi is still recomputed, matching FLA's recompute contract. The saved set stays read-only (mutates_args=()), so gradients land in fresh buffers instead of recycling the intermediates. CP keeps the recompute path (its boundary pre-process reads u) and the dr=0 plain path is untouched. A/B vs HEAD: D128 bit-exact; D64 shifts by bf16 rounding noise only (backward now uses the forward's own cumsum gi rather than from_gk's internally re-derived one). Gates: tilelang 52p/2s, backends 16p, CP 5p/2s. dr=1 fwd+bwd tl-ms: rect h1024 B8 8.77 -> 8.19, rect h2560 B4 11.07 -> 10.37, rect h4096 B4 17.61 -> 16.55, varlen h2560 N5 5.49 -> 4.87, varlen h1024 N5 2.15 -> 1.95; routed dr=1 fwd+bwd now 1.40-1.62x vs Triton (was 1.31-1.40x); dr=0 and D128 control rows unchanged.
Two changes to the generic prepare_wy_repr inversion kernel (BT 16/32): the inner j-reduction now stops at row_i instead of running the full BT range (v[j>=row_i] is identically zero, so the skipped terms contribute nothing), and interior chunks bulk-copy A_ab into shared memory and apply the strict-lower mask there while boundary chunks keep the scalar predicated path (a bulk store measured slower, so stores stay scalar). Also drop the vestigial single-entry _WY_INV_CONFIGS table. The BT=64 blocked kernel already matches Triton and the same probes measured flat there, so it is untouched. w/u/A_ab_inv bit-exact vs HEAD at BT 16/32/64, rect+varlen including boundary chunks. Kernel level (B8 T4096 H32/H64 grids): BT16 0.170/0.325 -> 0.106/0.259 ms, BT32 0.326/0.623 -> 0.196/0.384 ms (exact parity with Triton's 0.193/0.383). Gates: tilelang 52p/2s, backends 16p, CP 5p/2s. End-to-end cs32 routed dr=0 fwd+bwd: 1.34-1.36x rect h2560/h4096 B8, 1.29x varlen h2560 N3.
…k A-stage The ctx forward (the disable_recompute path) was the only entry still doing the standalone cumsum + tensorcore intra split at head_dim 64; the plain forward has used the fused from_gk stage for a long time. Besides the extra cumsum kernels, the tensorcore intra's torch-side gate preparation showed up as ~1.4 ms/iter of elementwise work in profiles, which was most of the dr=1 forward deficit against Triton. A/B: D128 bit-exact; D64 differs at the bf16 rounding class already accepted with the wy save set (forward from_gk gi vs backward cumsum gi). Gates: tilelang 52p/2s, backends 16p, CP 5p/2s. Routed dr=1: cs32 rect fwd 0.94-0.95x -> 1.29-1.34x, fwd+bwd 1.02-1.03x -> 1.13-1.14x; cs32 varlen 1.12/1.17x -> 1.23/1.23x; cs64 fwd 1.53-1.57x -> 1.87-2.06x (near dr=0 parity), cs64 fwd+bwd now 1.55-1.70x. dr=0 cells unchanged.
The cc90 config already used 256 threads to hide unpipelined load latency; the same applies at BT=64 on cc120 (bwd_intra -14%), while BT<=32 regresses (+16%) since occupancy already hides latency there, so the extra warps are gated to BT>=64. The BT=64 bulk-copy path stays cc90-only: its staging tiles need 106496B, over cc120's 99KB cap. Bit-exact vs HEAD at cs64-D64 and cs32-D64, rect+varlen. Gates: tilelang 52p/2s, backends 16p, CP 5p/2s. cs64 fwd+bwd: dr=0 1.86-1.90x -> 1.91-1.95x rect, dr=1 1.55-1.62x -> 1.60-1.67x.
- remove the never-referenced _H_FWD_CONFIGS/_HO_FWD_CONFIGS constants and trim the config helpers to the parameters their bodies actually use - remove the allocate_state_cache escape hatch (no caller) and the dead scale_value/num_stages/USE_SWIZZLE kernel parameters (USE_SWIZZLE was a provable no-op: tilelang's use_swizzle returns None when disabled; num_stages=0 disables pipelining, so the loop is plainly serial) - merge the two fwd_ho ctx wrappers into one store_context entry point - drop the unused cc parameter from stream_bwd_schedule_or_none and the unused qside_bv override from stream_low_bwd_config - fix the stale from_gk docstring and the double-assigned h_out
…ang backend With the backend active, unbounded-gate calls must stay on Triton's robust sub_intra path and return correct, finite results at every supported chunk size.
chunk_dplr_delta_rule gains a lower_bound kwarg asserting lower_bound <= gk < 0. When the bound fits the chunk size (per-row centered exponents stay in fp32 range) it canonicalizes to safe_gate=True on the Triton path and lets the TileLang verifier accept the call without safe_gate. A non-negative bound raises ValueError on every route.
low_v2 -> low (the high/mid/low smem ladder), the fused stream backward drops the dhu_o suffix (chunk_dplr_bwd_stream_*), and the BT=64 WY-inverse factory spells out _bt64.
|
Thanks for the thorough review — all five items are now addressed in the pushed commits. Details per item: 1. Gate-bound routing. The verifier now requires a caller-asserted gate bound, exactly as you framed it. 2. Coverage of the RWKV7-training branches. 3. Robustness. (a) 4. Benchmarks. Same-hardware before/after, dense + varlen, now in the PR body: PRO 6000 (sm_120) and H800 (sm_90), bf16, routed 5. Style. Device queries now go through small Inline nits (dummy 1-elem alloc in Rebased onto current |
zhiyuan1i
left a comment
zhiyuan1i
left a comment
There was a problem hiding this comment.
Thanks for the thorough follow-up — the A100 smem gate now correctly rejects (K=128, cs64) and accepts (K=64, cs64), the device queries are properly routed through fla.utils helpers, and the new parity/opcheck/compile/checkpoint tests are real additions. Four items still block merge:
-
Numerical-safety hole (inline): the verifier skips the gate-bound check entirely when
safe_gate=True, sosafe_gate=True, chunk_size=64— a combination the public docstring still advertises as ~[-5, 0) — enters the centered TileLang A-stage with worst-case exponents of ~231–238 log2, past fp32 range. The new tests even assert this unsafe combination is accepted. The RWKV7 built-in caller is safe (clamped w → ~28 log2), but the public path is not. -
Positional-argument break (inline):
lower_boundis inserted betweensafe_gateandchunk_sizein the public signature, so existing positional callers silently bind their chunk_size into lower_bound. Please append it after the existing parameters (keyword-style, per the GDN2 convention). -
CP rank-local fallback: the runtime TileLang→Triton fallback is decided per rank. Under context parallelism, one rank failing over while its peers stay on TileLang can diverge the collective sequence or hang. The fallback decision needs to be collective-uniform (e.g. all-reduce the verifier result), or disabled under CP.
-
H800 cs64 regression still routed: your own table shows H800 cs64 fwd+bwd at 0.54–0.81x, yet the verifier accepts cs64 on sm90 — contradicting the "only accepts when measurably faster" claim. Please either reject cs64 on sm90 (fall back to Triton) or provide per-shape data justifying the accepted subset; the Impact section's "cs64 up to ~2x on both architectures" also needs correcting against your H800 numbers.
Two smaller items:
gate_bound_is_safeis off by one: it useschunk_size/2rows, butcentered_gespanschunk_size/2 + 1rows. At cs16 the helper accepts lower_bound≈-10.4, where the real exponent reaches ~135 log2 and overflows. Please tighten the formula.- The varlen benchmark numbers aren't reproducible: only total tokens and sequence count N are given — no cu_seqlens, generation rule, seed, or script. Please add the harness details (and the 7.08x PRO 6000 outlier deserves a sanity note).
Once the verifier gates both safe_gate paths on the (corrected) bound formula, the signature keeps positional compatibility, the CP fallback is made uniform, and sm90 cs64 stops routing to a slower backend, this should be in good shape.
| chunk_size = 16 if chunk_size is None else chunk_size | ||
| if chunk_size not in (16, 32, 64): | ||
| return False, f"TileLang backend supports chunk_size 16/32/64, got {chunk_size}; fall back to Triton" | ||
| if not safe_gate: |
There was a problem hiding this comment.
Blocking: when safe_gate=True, the gate-bound check is skipped entirely, so safe_gate=True, chunk_size=64 enters the mid-chunk-centered A-stage unconditionally. The public docstring still describes safe_gate gates as roughly [-5, 0), which at cs64 reaches 5×32×log2e ≈ 231 log2 — well past the fp32 exp2 range (~127). safe_gate=True callers assert a bound, but nothing here enforces that the asserted bound fits the chunk size. The verifier should apply gate_bound_is_safe to both paths: for safe_gate=True use the documented default bound (-5), and reject cs64 for it, same as the lower_bound path.
| cu_seqlens: torch.LongTensor | None = None, | ||
| cu_seqlens_cpu: torch.LongTensor | None = None, | ||
| safe_gate: bool = False, | ||
| lower_bound: float | None = None, |
There was a problem hiding this comment.
Blocking: inserting lower_bound between safe_gate and chunk_size breaks every existing positional caller — a positional chunk_size now binds into lower_bound (a float), silently changing behavior instead of raising. Please append new parameters after the existing public ones (or make it keyword-only), per the GDN2 convention.
…bound The verifier trusted safe_gate=True unconditionally, but the documented safe_gate range [-5, 0) overflows the mid-chunk-centered A-stage at chunk_size 64: per-row exponents reach (64/2)*5*log2(e) = 231 log2, past the fp32 limit of 128. With all gates pinned, the measured NaN onset matches a chunk_size/2-row half-range exactly (cs64 onset between 2.75 and 2.80 vs 2.77 predicted; cs16 between 11.0 and 11.5 vs 11.09 predicted), so the cs/2 formula in gate_bound_is_safe needs no change. Map safe_gate=True to its documented -5 bound and apply gate_bound_is_safe on both gate paths, so safe_gate=True + chunk_size=64 falls back to Triton. RWKV7 keeps its cs64 route by asserting its architectural w clamp via lower_bound=-exp(-0.5). Sustained-pin stress tests cover the licensed bounds at cs16/32/64: the route stays finite and tracks the fp32 recurrence (at the -5 pin, dgk RMS err 0.12-0.17 vs 0.34-0.48 for the Triton reference, so parity is asserted against the baseline). Also correct the varlen layout cache comment: fla.utils.tensor_cache is not reusable there since it identity-keys every argument.
lower_bound landed between safe_gate and chunk_size, which breaks positional callers that pass chunk_size. Move it after cp_context in the public op and in the TileLang backend protocol methods; all in-tree callers use keywords.
A rank-local fallback to Triton would diverge the cross-rank collectives inside the CP forward and hang the job; surface the error instead.
|
Thanks for the careful review. Items 1–3 and the two minors are addressed in d24c7e9, fd7ee74, 43beb65 (pushed); item 4 needs a Hopper box, plan below. 1. safe_gate + cs64 overflow — confirmed and fixed. The verifier now maps 5. Off-by-one in 2. Positional break — 3. CP fallback — the forward now re-raises under 6. Varlen repro — generation rule added to the benchmark preamble: 8192 tokens flattened to 4. H800 cs64 — agree with the data (0.54–0.81×). Plan: gate cs64 to devices with measured wins (cc12x) in the verifier, one cc check next to the existing schedule arithmetic — unless a Hopper re-measurement justifies keeping sm_90. I'll run the Hopper pass first (machine access incoming) and follow up either way. |
…can kernel The rectangular-batch gate cumsum ran as a vectorized PyTorch chain (fp32 cast, pad, cumsum(dim=2), scale, zeros_like + shift), which profiled at ~2.3-2.8 ms/iter of direct_copy / scan_outer_dim / fill glue on H800 at h4096. A segmented fp32 scan kernel (SEG lane groups scan their segment serially in registers, then apply a per-channel exclusive prefix of segment sums) emits gi and the shifted ge in one pass at ~1.3 ms, ulp-level parity with the chain. chunk_local_cumsum gains output_ge=False: the disable_recompute backward at head dim 64 derives ge in-kernel, so the ge write (and its fp32 allocation) is now skipped there; both scan kernels take an OUTPUT_GE constexpr and the dead buffer aliases gi.
Python autograd.Function backward receives dense zero grads for every forward output slot. The ctx forward returns 13 save-for-backward intermediates alongside o/final_state, so every backward step materialized ~2.8GB of zeros (h4096, bf16) that the backward discards unread. set_materialize_grads(False) passes None instead; both backward entry points already treat a missing dht as None.
cs64 measured 0.68-1.05x vs Triton on sm_90 (H800 sweep, rect and varlen, both head dims, pre- and post-cumsum-fix): the BT=64 wy/wu UT-transform stages are latency-bound there and the fused h+o pass does not amortize at twice the chunk serial length. Schedulability was never the blocker on Hopper's 228KB optin — the route simply loses, so the verifier now rejects it below cc12 and the call falls back to Triton (1.00x). sm_120 keeps cs64 (validated 1.30-2.18x). The sm90-accepts-cs64 verifier test flips to assert the rejection, and the parity-suite skip helper mirrors the gate (renamed _cs64_launchable -> _cs64_routed).
The a-side centered operand uses the exclusive cumsum ge (ge[0]=0) against the inclusive gi[mid], so its worst positive exponent spans chunk_size/2+1 rows, not chunk_size/2. The old helper accepted bounds whose true exponent overflows fp32 (e.g. lower_bound=-10.39 at cs16 reaches 135 log2). Use the exact span and set the threshold to 124 = 128 log2 (fp32 limit) minus 4 log2 of headroom for the activation multiply; the documented safe_gate range [-5, 0) stays licensed at cs16/cs32 (65/123 log2) and rejected at cs64 (238 log2), and RWKV7's -0.61 clamp at cs64 (29 log2) is unaffected.
Measured 0.83-1.02x vs Triton on sm_90 (H800, D64, all varlen sizes, both disable_recompute policies): the BT=16 serial state pass does not amortize ragged chunks there, so varlen cs16 now falls back to Triton below cc12. cc12x keeps the route (1.08-1.23x at h2560/h4096 on the PRO 6000), as does rect cs16 on both architectures (1.06-1.35x).
…havior Fresh probes on sm_120 (every gate pinned at the -5 bound, cs32, h0+dht backward) show the Triton reference stays finite too; what actually breaks parity is its dgk diverging from the fp32 recurrence at an err ratio of ~0.5 vs TileLang's ~0.17. State that instead of the earlier unreproduced fp32-overflow claim.
|
Thanks for the detailed review — all four blockers and both smaller items are addressed in the current head. 1. Gate-bound hole. Both 2. Signature order. 3. CP fallback. Under 4. sm90 cs64. cs64 now routes only on cc12x+ (acceptance still implies launchability via shared schedule arithmetic); varlen cs16 measured the same and got the same gate:
Rejected configs fall back at 1.00×; the Impact section is corrected. Why it lags on Hopper (per-kernel attribution + IKET warp-level tracing): the deficit is confined to the two serial state-scan stages (fused h+o forward, stream backward), which walk T/64 chunks serially on grids of N·H·(V/BV) and N·H blocks — at h1024 D128 that's 256/64 blocks against 132 SMs, while Triton's split passes re-parallelize over chunks and V-blocks (stream 1.13 vs. 0.65 ms; h+o 0.74 vs. 0.53), even though our fused intra stage wins there (0.61 vs. 0.92 ms). IKET confirms it: per ~8 µs chunk chain, 54% of warp time is global staging vs. 37% for all mma phases, with no next-chunk prefetch and ~12% occupancy — prefetch/pipelining is the follow-up. Once the grid fills, the scans reach parity and the residual is the intra stage (2.96 vs. 2.24 ms), Triton's strongest. On cc12x the same structure holds up: the stream still wins (3.21 vs. 5.01 ms) and h+o is at parity — Triton's split passes are just as latency-bound and gain nothing from the extra 56 SMs. The total flips because Triton's sm_90a-tuned intra regresses 5.5× on sm_120 (2.2 → 12.3 ms) while ours stays flat (≤1.3×, tracking clocks): 0.76× deficit → 3.3× win, deciding the 1.88× total. Same ISA class on both GPUs (mma.sync m16n8k16, no wgmma/TMA) — scheduling, not tensor cores. The cumsum/stream fixes also moved the worst rect cell from 0.80× to 0.92× on H800; what remains is the structural underfill, hence the gate. Off-by-one. Varlen reproducibility. The harness is documented: 8192 tokens split into N sequences with lengths ∝ Accuracy. Pinned by the new tests: route parity holds TileLang to Triton's own err-ratio bars vs. the fp32 recurrence (0.007/0.008) across rect+varlen, bf16/fp16, fp32 gates, cs16/32/64, both Suite: |
Summary
A TileLang implementation of the DPLR chunk pipeline, registered as a
dispatch backend on
chunk_dplr_delta_rule(chunk_rwkv7is an alias).A verifier accepts a call only where the TileLang path is correct and
measurably faster; everything else falls back to the default Triton
implementation unchanged. On by default when
tilelangis present;FLA_TILELANG=0forces Triton. Context parallelism (cp_context) worksthrough FLA's existing boundary-state scheme.
1. Differences from the Triton implementation
Structure
hand the outputoin a single sweep, replacing Triton's separatefwd_h/fwd_opasses and their global-memory round trip of
h/v_new.dhscan together with itsconsumers (
dq/dk/dw/db/dv2/dgk_last/dh0);dhis nevermaterialized. Triton runs separate
bwd_dhu/dv/dqkwgkernels.gkatK = 64; atK = 128asegmented fp32 scan kernel replaces a PyTorch cumsum chain that
profiled at ~2.3–2.8 ms/iter of glue (h4096).
torch.library.custom_op+register_fake+register_autograd— opaque totorch.compile, no graph break.Triton uses
autograd.Functionunder@torch.compiler.disable.Dead zero-grad fills for the 13 saved intermediates are disabled via
set_materialize_grads(False)(~2.8 GB/step at h4096).disable_recomputekeeps the Triton semantics (save vs. recompute);both policies are supported and benchmarked.
Numerics — same conventions as Triton (log2 gates, mid-chunk
gate-centering for tensor-core operands, fp32 state accumulation, fp32
A_ab_inv), exceptA_ab/A_akare stored in fp16 (Triton: fp32),validated within the existing tolerances. Context parallelism uses the
exact same boundary-state exchange as the Triton CP path.
Accuracy vs. the fp32 per-token recurrence (max err ratio over o, final
state, and all input grads; B8 T256 H32,
initial_state+dhtbackward, measured on sm_120):
TileLang is at or below Triton's error in every config (up to ~1.8×
closer at cs32 bf16); the fp16
A_ab/A_akstorage costs nothingmeasurable (fp16 rows). Under sustained −5 gate saturation both stay
finite, but Triton's
dgkerr ratio is ~3× worse (0.51 vs. 0.17).Routing
safe_gate=True, or alower_boundthat fits the chunk size (GDN2's convention): theA-stage's mid-chunk-centered
exp2operands must stay in fp32 range.Unbounded calls stay on Triton, which handles arbitrary gates.
(one log line per reason). For
chunk_size=64it shares schedulearithmetic with the launcher, so acceptance implies launchability.
2. Benchmark results
bf16,
triton.testing.do_bench(warmup 10, rep 50),FLA_TILELANG=0vs.
=1through the dispatched public op, idle GPUs: RTX PRO 6000(sm_120, 99KB smem) and H800 (sm_90, 228KB smem).
dr=disable_recompute. † = verifier reroutes to Triton (~1.00× byconstruction). Headline cells were additionally re-measured with
CUDA-event wall timing and profiler attribution — same ratios. Varlen
harness: each row splits 8192 tokens into N sequences with lengths
∝
rand(N) + 0.35(seed = 17 + N),[1, 8192]layout,gkdrawn as-0.61·sigmoid(5x)in fp32; the timing script is the same do_benchharness for every table and is included in the PR discussion.
PRO 6000 (sm_120)
Rectangular, cs32, T=4096 (
hidden= H·D):Varlen, cs32, 8192 tokens total (
[1, 8192]layout; N = sequence count):H800 (sm_90)
Rectangular, cs32, T=4096:
Varlen, cs32, 8192 tokens total:
Reading the numbers:
below the fwd-only ratio: the forward carries the large wins, the
backward (bwd-only, derived as fb − fwd) is 1.0–1.5× at dr=0 on both
devices. At dr=1 on sm_90 the TileLang backward alone loses to
Triton's dr=1 backward (0.6–0.9× — Triton skips its recompute too),
and the fwd win carries the total; dr=0 is the default and the
recommended mode.
intra tensorcore kernel degrading on ragged varlen on sm_90 (14.4 of
17.7 ms over fwd+bwd at 2560/D128/N5 — two calls, ~7.2 ms each). On
sm_120 the same kernel is healthy (2.7 of 7.5 ms) and the fused
TileLang A-stage still wins ~1.9× fwd there.
(1.11–1.37× on sm_120); varlen cs16 measured 0.83–1.02× on sm_90 (no
win at any size) and is now rerouted to Triton below cc12 — cc12x
keeps it (1.08–1.23× at h2560/h4096), with the smallest cell (h1024,
N6) at 0.92–0.98× there; one rect D64 dr=1 cell (h2560, B8) is 0.96×
on sm_90. Grids below half the SMs reroute to Triton at ~1.00×.
launch overhead in the tvm_ffi adapter (dynamic output-shape
resolution per launch), fixed by caller-allocated outputs — not a
kernel effect.
does not reproduce on an idle GPU — it was measured while two sweep
shards shared one device, which inflated the Triton baseline; all
tables above were re-run one shard per GPU. The current maximum ratio
(5.96×, H800 varlen D128) is profiler-attributed above.
chunk_size=64 (PRO 6000 only)
Two footprint fixes (staging-pair aliasing in the A-backward,
global-
gkreads in the A-forward) made every BT=64 stage launchable in99KB, bit-exact against their predecessors. The routed sweep:
(h1024 D64 spot cells: 1.93–2.20× rect, 1.66–1.92× varlen.) On H800
every cs64 schedule launches but measured 0.68–1.05× vs. Triton, so the
verifier accepts cs64 only on cc12x+; below that it falls back at 1.00×.
Why cs64 loses on sm_90 (per-kernel attribution via
torch.profileron both GPUs at the same rect h4096/h1024 D64/D128 dr1 configs,
generated code inspected from the TileLang/Triton kernel caches):
fused h+o forward and the stream backward — which walk T/64 chunks
serially per block on grids of N·H·(V/BV) and N·H blocks. At h1024
D128 that is 256/64 blocks against 132 SMs, so the scans run long
underfilled chains while Triton's split passes (dhu + dqkwg + o)
re-parallelize over chunks and V-blocks (stream 1.13 ms vs. 0.65 ms
combined; h+o 0.74 ms vs. 0.53 ms — together most of the 0.92×
total, while the fused intra/A stage itself wins 0.61 ms vs.
0.92 ms). Once the grid fills, the scan stages reach parity
(stream 2.50 ms vs. 2.64 ms at h4096 D64 dr1) and the residual
delta is the intra/A stage (2.96 ms vs. 2.24 ms), where Triton's
sm_90a schedule is at its strongest.
event tool,
nvidia-cutlass-dsl) confirms the mechanism at 64blocks / 8 warps on 132 SMs (h1024 D64, cs64): per chunk iteration
(~8 µs of dependent phases), global staging dominates — input
loads 33% + output row stores 21% = 54% of warp time vs. 37% for
all mma phases (dv2 11% + state gemms 23% + dh update 4%), with
dgk reduction at 8%. The next chunk's loads are not prefetched
ahead of the dh dependency chain, and 1 block/SM leaves ~12% warp
occupancy to hide load latency. That points at next-chunk
prefetch/pipelining as the follow-up optimization, on all
architectures.
the same stage takes 2.2 ms on H800 vs. 12.3 ms on the PRO 6000 —
a 5.5× regression of the Triton reference itself on sm_120 (2.4×
for the backward intra). Our kernels' absolute times are flat across
architectures (≤1.3×, tracking the 1980 vs. 3090 MHz clocks). Both
backends emit the same instruction class on both GPUs (mma.sync
m16n8k16 + ldmatrix; no wgmma/TMA on either side), so the gap is
scheduling, not tensor-core ISA.
FLOPS, ~15–18% of peak bandwidth), which is why fusion pays once
the GPU is filled, and why sm_120 — where the reference's intra
path collapses — flips the same h4096 config from 0.89× to 1.88×.
stage the scans hold up there (stream 3.21 ms vs. Triton's combined
dh-side passes 5.01 ms = 1.56×; fused h+o at parity, 2.42 vs.
2.54 ms) — Triton's split passes are just as latency-bound and
gain nothing from the extra 56 SMs. What flips the total is the
intra/A stage turning a 0.76× deficit into a 3.3× win (3.72 ms
vs. 12.30 ms): the stages the design is weak in stay at or above
parity, and the stage it is strong in decides the total
(24.85 ms vs. 46.69 ms = 1.88×).
Correctness gates: route-parity vs. Triton within
assert_close0.008–0.01 on all outputs and gradients (rect + varlen, fwd + bwd,
fp16/bf16, fp32 gates, K=V∈{64,128}, chunk 16/32/64); parity vs. the
per-token fp32 recurrence at the same 0.007/0.008 ratios
test_dplr_deltaholds the Triton path to (worst observed 0.0051);exact 0.0000 match for bf16 + fp32 gates on rectangular inputs; the
2-GPU CP test against the naive recurrence. The full gate
(
tests/ops/test_dplr_tilelang.py+tests/ops/test_dplr_delta.py,99 tests) passes on both devices: 67 passed + 2 skipped in the backend
file, and
test_dplr_delta.pyapart from the oneD100-chunk_size64fp16 failure that reproduces identically on
main(pre-existing; headdim 100 never routes to TileLang). With cs64 gated to cc12x+, the sm_90
run skips the cs64 parity rows and exercises the rejection unit test
instead.
3. Current incompatibilities
The verifier rejects these and the call falls back to Triton (one log
line per reason):
safe_gate=Falsewithout a fittinglower_boundexp2operands in fp32 range; unbounded calls stay on Triton. RWKV7's clamp (w ∈ (−0.61, 0)) qualifies as either.gkdtype ∉ {fp16, bf16, fp32};k/v/a/bdtype ≠qdtypeK != Vchunk_size∉ {16, 32, 64}chunk_size16 + head dim 128; varlenchunk_size16 below cc12chunk_size64N·H·(V/BV) < SMs/2)cp_contextis supported (same boundary-state exchange and the samecontract as the Triton entry: no user
initial_state/output_final_stateunder CP, loud failure on a userdht).One stress-test note: under sustained saturation at the −5 gate bound
(every gate pinned, cs32) the Triton reference's
dgkdiverges from thetrue recurrence far beyond parity tolerance (measured err ratio ~0.5 vs.
TileLang's ~0.17 on sm_120; both stay finite), so stress rows pin ~1% of
gates. The
lower_boundkwarg is wired through both routes.Test plan
tests/ops/test_dplr_tilelang.py(new): verifier accept/reject perbranch (incl.
safe_gate=False,lower_boundlicensing, cs64 devicegates, small-grid guard), route parity vs. Triton (fwd + all grads +
dh0, rect + varlen, bf16/fp16, fp32 gates, chunk 16/32/64, both drpolicies), gate-stress rows, PyTorch-recurrence baseline rows, and
opcheck/torch.compilefullgraph /torch.utils.checkpointsmokes.
tests/ops/test_backends.py: nvcc/tilelang gating parametrizationextended to the new backend.
tests/context_parallel/test_cp_dplr.py: new 2-GPUtest_cp2_tilelang_routewith per-rank route assertion.test_dplr_delta.py,test_rwkv7.py,test_modeling_rwkv7.py,test_layer_cache_layer_idx.pypass (the 4channel_mixing_gradientsand 1D100-chunk_size64failuresreproduce identically on
main— pre-existing).Breaking changes
None. The change is additive: where the verifier does not accept a call,
the default Triton implementation is used unchanged.
Developed on
tilelang-dplr-backend; benchmarks on idle RTX PRO 6000(sm_120) and H800 (sm_90). Correctness gates pass on both.