Add GDN benchmarking

benchmark/linear_attention/Dockerfile

@@ -19,5 +19,12 @@ RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/sbs
 RUN git clone https://github.com/NVIDIA/cudnn-frontend.git
 RUN pip install -v /workspace/cudnn-frontend
 
-# Install the Cutlass DSL runtime (cuDNN FROST engines) and FLA
-RUN pip install nvidia-cutlass-dsl[cu13]==4.5.0 flash-linear-attention
+# Install the Cutlass DSL runtime (cuDNN FROST engines) and FLA.
+RUN pip install nvidia-cutlass-dsl[cu13]==4.7.0 apache-tvm-ffi flash-linear-attention
+
+# Install FlashQLA from source.
+RUN git clone https://github.com/QwenLM/FlashQLA.git
+RUN pip install -v /workspace/FlashQLA
+
+# Install the chart dependencies for plot_results.py
+RUN pip install pandas matplotlib seaborn

benchmark/linear_attention/README.md

@@ -57,6 +57,13 @@ python benchmark_single_linear_attention.py \
     --la_backend fla --variant gdn --data_type bfloat16 \
     --skip_ref --fwd_bwd
 
+# FlashQLA (TileLang) comparison point (gdn variant only)
+python benchmark_single_linear_attention.py \
+    --batch_size 1 --seqlen 8192 \
+    --num_q_heads 8 --num_kv_heads 64 --head_dim 128 \
+    --la_backend flash_qla --variant gdn --data_type bfloat16 \
+    --skip_ref --fwd_bwd
+
 # Recurrent state ports: seed with an initial state and request the final
 # state (its gradient feeds the backward pass)
 python benchmark_single_linear_attention.py \
@@ -76,6 +83,7 @@ Dropping `--skip_ref` validates the forward output against FLA (the same way the
 |---------|-------------|
 | `cudnn` | cuDNN (native, via the cuDNN Frontend torch custom ops) |
 | `fla`   | FLA (flash-linear-attention, Triton) |
+| `flash_qla` | FlashQLA (TileLang fused GDN kernels, `gdn` variant only) |
 
 The cuDNN backend routes through the pygraph engines: FROST (Cutlass DSL) on SM100-class devices, the cuTile engines elsewhere. Both passes run through autograd, exactly like a training step.
 
@@ -89,9 +97,11 @@ The cuDNN backend routes through the pygraph engines: FROST (Cutlass DSL) on SM1
 
 The benchmark runs `kda` and `gdn2` with the in-kernel q/k L2 normalization off (`use_qk_l2norm_in_kernel=False`) on every backend, for an apples-to-apples comparison.
 
+Recent `fla` releases dispatch `chunk_gated_delta_rule` to FlashQLA whenever `flash_qla` is importable; the benchmark sets `FLA_DISABLE_BACKEND_DISPATCH=1` (unless already set) so the `fla` backend always measures FLA's own Triton kernels and the two backends stay distinct.
+
 ## Notes
 
-- Head convention: `--num_q_heads` counts the query/key heads and `--num_kv_heads` counts the value heads; the gates, output, and recurrent state live at `max(num_q_heads, num_kv_heads)` heads. Both grouping directions are supported for `gdn`: grouped-value attention (`num_kv_heads > num_q_heads`, v-heads grouped over q-heads) and GQA (`num_q_heads > num_kv_heads`, q-heads grouped over v-heads, e.g. `--num_q_heads 64 --num_kv_heads 8`). The two counts must be equal or one a multiple of the other; `kda` and `gdn2` support the GVA direction only.
+- Head convention: `--num_q_heads` counts the query/key heads and `--num_kv_heads` counts the value heads; the gates, output, and recurrent state live at `max(num_q_heads, num_kv_heads)` heads. Both grouping directions are supported for `gdn`: grouped-value attention (`num_kv_heads > num_q_heads`, v-heads grouped over q-heads) and GQA (`num_q_heads > num_kv_heads`, q-heads grouped over v-heads, e.g. `--num_q_heads 64 --num_kv_heads 8`). The two counts must be equal or one a multiple of the other; `kda` and `gdn2` support the GVA direction only, and so does the `flash_qla` backend.
 - The cuDNN ops use the THD (token-packed) layout internally; the benchmark expresses the dense batch as `cu_seqlens = [0, T, 2T, ...]`.
 - `--initial_state` provides a per-sequence fp32 recurrent state (its gradient is produced in the backward pass); `--store_on` requests the per-sequence final state from the forward pass and feeds its gradient in the backward pass. Both are once-per-kernel I/O ports (one `[head_dim_qk, head_dim_vo]` tile per sequence per state head).
 - Performance is measured with the torch profiler (device time of the matched kernels), with a 256 MB L2 flush before each timed iteration and the median reported. TFLOPS use the chunked-BMM FLOPs model documented in the script's `flops()`.

benchmark/linear_attention/benchmark_single_linear_attention.py

@@ -197,6 +197,7 @@ def parse_args():
         help="Linear attention backend to use",
         choices=[
             "fla",
+            "flash_qla",
             "cudnn",
         ],
     )
@@ -249,7 +250,7 @@ def run_benchmark(
         head_dim_qk: Head dimension for Q/K (optional, for asymmetric)
         head_dim_vo: Head dimension for V/O (optional, for asymmetric)
         data_type: Data type ("bfloat16", "float16")
-        backend: Backend name ("cudnn", "fla")
+        backend: Backend name ("cudnn", "fla", "flash_qla")
         variant: Linear attention variant ("gdn", "kda", "gdn2")
         profile_pass: Which pass to profile ("fwd", "bwd", "both")
         num_iterations: Number of benchmark iterations
@@ -436,6 +437,11 @@ def run_benchmark(
         raise ValueError("gdn2 is forward only (the backward kernel is a stub); use --profile_pass fwd")
     if args.variant == "gdn2" and args.la_backend == "fla":
         raise ValueError("gdn2 is only supported with the 'cudnn' backend")
+    if args.la_backend == "flash_qla":
+        if args.variant != "gdn":
+            raise ValueError("flash_qla only supports the 'gdn' variant")
+        if num_q_heads > num_kv_heads:
+            raise ValueError("flash_qla does not support GQA (num_q_heads > num_kv_heads)")
 
     l2_flush_size_mb = 256
     l2_flush_size = l2_flush_size_mb * 1024 * 1024
@@ -507,9 +513,33 @@ def cudnn_linear_attention(query, key, value, gate, beta, write_gate, s0):
                     use_qk_l2norm_in_kernel=False,
                 )
 
+    if args.la_backend == "flash_qla":
+        attn_scale = head_dim_qk ** (-0.5)
+
+        from flash_qla import chunk_gated_delta_rule as fqla_chunk_gated_delta_rule
+
+        if args.verbose:
+            import flash_qla
+
+            print(f"[INFO] FlashQLA Version: {getattr(flash_qla, '__version__', 'unknown')}")
+
+        def flash_qla_linear_attention(query, key, value, gate, beta, write_gate, s0):
+            return fqla_chunk_gated_delta_rule(
+                query,
+                key,
+                value,
+                gate,
+                beta,
+                scale=attn_scale,
+                initial_state=s0,
+                output_final_state=args.store_on,
+            )
+
     if args.la_backend == "fla" or (not args.skip_ref):
         attn_scale = head_dim_qk ** (-0.5)
 
+        os.environ.setdefault("FLA_DISABLE_BACKEND_DISPATCH", "1")
+
         if args.variant == "gdn":
             from fla.ops.gated_delta_rule import chunk_gated_delta_rule
         elif args.variant == "kda":
@@ -549,21 +579,23 @@ def fla_linear_attention(query, key, value, gate, beta, write_gate, s0):
     def get_linear_attention_function(backend):
         if backend == "fla":
             return fla_linear_attention
+        elif backend == "flash_qla":
+            return flash_qla_linear_attention
         elif backend == "cudnn":
             return cudnn_linear_attention
         else:
             raise ValueError(f"Invalid backend: {backend}")
 
     # Util function for addressing different qkv formats for each backend
-    # (cudnn is THD [B*T, H, D]; fla is dense [B, T, H, D])
+    # (cudnn is THD [B*T, H, D]; fla and flash_qla are dense [B, T, H, D])
     def preprocess_qkv(query, key, value, backend):
         if backend == "cudnn":
             return (
                 query.reshape(batch_size * seqlen, *query.shape[2:]),
                 key.reshape(batch_size * seqlen, *key.shape[2:]),
                 value.reshape(batch_size * seqlen, *value.shape[2:]),
             )
-        elif backend == "fla":
+        elif backend in ("fla", "flash_qla"):
             return query, key, value
         else:
             raise ValueError(f"Invalid backend: {backend}")
@@ -575,7 +607,7 @@ def preprocess_gates(gate, beta, write_gate, backend):
                 beta.reshape(batch_size * seqlen, *beta.shape[2:]),
                 write_gate.reshape(batch_size * seqlen, *write_gate.shape[2:]) if write_gate is not None else None,
             )
