Findings from integrating FlashQLA's Blackwell GDN kernels into a JAX trainer

[celiolarcher]

Hi!

We integrated FlashQLA's sm_100 Gated Delta Rule kernels as a training backend on B200s. It works great (~2× over our previous Triton path at kernel level, bit-reproducible backward), and along the way we hit three bugs and one optimization opportunity worth reporting.

1. tensor_cache keyed by id(tensor) returns stale results for recycled tensor addresses

flash_qla/utils/index.py's tensor_cache decorator uses id(tensor) as the cache key. Python recycles object addresses aggressively: when a cached tensor is freed and a NEW tensor is allocated at the same address (guaranteed in tight loops, and in any embedding host where buffers are wrapped repeatedly, e.g. torch.Tensor views of external memory), the cache serves the PREVIOUS tensor's result.

For prepare_chunk_indices / prepare_chunk_offsets this means a wrong chunk grid for the new cu_seqlens → out-of-bounds writes, non-deterministic outputs and (in our case) glibc heap corruption aborts.

Repro sketch:

import torch
from flash_qla.ops.utils.index import prepare_chunk_indices

for i in range(1000):
    # different cu each iteration, but allocated at recycled addresses
    lens = torch.randint(1, 512, (8,), device="cuda")
    cu = torch.cat([torch.zeros(1, device="cuda", dtype=torch.int32),
                    lens.cumsum(0).to(torch.int32)])
    idx = prepare_chunk_indices(cu, 64)
    ref = prepare_chunk_indices.__wrapped__(cu, 64)
    assert torch.equal(idx, ref), f"stale cache hit at iter {i}"

Suggested fix: key by tensor content (e.g. bytes of a small cpu copy) or by (data_ptr, shape, version_counter); id() alone is unsound for dead-object reuse.

2. Length-0 segments in cu_seqlens produce non-deterministic garbage (fused_fwd Store-O epilogue)

With repeated offsets in cu_seqlens (length-0 segments — natural when padding a static-shape cu buffer), tilelang_fused_chunk_gdr_fwd corrupts the LAST REAL chunk's output non-deterministically. Root cause: the Store-O epilogue runs even when num_iters == 0; seq_split_idx underflows to seq_end - 64 and the ghost CTA races the real CTA that owns that range, writing its uninitialized o_shared over real output. Whether corruption appears depends on CTA scheduling (wave timing), which is why it looks random and quantity-dependent.

One-line fix, validated bit-exact with up to 120 ghost segments appended, and your full test suite still passes (120/120):

                     # Store O
                     T.barrier_wait(bar_o, 0)
-                    if store_o:
+                    if store_o and num_iters > 0:

(flash_qla/ops/gated_delta_rule/chunk/blackwell/fused_fwd.py, Store-O epilogue.)

The backward kernels already tolerate length-0 segments. Supporting them in fwd makes static-shape integration (JAX/XLA, CUDA graphs, TensorRT-style engines) much easier: callers can pad cu_seqlens to a fixed length with trailing ghosts instead of recompiling per segment count.

3. TVM host stubs reject strides == nullptr DLTensors

The generated host stubs validate strides explicitly and fail ("input g_raw strides[1] constraint not satisfied"-style errors) when a DLTensor carries strides == nullptr — which the DLPack spec defines as valid shorthand for compact row-major. Callers constructing DLTensors manually (any C-ABI embedding) must materialize explicit strides. Either accepting nullptr in the generated checks or documenting the requirement would save integrators a debugging session.

4. (Enhancement) Skip the final_correction store in the no-CP specialization

prepare_h unconditionally writes the CP final_correction tensor through a TMA map even when compiled for the no-CP path, where nothing reads it. At our shape (packed 64k tokens, 32 v-heads, d=128) that's a ~0.5 GiB write-only buffer the caller must allocate per backward call. Gating the store on the CP specialization (like the dead h0/ht parameters, which TileLang already eliminates from the device signature) would remove the allocation and the write bandwidth entirely.

---

Context for reproducibility: B200 (sm_100), CUDA 13, tilelang==0.1.9, FlashQLA @ 2ef89a7.

Thanks for the great kernel!

[Starmys]

Thank you for your detailed and valuable feedback! We're very pleased that FlashQLA's Blackwell kernels are working well in your JAX trainer.

Regarding the four points you raised:

1. tensor_cache keyed by id(tensor): We agree that using id() is unsound due to object address recycling. In recent commits, we have re‑implemented the cache using a FLA‑style tensor cache that directly caches based the input tensor data itself (i.e., cu_seqlens), rather than any address‑derived identifier.

2. Length‑0 segments in cu_seqlens: We appreciate your precise diagnosis and the one‑line fix. We have applied the suggested change (if store_o and num_iters > 0) in the forward kernel.

3. TVM host stubs rejecting strides == nullptr: We have addressed this by adding a clear note in the docstring of chunk_gated_delta_rule(), explicitly warning integrators that TVM does not accept strides == nullptr and that explicit strides must be provided when constructing DLTensors manually. This should help users avoid the debugging overhead.

4. Skip the final_correction store in the no‑CP specialization: This is a valuable optimization. We have added this to our roadmap and plan to address it in the future.

Thank you again for your thorough investigation and for sharing your findings!