Linear-attention examples and Helion-vs-FLA benchmark (ported from #1941)#2966
Linear-attention examples and Helion-vs-FLA benchmark (ported from #1941)#2966tarinduj wants to merge 37 commits into
Conversation
73b4666 to
4dbc365
Compare
There was a problem hiding this comment.
I see a lot of duplication between the accuracy(), test(), and benchmark() functions across the examples, can we restructure / deduplicate while preserving the semantics?
An better structure could be something like:
- For the helion kernels: make sure they have the same interface as
FLA, i.e. for everyfla_kernel there should be ahelion_kernel - We introduce a simple
LinearAttentionKernelVariantclass, which includes flags/enums that narrows down to a specific variant (e.g. vanilla / kda / rwkv) - We introduce kernel getters e.g.
get_helion_fwd_kernel,get_fla_fwd_kernel,get_helion_bwd_dkwhich takes aLinearAttentionKernelVariantand returns a kernel function - We only have one copy of
accuracy,test,benchmark, which accepts aLinearAttentionKernelVariantas arg, and then fetches the kernels and run them to compare accuracy / perf
| try: | ||
| from fla.ops.linear_attn import ( # pyrefly: ignore[missing-import] | ||
| chunk_linear_attn as _fla_fn, | ||
| ) | ||
|
|
||
| fla_chunk_linear_attn: Any = _fla_fn | ||
|
|
||
| o_fla = fla_chunk_linear_attn( | ||
| _htf(q), | ||
| _htf(k), | ||
| _htf(v), | ||
| scale=scale, | ||
| normalize=False, | ||
| ) | ||
| if isinstance(o_fla, tuple): | ||
| o_fla = o_fla[0] | ||
| o_fla_hf = o_fla.transpose(1, 2).contiguous() | ||
| fla_err = _rel_error(out, o_fla_hf) | ||
| print( | ||
| f" fwd vs FLA: {fla_err:.4e} {'PASS' if fla_err < 0.02 else 'FAIL'}" | ||
| ) | ||
| _has_fla = True | ||
| except ImportError: | ||
| warnings.warn("fla not installed, skipping FLA comparisons", stacklevel=1) | ||
| _has_fla = False |
There was a problem hiding this comment.
this section looks duplicated across different examples, with the only distinction being the fla_chunk_linear_attn which differs from example to example. Can we try to deduplicate?
| grad_out = torch.randn(B, H, T, DV, device=DEVICE, dtype=DTYPE) | ||
| q1 = q.clone().requires_grad_(True) | ||
| k1 = k.clone().requires_grad_(True) | ||
| v1 = v.clone().requires_grad_(True) | ||
| o1 = chunked_linear_attn(q1 * scale, k1, v1, g, C=C) | ||
| o1.backward(grad_out) | ||
|
|
||
| q2 = q.clone().requires_grad_(True) | ||
| k2 = k.clone().requires_grad_(True) | ||
| v2 = v.clone().requires_grad_(True) | ||
| o2 = chunked_linear_attn_reference(q2 * scale, k2, v2, g, C=C) | ||
| o2.backward(grad_out) |
There was a problem hiding this comment.
Same here, this also looks heavily duplicated, with the only distinction being which function to call as the kernel as the kernel / reference
dcd9b1f to
cf0335c
Compare
03cb428 to
808d4f0
Compare
c3e12b4 to
a1ed2a3
Compare
| | simple_gla | ok/ok | ok/ok | ok/ok | ok/ok | ok/ok | ok/ok | | ||
| | retention | ok/ok | ok/ok | ok/ok | ok/ok | ok/ok | ok/ok | | ||
| | full_gla | ok/ok | ok/ok | ok/ok | ok/ok | ok/ok | ok/ok | | ||
| | delta_rule | ok/ok | ok/ok | ok/ok | ok/ok | ok/FAIL | ok/ok | |
There was a problem hiding this comment.
Can you provide some context/analysis on what's happening over these two FAIL cases?
There was a problem hiding this comment.
Here the reference kernel OOMs on big shapes so so there's nothing to compare against. I will add a more specific label like: REF-ERR instead of FAIL.
| return chunked_linear_attn_reference(i.q * i.scale, i.k, i.v, i.g, beta=i.beta, C=C) | ||
|
|
||
|
|
||
| VARIANT = LinearAttentionVariant( |
There was a problem hiding this comment.
This is quite helpful, but it is somewhat different from what I was alluding to in this comment.
What I was meant is for the LinearAttentionVariant class to specify a variant of a kernel, instead of a variant of an ExampleHarness (which is also helpful). What I had in mind was to provide something like:
delta_rule_variant = LinearAttentionVariant(..) # I'll let you decide which params go here, could even be just an Enum
delta_rule_fwd_kernel = linear_attention_engine.get_fwd_kernel(delta_rule_variant)
result = delta_rule_fwd_kernel(q, k, v, ...)
So this would be a convenient method for users to obtain Helion kernels of any variant, without having to manually write code like:
def _helion_fwd(i: Inputs, C: int) -> torch.Tensor:
return chunked_linear_attn(i.q * i.scale, i.k, i.v, i.g, beta=i.beta, C=C)
for each variant.
Notice that, one such "user" is exactly these example runner files.
I expect that once we have that mechanism, we can rename your existing LinearAttentionVariant class into a LinearAttentionExampleHarness class, and fold even more logic into that class, so that each of these example files is less than 50 lines of code
There was a problem hiding this comment.
Makes sense. I will do another refactor.
| B, H, T, D, DV = 2, 4, 128, 32, 32 | ||
| C = 32 | ||
| DTYPE = torch.bfloat16 | ||
| BENCH_CONFIGS = [(1, 32, 2048, 128, 128), (1, 32, 4096, 128, 128)] | ||
| BENCH_C = 64 |
There was a problem hiding this comment.
Looks like all of these are duplicated and be folded into a common harness? (see this comment)
There was a problem hiding this comment.
Was going to move this into the harness, but some examples use DV=16 (eg: vanilla) while others use DV=32. I wasn't sure if we wanted to test variants on different shapes. But yes, making them uniform and folding into the harness is cleaner.
There was a problem hiding this comment.
What's the motivation for using different DV values for testing different variants?
There was a problem hiding this comment.
This is an artifact from the original PR I didn't change. I think testing on the same DV value is fine.
| ## Summary | ||
|
|
||
| `chunk_bwd_dh_diag_fused` and similar kernels using `hl.grid(N)` where N is | ||
| a dynamic (non-specialized) value produce an illegal memory access (IMA) when: |
There was a problem hiding this comment.
is this bug still happening?
There was a problem hiding this comment.
I can't reproduce it anymore. I think Markus hit this on an H200, but I am on a B200 now (some FLA kernels weren't supported on Hopper). All the benchmarking CIs were written for B200 as well, so I could drop this from the PR.
There was a problem hiding this comment.
some FLA kernels weren't supported on Hopper
Which ones are not supported on Hopper, and which ones are? For the ones that are, I think we should have CIs and benchmarks for them on Helion as well.
There was a problem hiding this comment.
FLA had a miscompilation for this kernel. But also looks like it was fixed in Triton 3.7.1. This affects simple_gla and gated_delta_rule.
| @@ -0,0 +1,284 @@ | |||
| """ | |||
| Mamba-2 SSD Example | |||
There was a problem hiding this comment.
Can you share some context on the relationship/differences between these mamba2 kernels versus the ones we have in the repo?
There was a problem hiding this comment.
The upstream ones are standalone kernels that reimplement Mamba2's kernels directly, but forward-only. This PR instead runs Mamba2 through the shared engine: same kernel/ recurrence as simple_gla (scalar decay, no correction), just with Mamba2's inputs, but it has forward and backward. FLA has no Mamba-2, so it isn't in the benchmark CI either. Since it's redundant with the upstream examples and can't be benched against FLA, I could drop it from this PR as well.
| | B2_T16384_H16_D128 | ok/ok | 0.222ms | 0.317ms | 142.8% | 1.107ms | 2.155ms | 194.7% | | ||
| | B4_T2048_H16_D128 | ok/ok | 0.106ms | 0.091ms | 85.8% | 0.948ms | 1.039ms | 109.6% | | ||
| | B4_T4096_H64_D128 | ok/ok | 0.428ms | 0.576ms | 134.6% | 2.002ms | 3.753ms | 187.5% | | ||
| | B8_T2048_H32_D256 | ok/ok | 0.658ms | 0.803ms | 122.0% | 3.007ms | 62.303ms | 2071.9% | |
There was a problem hiding this comment.
Any idea what's special about this shape? 2000% is a lot! While most others are in the 100% to 200% region.
(not high-pri to resolve for now, as the top priority is to get kernels + tests + benchmarks landed first, but curious to know if you have any ideas)
There was a problem hiding this comment.
I think this is some Triton + B200 artifact on the FLA kernel. The numbers were not this crazy on H100 (see here).
a1ed2a3 to
6e22c67
Compare
|
|
||
|
|
||
| def run_test( | ||
| v: LinearAttentionHarness, |
There was a problem hiding this comment.
tiny nit, this perhaps shouldn't be named v anymore?
|
|
||
|
|
||
| @dataclass | ||
| class LinearAttentionHarness: |
There was a problem hiding this comment.
tiny nit, lets rename this to LinearAttentionExampleHarness just to make it clear that this is for running the example, i.e. not essential to getting the kernel
| decay: DecayType = DecayType.NONE | ||
| correction: bool = False | ||
|
|
||
| def get_fwd_kernel(self) -> Callable[..., torch.Tensor]: |
There was a problem hiding this comment.
nit, Can the input types be specified?
| def _fla_fwd(i: Inputs, scale: float) -> torch.Tensor: | ||
| o, _ = _fla_chunk_delta_rule(i.q, i.k, i.v, i.beta, scale=scale) | ||
| return o |
There was a problem hiding this comment.
instead of having to declare this, can we have FLA kernel getters given an LinearAttentionVariant as well?
also btw I think it would be cleaner for the getters to be module-level functions rather than class methods. So we could have a module-level get_helion_fwd_kernel(variant: LinearAttentionVariant), as well as get_fla_fwd_kernel(variant: LinearAttentionVariant)
There was a problem hiding this comment.
Let me see, should be doable.
| decay: DecayType = DecayType.NONE | ||
| correction: bool = False | ||
|
|
||
| def get_fwd_kernel(self) -> Callable[..., torch.Tensor]: |
There was a problem hiding this comment.
also I wonder why the bwd kernels aren't here?
There was a problem hiding this comment.
I didn't add a separate bwd kernel since it's always fwd().backward(grad_out); like:
def helion_fwd(self, i: Inputs, C: int) -> torch.Tensor:
fwd = self.variant.get_fwd_kernel()
return fwd(i.q, i.k, i.v, i.g, i.beta, C=C, scale=i.scale)
def helion_fb(self, i: Inputs, grad_out: torch.Tensor, C: int) -> None:
self.helion_fwd(i, C).backward(grad_out)
| @@ -0,0 +1,95 @@ | |||
| # Linear Attention Kernels: Helion vs FLA (ported from #1941) | |||
There was a problem hiding this comment.
nit, lets move these as part of the summary function instead of an md file to be commited
| @@ -0,0 +1,13 @@ | |||
| # Linear Attention Kernels: Accuracy | |||
There was a problem hiding this comment.
nit, lets move these as part of the summary function instead of an md file to be commited
6e22c67 to
03501d9
Compare
| # Module API consumed by run_linattn.py: test / benchmark / accuracy. | ||
| test = HARNESS.test | ||
| benchmark = HARNESS.benchmark | ||
| accuracy = HARNESS.accuracy |
There was a problem hiding this comment.
nit, don't really these tmp vars, lets just directly call harness.test, test.benchmark
| print("All tests passed.") | ||
|
|
||
|
|
||
| def benchmark() -> None: |
There was a problem hiding this comment.
For my own understanding, why doesn't this mamba2_ssd example share the sample variant-harness pattern that the rest of the examples use?
There was a problem hiding this comment.
The harness's comparison path is built around FLA, but FLA doesn't have mamba2. So, I left mamba2 out of the harness.
| "b200_cute": "kernels_cute", | ||
| "mi350x": "kernels", | ||
| "tpu": "kernels_tpu", | ||
| "b200_fla": "kernels_linattn", |
There was a problem hiding this comment.
for my own understanding:
- Why do we use "fla" as part of the name? isn't the main point to test Helion rather than the fla library?
- Are we also running these on h100? (if not, we should)
There was a problem hiding this comment.
Additionally, why does linear attention need its own special entry here? Why can't it be part of the existing "h100" and "b200" dashboards?
There was a problem hiding this comment.
I wasn't actually sure what's the best way to handle this.
- Regular b200 runs benchmark.yml with tritonbench as the baseline.
- b200_fla runs benchmark_linattn.yml with FLA as the baseline, and our linear-attention variants aren't tritonbench ops.
Having this new b200_fla lane felt like the least invasive way to add the benchmarks. I did consider merging the results into the b200 column, but that felt more invasive.
There was a problem hiding this comment.
Oh, and I will add h100 support as well.
| @@ -0,0 +1,160 @@ | |||
| name: Benchmark Linear Attention | |||
There was a problem hiding this comment.
Why do we need a new workflow here, instead of reusing existing h100 and b200 workflows?
There was a problem hiding this comment.
I might be able to parameterize the existing workflow to run what I want. I will take a look again.
There was a problem hiding this comment.
FLA kernels has this flag return_final_state: bool = False (set to False by default), if it's true it will return a tuple. I am explicitly casting so the type checker doesn't complain.
There was a problem hiding this comment.
Its not very idea that we have both HELION_FWD and a def helion_fwd. I'd recommend not doing the cast and the wrapper definition. Instead, lets just add assert isintance(out, torch.Tensor) after the kernel invocation to make the type checker happy
There was a problem hiding this comment.
Oops yeah, sorry, I read the comment on my email and replied to the wrong thread here. I will fix that.
03501d9 to
aa00b0a
Compare
Adds a chunked linear-attention engine (examples/linear) and the variants built on it: vanilla, simple_gla, retention, full_gla, delta_rule, gated_delta_rule, kda, rwkv6, mamba2_ssd.
Convert the engine's kernels from aot_kernel to helion.kernel so they autotune through the standard path.
Add a scale arg to the forward and backward so the no-decay path passes q unscaled and folds scale into the output and gradient kernels, similar to FLA. Defaults to 1.0 so the decay variants that pre-scale q are unaffected. Fuse the cross/state adds into the matmul accumulators in chunk_fwd_o_helion and chunk_bwd_dqk_helion, and run the chunk_bwd_dh_diag_fused matmul as a bf16 hl.dot instead of an fp32 torch.mm, similar to FLA.
Reran all eight variants at full autotune across the six benchmark shapes.
Add linear_attention_harness with LinearAttentionVariant plus run_test / run_benchmark / run_accuracy, and move example_vanilla_linear_attn onto it.
The recurrent reference now floats g in test() as well as accuracy(); the fixed per-head decay (~0.97-0.9997) needs fp32 to stay accurate over long T.
Add make_kda_inputs alongside the other input generators; kda was the only variant building inputs inline.
example_rwkv6 is not fully implemented to match FLA's chunk_rwkv6 yet. Drop it until it does.
…hods Inline each variant's input construction into its own _make_inputs and drop the make_*_inputs helpers from linear_attention_utils (make_mamba2_inputs stays; mamba2 is off-harness). Make test/benchmark/accuracy methods on LinearAttentionVariant with the shapes as fields.
run_accuracy runs the Helion and reference calls in separate try blocks, so a kernel error, a reference error, and an over-tolerance mismatch are no longer conflated into one bool.
Add DecayType + LinearAttentionVariant to linear_attention_engine, with a get_fwd_kernel() method that returns a Helion forward kernel for that variant. Rename the harness's LinearAttentionVariant to LinearAttentionHarness, which now takes a variant and builds helion_fwd, helion_fb, fla_fb, and the two fp32 references. Each example names its variant and supplies only make_inputs and the FLA forward.
Add linear_attention_fla with get_fla_fwd_kernel(variant), an enum-keyed lookup of each variant's FLA op, mirroring get_helion_fwd_kernel.
Add per-variant make_*_inputs to linear_attention_harness and have LinearAttentionExampleHarness resolve make_inputs (and the title) from its variant in __post_init__. Each example now just names its variant and builds the harness
Fold the linear-attention benchmark path into the shared benchmark workflow and publish its results under the existing B200 dashboard column. Label linattn rows as FLA-baseline rows, exclude them from eager-baseline speedup aggregates, and drop unsupported rwkv6 entries from the nightly linattn default.
…ark snapshots and have run_linattn call each example's HARNESS directly
aa00b0a to
2665c45
Compare
| return cast( | ||
| "torch.Tensor", HELION_FWD(qi, ki, vi, gi, C=BENCH_C, scale=scale) | ||
| ) |
There was a problem hiding this comment.
what's happening here? I'd expect the original kernel returned by get_helion_fwd_kernel to already return a torch.Tensor hmm. What does it return instead?
2af6e74 to
726e561
Compare
726e561 to
31cc574
Compare
Picks up the linear-attention examples from #1941 (left unfinished) and adds correctness checks and a benchmark against flash-linear-attention (FLA).
What's here
examples/linear) and the variants on it: vanilla, simple_gla, retention, full_gla, delta_rule, gated_delta_rule, kda, rwkv6, mamba2_ssd. Each has a correctness test vs a pure PyTorch reference.benchmarks/run_linattn.py) that times each variant's forward and forward+backward against FLA (mamba2_ssd excluded for now). These kernels have no TritonBench operator, so FLA is the baseline and speedup is Helion vs FLA.accuracy()per variant) vs the fp32 reference, emitted ashelion_accuracy. Accuracy is a separate metric; timings are reported regardless of whether it passed.examples/linear/BENCHMARKS.md) with the B200 per-variant fwd and fwd+bwd numbers vs FLA.b200_flaplatform) so the forward and-bwdrows render alongside the existing platforms.Notes
benchmark_linattn.yml.