Skip to content

Linear-attention examples and Helion-vs-FLA benchmark (ported from #1941)#2966

Draft
tarinduj wants to merge 37 commits into
mainfrom
tarindu/markus-linear
Draft

Linear-attention examples and Helion-vs-FLA benchmark (ported from #1941)#2966
tarinduj wants to merge 37 commits into
mainfrom
tarindu/markus-linear

Conversation

@tarinduj

@tarinduj tarinduj commented Jun 29, 2026

Copy link
Copy Markdown
Contributor

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

  • A chunked linear-attention engine (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.
  • A benchmark runner (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.
  • A per-shape accuracy check (accuracy() per variant) vs the fp32 reference, emitted as helion_accuracy. Accuracy is a separate metric; timings are reported regardless of whether it passed.
  • A results writeup (examples/linear/BENCHMARKS.md) with the B200 per-variant fwd and fwd+bwd numbers vs FLA.
  • Dashboard + CI wiring (b200_fla platform) so the forward and -bwd rows render alongside the existing platforms.

Notes

  • Runs on B200 via benchmark_linattn.yml.
  • Kernels are not yet perf-tuned; this lands the examples, correctness, and the benchmark harness.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jun 29, 2026
@tarinduj tarinduj force-pushed the tarindu/markus-linear branch from 73b4666 to 4dbc365 Compare June 30, 2026 00:20
@tarinduj tarinduj requested a review from AmesingFlank June 30, 2026 00:21

@AmesingFlank AmesingFlank left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 every fla_ kernel there should be a helion_ kernel
  • We introduce a simple LinearAttentionKernelVariant class, 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_dk which takes a LinearAttentionKernelVariant and returns a kernel function
  • We only have one copy of accuracy, test, benchmark, which accepts a LinearAttentionKernelVariant as arg, and then fetches the kernels and run them to compare accuracy / perf

Comment on lines +56 to +80
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Comment on lines +83 to +94
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here, this also looks heavily duplicated, with the only distinction being which function to call as the kernel as the kernel / reference

@tarinduj tarinduj force-pushed the tarinduj/stack/1 branch 3 times, most recently from dcd9b1f to cf0335c Compare June 30, 2026 17:50
@tarinduj tarinduj changed the title Linear-attention examples and Helion-vs-FLA benchmark Linear-attention examples and Helion-vs-FLA benchmark (ported from #1941) Jun 30, 2026
@tarinduj tarinduj force-pushed the tarindu/markus-linear branch from 03cb428 to 808d4f0 Compare June 30, 2026 17:59
@tarinduj tarinduj changed the base branch from tarinduj/stack/1 to main July 6, 2026 21:10
@tarinduj tarinduj force-pushed the tarindu/markus-linear branch from c3e12b4 to a1ed2a3 Compare July 6, 2026 21:10
Comment thread examples/linear/ACCURACY.md Outdated
| 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 |

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you provide some context/analysis on what's happening over these two FAIL cases?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread examples/linear/example_delta_rule.py Outdated
return chunked_linear_attn_reference(i.q * i.scale, i.k, i.v, i.g, beta=i.beta, C=C)


VARIANT = LinearAttentionVariant(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense. I will do another refactor.

Comment thread examples/linear/example_delta_rule.py Outdated
Comment on lines +33 to +37
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like all of these are duplicated and be folded into a common harness? (see this comment)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the motivation for using different DV values for testing different variants?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this bug still happening?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you share some context on the relationship/differences between these mamba2 kernels versus the ones we have in the repo?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread examples/linear/BENCHMARKS.md Outdated
| 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% |

@AmesingFlank AmesingFlank Jul 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is some Triton + B200 artifact on the FLA kernel. The numbers were not this crazy on H100 (see here).

@tarinduj tarinduj force-pushed the tarindu/markus-linear branch from a1ed2a3 to 6e22c67 Compare July 7, 2026 20:48


def run_test(
v: LinearAttentionHarness,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

tiny nit, this perhaps shouldn't be named v anymore?



@dataclass
class LinearAttentionHarness:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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]:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit, Can the input types be specified?

Comment thread examples/linear/example_delta_rule.py Outdated
Comment on lines +55 to +57
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let me see, should be doable.

decay: DecayType = DecayType.NONE
correction: bool = False

def get_fwd_kernel(self) -> Callable[..., torch.Tensor]:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

also I wonder why the bwd kernels aren't here?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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) 

Comment thread examples/linear/BENCHMARKS.md Outdated
@@ -0,0 +1,95 @@
# Linear Attention Kernels: Helion vs FLA (ported from #1941)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit, lets move these as part of the summary function instead of an md file to be commited

Comment thread examples/linear/ACCURACY.md Outdated
@@ -0,0 +1,13 @@
# Linear Attention Kernels: Accuracy

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit, lets move these as part of the summary function instead of an md file to be commited

@tarinduj tarinduj force-pushed the tarindu/markus-linear branch from 6e22c67 to 03501d9 Compare July 8, 2026 19:39
Comment thread examples/linear/example_delta_rule.py Outdated
Comment on lines +16 to +19
# Module API consumed by run_linattn.py: test / benchmark / accuracy.
test = HARNESS.test
benchmark = HARNESS.benchmark
accuracy = HARNESS.accuracy

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit, don't really these tmp vars, lets just directly call harness.test, test.benchmark

print("All tests passed.")


def benchmark() -> None:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For my own understanding, why doesn't this mamba2_ssd example share the sample variant-harness pattern that the rest of the examples use?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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",

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Additionally, why does linear attention need its own special entry here? Why can't it be part of the existing "h100" and "b200" dashboards?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, and I will add h100 support as well.

Comment thread .github/workflows/benchmark_linattn.yml Outdated
@@ -0,0 +1,160 @@
name: Benchmark Linear Attention

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need a new workflow here, instead of reusing existing h100 and b200 workflows?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I might be able to parameterize the existing workflow to run what I want. I will take a look again.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm assuming this is a reply to this comment?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oops yeah, sorry, I read the comment on my email and replied to the wrong thread here. I will fix that.

@tarinduj tarinduj force-pushed the tarindu/markus-linear branch from 03501d9 to aa00b0a Compare July 8, 2026 23:27
tarinduj added 2 commits July 9, 2026 18:49
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.
tarinduj added 24 commits July 9, 2026 18:49
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
@tarinduj tarinduj force-pushed the tarindu/markus-linear branch from aa00b0a to 2665c45 Compare July 9, 2026 18:49
Comment thread examples/linear/example_mamba2_ssd.py Outdated
Comment on lines +244 to +246
return cast(
"torch.Tensor", HELION_FWD(qi, ki, vi, gi, C=BENCH_C, scale=scale)
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

@tarinduj tarinduj force-pushed the tarindu/markus-linear branch from 2af6e74 to 726e561 Compare July 9, 2026 22:08
@tarinduj tarinduj force-pushed the tarindu/markus-linear branch from 726e561 to 31cc574 Compare July 9, 2026 23:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants