Widen forward head-dim support on Blackwell (SM100)

[Costa-SM]

Hi folks.

I really wanted to try FlashQLA out with some different head dimension values (e.g. GDN-hybrid expand_v=2.0 layout w/ head_dim_k=64, head_dim_v=128) and I noticed that although the chunk gated_delta_rule kernels are parameterized on the head dims (DK/DV), the Python entry points assert head_dim == 128. So here's a PR to widen the accepted head_dim values during forward passes on Blackwell GPUs :D

I ran a parity sweep against the fp64 reference over the full config matrix (GQA, varlen, non-null h0, both state layouts) to get the head dims that are actually correct on the forward/inference path. During this I discovered that a few head_dim_k values (96 / 160 / 192) run without erroring and return incorrect results. So this PR also swaps the bare assert for a real guard that rejects these values.

Measurements

The parity sweep demonstrated that only selected values of head_dim_k work while head_dim_v basically accepts anything you throw at it. The parity checks for head_dim_k were run for values {16, 32, 48, 64, 80, 96, 128, 160, 192, 256} using head_dim_v=128. For head_dim_v, every multiple of 16 in [32, 256] was tested and is correct.

head_dim_k	Result	Notes
64, 128	✅ correct	matches the fp64 ref across the full config matrix
16, 48, 80, 256	🛑 raises	no valid MMA / shared-memory config
96, 160, 192	silently wrong	runs with no error, rel-err ~0.4
32	silently wrong (vk only)	correct for the kv layout, wrong for vk

So the final guards are head_dim_k==64 or head_dim_k==128 and head_dim_v % 16 == 0 and 32 <= head_dim_v <= 256.

Performance

In practice, we could still run head_dim_k=64 before by zero-padding it to 128. But since this PR adds support for 64, below is a comparison between the two. It was measured on B200/SM100, H=12, head_dim_v=128, causal, via tilelang.profiler.do_bench (cudagraph backend):

batch × seqlen	pad K→128	native K=64	saved
64 × 4096	1602 µs	1396 µs	~13%
32 × 8192	1596 µs	1378 µs	~14%
8 × 16384	1268 µs	1117 µs	~12%

Benchmark script

import torch, tilelang
import torch.nn.functional as F
from flash_qla import chunk_gated_delta_rule as qla
from flash_qla.utils import l2norm

H, V = 12, 128
kw = dict(warmup=10, rep=100, backend="cudagraph")

def run(q, k, v, g, beta, scale):
    o = qla(q=q, k=k, v=v, g=g, beta=beta, scale=scale, output_final_state=True)
    return o[0] if isinstance(o, tuple) else o

for b, t in [(64, 4096), (32, 8192), (8, 16384)]:
    q = l2norm(torch.randn(b, t, H, 64, device="cuda", dtype=torch.bfloat16))
    k = l2norm(torch.randn(b, t, H, 64, device="cuda", dtype=torch.bfloat16))
    v = torch.randn(b, t, H, V, device="cuda", dtype=torch.bfloat16)
    g = torch.nn.functional.logsigmoid(torch.randn(b, t, H, device="cuda", dtype=torch.float32)) / 16
    beta = torch.rand(b, t, H, device="cuda", dtype=torch.float32)
    scale = 64 ** -0.5
    qp, kp = F.pad(q, (0, 64)).contiguous(), F.pad(k, (0, 64)).contiguous()  # pad-to-128 workaround
    native = tilelang.profiler.do_bench(lambda: run(q, k, v, g, beta, scale), **kw)
    padded = tilelang.profiler.do_bench(lambda: run(qp, kp, v, g, beta, scale), **kw)
    print(f"b={b} t={t}: pad->128={padded*1e3:.0f}us native64={native*1e3:.0f}us saved={100*(padded-native)/padded:.0f}%")

Correctness

I added forward parity tests against the fp64 reference (RTOL=0.02). For the dims we reject, the guard fires before the kernel runs. And for the dims we accept, they match the reference: head_dim_k in {64, 128} across the whole CORE_CONFIGS matrix (GQA, varlen, h0, both kv/vk layouts) in test_fwd_head_dim_k, head_dim_v at multiples of 16 spanning [32, 256] in test_fwd_head_dim_v, and the head_dim_k=128 x head_dim_v!=128 corner that neither single-axis sweep hits in test_fwd_head_dim_kv_cross.

Scope / limitations

This PR only widens the forward pass on Blackwell (SM100). So there's still the need for a follow-up on the backward pass, and running the sweep on Hopper (SM90), and on SM120.