-        elif backend == "fla":
+        elif backend in ("fla", "flash_qla"):
             return gate, beta, write_gate
         else:
             raise ValueError(f"Invalid backend: {backend}")
@@ -584,7 +616,7 @@ def preprocess_gates(gate, beta, write_gate, backend):
     def postprocess_o(output, backend):
         if backend == "cudnn":
             return output.reshape(batch_size, seqlen, num_o_heads, head_dim_vo)
-        elif backend == "fla":
+        elif backend in ("fla", "flash_qla"):
             return output
         else:
             raise ValueError(f"Invalid backend: {backend}")
@@ -800,8 +832,12 @@ def generate_gates(io_dtype):
                 query_ref = query.detach().reshape(batch_size, seqlen, num_q_heads, head_dim_qk)
                 key_ref = key.detach().reshape(batch_size, seqlen, num_q_heads, head_dim_qk)
                 value_ref = value.detach().reshape(batch_size, seqlen, num_kv_heads, head_dim_vo)
-                gate_ref = gate.detach().reshape(batch_size, seqlen, *gate.shape[1:])
-                beta_ref = beta.detach().reshape(batch_size, seqlen, *beta.shape[1:])
+                if args.la_backend == "cudnn":
+                    gate_ref = gate.detach().reshape(batch_size, seqlen, *gate.shape[1:])
+                    beta_ref = beta.detach().reshape(batch_size, seqlen, *beta.shape[1:])
+                else:
+                    gate_ref = gate.detach()
+                    beta_ref = beta.detach()
                 s0_ref = s0.detach() if s0 is not None else None
                 output_ref, _ = fla_linear_attention(query_ref, key_ref, value_ref, gate_ref, beta_ref, None, s0_ref)
 

benchmark/linear_attention/plot_results.py

@@ -0,0 +1,164 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Chart generation for linear attention benchmark results.
+
+Reads the CSV emitted by the shmoo runner (one --format_output line per
+(backend, batch, seqlen) case) and generates one comparison chart per batch
+size, with Forward and Backward TFLOPS panels side by side — same style as
+the SDPA training benchmark charts.
+
+    python plot_results.py results/gdn/b300/gdn_labench.csv \
+        --output-dir results/gdn/b300 --gpu-name B300 --cudnn-version 9.24.0
+"""
+
+import argparse
+from pathlib import Path
+
+import pandas as pd
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+# Backend display configuration; `order` fixes both bar grouping and legend order.
+BACKEND_CONFIG = {
+    "fla": {"name": "FLA (Triton)", "color": "#FF8C00", "order": 0},
+    "flash_qla": {"name": "FlashQLA (TileLang)", "color": "#6495ED", "order": 1},
+    "cudnn": {"name": "cuDNN", "color": "#76b900", "order": 2},
+}
+
+LABEL_FONT_SIZE = 10
+LEGEND_FONT_SIZE = 8
+TITLE_FONT_SIZE = 12
+BAR_LABEL_FONT_SIZE = 6
+
+CSV_COLUMNS = [
+    "case_tag",
+    "backend",
+    "variant",
+    "batch_size",
+    "seqlen",
+    "num_q_heads",
+    "num_kv_heads",
+    "head_dim",
+    "fwd_ms",
+    "bwd_ms",
+    "fwd_tflops",
+    "bwd_tflops",
+    "max_diff",
+    "num_iters",
+]
+
+
+def get_backend_display_name(backend: str, cudnn_version: str = None) -> str:
+    base_name = BACKEND_CONFIG.get(backend, {}).get("name", backend)
+    if backend == "cudnn" and cudnn_version:
+        base_name = f"{base_name} {cudnn_version}"
+    return base_name
+
+
+def generate_charts(df: pd.DataFrame, output_dir: Path, gpu_name: str = "", cudnn_version: str = None, variant: str = "gdn", batch_sizes: list = None) -> list:
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    df = df[df["variant"] == variant].copy()
+    if batch_sizes:
+        df = df[df["batch_size"].isin(batch_sizes)].copy()
+    if df.empty:
+        raise ValueError(f"No rows for variant {variant!r}")
+
+    df["backend_display"] = df["backend"].map(lambda b: get_backend_display_name(b, cudnn_version=cudnn_version))
+    df["backend_order"] = df["backend"].map(lambda b: BACKEND_CONFIG.get(b, {}).get("order", 99))
+
+    palette = {}
+    for _, row in df[["backend", "backend_display"]].drop_duplicates().iterrows():
+        palette[row["backend_display"]] = BACKEND_CONFIG.get(row["backend"], {}).get("color", "gray")
+
+    saved_paths = []
+    for batch_size in sorted(df["batch_size"].unique()):
+        sub = df[df["batch_size"] == batch_size].copy()
+        sub.sort_values(["seqlen", "backend_order"], inplace=True)
+        hue_order = list(sub.sort_values("backend_order")["backend_display"].drop_duplicates())
+
+        fwd_df = sub[sub["fwd_tflops"] > 0]
+        bwd_df = sub[sub["bwd_tflops"] > 0]
+        has_fwd = not fwd_df.empty
+        has_bwd = not bwd_df.empty
+
+        if has_fwd and has_bwd:
+            fig, (ax_fwd, ax_bwd) = plt.subplots(1, 2, figsize=(14, 6), dpi=150)
+        elif has_fwd:
+            fig, ax_fwd = plt.subplots(1, 1, figsize=(10, 6), dpi=150)
+            ax_bwd = None
+        else:
+            fig, ax_bwd = plt.subplots(1, 1, figsize=(10, 6), dpi=150)
+            ax_fwd = None
+
+        heads = sub["num_q_heads"].iloc[0]
+        head_dim = sub["head_dim"].iloc[0]
+        gpu_info = f" ({gpu_name})" if gpu_name else ""
+        fig.suptitle(
+            f"{variant.upper()} Linear Attention (BF16) — Batch = {batch_size}, Heads = {heads}, d = {head_dim}{gpu_info}",
+            fontsize=TITLE_FONT_SIZE,
+        )
+
+        for ax, pass_df, pass_name, y_col in (
+            (ax_fwd, fwd_df, "Forward", "fwd_tflops"),
+            (ax_bwd, bwd_df, "Backward", "bwd_tflops"),
+        ):
+            if ax is None or pass_df.empty:
+                continue
+            sns.barplot(
+                data=pass_df,
+                x="seqlen",
+                y=y_col,
+                hue="backend_display",
+                hue_order=hue_order,
+                ax=ax,
+                palette=palette,
+                edgecolor="black",
+                linewidth=0.5,
+                errorbar=None,
+            )
+            ax.set_xlabel("Sequence Length", fontsize=LABEL_FONT_SIZE)
+            ax.set_ylabel("TFLOPS", fontsize=LABEL_FONT_SIZE)
+            ax.set_title(pass_name, fontsize=TITLE_FONT_SIZE)
+            ax.legend(title="Backend", fontsize=LEGEND_FONT_SIZE)
+            ax.tick_params(axis="x", rotation=45)
+            for container in ax.containers:
+                ax.bar_label(container, fmt="%.0f", fontsize=BAR_LABEL_FONT_SIZE)
+
+        plt.tight_layout()
+        output_path = output_dir / f"{variant}_b{batch_size}.png"
+        plt.savefig(output_path, dpi=150, bbox_inches="tight")
+        plt.close()
+        saved_paths.append(output_path)
+        print(f"Chart saved to {output_path}")
+
+    return saved_paths
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("csv", type=Path, help="Results CSV from the shmoo runner")
+    parser.add_argument("--output-dir", type=Path, default=None, help="Output directory (default: alongside the CSV)")
+    parser.add_argument("--gpu-name", default="", help="GPU name for the chart title")
+    parser.add_argument("--cudnn-version", default=None, help="cuDNN backend version for the legend (e.g. 9.24.0)")
+    parser.add_argument("--variant", default="gdn", help="Linear attention variant to plot")
+    parser.add_argument("--batch-sizes", default=None, help="Comma-separated batch sizes to plot (default: all in the CSV)")
+    args = parser.parse_args()
+    batch_sizes = [int(b) for b in args.batch_sizes.split(",")] if args.batch_sizes else None
+
+    df = pd.read_csv(args.csv)
+    missing = [c for c in CSV_COLUMNS if c not in df.columns]
+    if missing:
+        raise ValueError(f"CSV is missing expected columns: {missing}")
+
+    output_dir = args.output_dir if args.output_dir is not None else args.csv.parent
+    generate_charts(df, output_dir, gpu_name=args.gpu_name, cudnn_version=args.cudnn_version, variant=args.variant, batch_sizes=batch_sizes)
+
+
+if __name__ == "__main__":
+    main()

benchmark/linear_attention/results/gdn/b300/gdn_20260806.csv

@@ -0,0 +1,16 @@
+case_tag,backend,variant,batch_size,seqlen,num_q_heads,num_kv_heads,head_dim,fwd_ms,bwd_ms,fwd_tflops,bwd_tflops,max_diff,num_iters
+fla_B1_T2048,fla,gdn,1,2048,64,64,128,0.217,0.649,89,89,0.000000,20
+flash_qla_B1_T2048,flash_qla,gdn,1,2048,64,64,128,0.102,0.352,190,165,0.000000,20
+cudnn_B1_T2048,cudnn,gdn,1,2048,64,64,128,0.098,0.279,196,208,0.000000,20
+fla_B1_T4096,fla,gdn,1,4096,64,64,128,0.416,1.260,93,92,0.000000,20
+flash_qla_B1_T4096,flash_qla,gdn,1,4096,64,64,128,0.186,0.677,208,171,0.000000,20
+cudnn_B1_T4096,cudnn,gdn,1,4096,64,64,128,0.129,0.428,301,271,0.000000,20
+fla_B1_T8192,fla,gdn,1,8192,64,64,128,0.796,2.490,97,93,0.000000,20
+flash_qla_B1_T8192,flash_qla,gdn,1,8192,64,64,128,0.347,1.342,223,173,0.000000,20
+cudnn_B1_T8192,cudnn,gdn,1,8192,64,64,128,0.187,0.721,414,321,0.000000,20
+fla_B1_T16384,fla,gdn,1,16384,64,64,128,1.567,4.957,99,94,0.000000,20
+flash_qla_B1_T16384,flash_qla,gdn,1,16384,64,64,128,0.674,2.648,229,175,0.000000,20
+cudnn_B1_T16384,cudnn,gdn,1,16384,64,64,128,0.305,1.309,506,354,0.000000,20
+fla_B1_T32768,fla,gdn,1,32768,64,64,128,3.119,9.976,99,93,0.000000,20
+flash_qla_B1_T32768,flash_qla,gdn,1,32768,64,64,128,1.329,5.294,233,175,0.000000,20
+cudnn_B1_T32768,cudnn,gdn,1,32768,64,64,128,0.529,2.494,585,372,0.000000,20