feat : optimized TileLang DSA sparse-MLA backward for GLM-5 by Butterfingrz · Pull Request #1625 · radixark/miles

Butterfingrz · 2026-07-10T15:06:10Z

Summary

Optimize the TileLang GLM-5 DSA sparse-MLA backward kernel without changing its mathematical behavior.

The optimized kernel preserves FP32 accumulation for dQ and dKV and achieves 4.03–4.30× kernel-level speedup over the original 128-thread implementation on H200 across sequence lengths from 2K to 16K.

Changes

Add an optimized sparse-MLA backward kernel with:
- split-D computation;
- pipelined and swizzled shared-memory staging;
- reduced redundant memory traffic and GEMM operand reuse.
Select threads adaptively:
- 256 threads when block_H >= 64;
- 128 threads otherwise.
Switch SparseMLA backward to the optimized kernel.
Add correctness and performance tests.

Validation

Tested on H200 (sm_90) with TileLang 0.1.11 and PyTorch 2.11.0+cu129.

dQ cosine similarity: 1.000000
dKV cosine similarity: 1.000000
Kernel speedup: 4.03–4.30×

gemini-code-assist

Code Review

This pull request introduces an optimized TileLang backward kernel implementation (tilelang_sparse_mla_bwd_opt.py) for GLM-5 DSA sparse-MLA, updates the main sparse_mla.py autograd path to use this optimized variant, and adds corresponding performance benchmarks and numerical parity tests. The reference backward kernel was also updated to support adaptive thread counts. The code reviewer provided feedback recommending the use of ValueError instead of assert for input validation, correcting a typo in the parameter name is_casual to is_causal, and dynamically deriving the model dimension D from the input tensor shape rather than hardcoding it.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-10T15:08:19Z

+    assert dtype == T.bfloat16
+    assert accum_dtype == T.float32


Use ValueError instead of assert for validating function arguments to prevent silent bypasses when assertions are disabled globally (e.g., via python -O).

Suggested change

assert dtype == T.bfloat16

assert accum_dtype == T.float32

if dtype != T.bfloat16:

raise ValueError("dtype must be T.bfloat16")

if accum_dtype != T.float32:

raise ValueError("accum_dtype must be T.float32")

References

Use ValueError instead of assert for validating function or constructor arguments (such as checking for positive or non-negative values).

Updated as requested.

gemini-code-assist · 2026-07-10T15:08:19Z

+    assert dtype == T.bfloat16
+    assert accum_dtype == T.float32


Use ValueError instead of assert for validating function arguments to prevent silent bypasses when assertions are disabled globally (e.g., via python -O).

Suggested change

assert dtype == T.bfloat16

assert accum_dtype == T.float32

if dtype != T.bfloat16:

raise ValueError("dtype must be T.bfloat16")

if accum_dtype != T.float32:

raise ValueError("accum_dtype must be T.float32")

References

Use ValueError instead of assert for validating function or constructor arguments (such as checking for positive or non-negative values).

Updated accordingly.

gemini-code-assist · 2026-07-10T15:08:19Z

+    assert is_causal == True, "non-casual is not supported now"
+    assert topk % block_size == 0, "otherwise will load some index=0 thus causing wrong kv to be loaded"
+    assert dtype == T.bfloat16
+    assert accum_dtype == T.float32
+    assert indices_dtype == T.int32


Use ValueError instead of assert for validating function arguments to prevent silent bypasses when assertions are disabled globally (e.g., via python -O).

Suggested change

assert is_causal == True, "non-casual is not supported now"

assert topk % block_size == 0, "otherwise will load some index=0 thus causing wrong kv to be loaded"

assert dtype == T.bfloat16

assert accum_dtype == T.float32

assert indices_dtype == T.int32

if not is_causal:

raise ValueError("non-causal is not supported now")

if topk % block_size != 0:

raise ValueError("otherwise will load some index=0 thus causing wrong kv to be loaded")

if dtype != T.bfloat16:

raise ValueError("dtype must be T.bfloat16")

if accum_dtype != T.float32:

raise ValueError("accum_dtype must be T.float32")

if indices_dtype != T.int32:

raise ValueError("indices_dtype must be T.int32")

References

Use ValueError instead of assert for validating function or constructor arguments (such as checking for positive or non-negative values).

Updated accordingly.

gemini-code-assist · 2026-07-10T15:08:19Z

+    return sparse_mla_bwd_kernel
+
+
+def sparse_mla_bwd(q, kv, o, do, indices, lse, sm_scale=None, is_casual=True, return_kernel=False, delta=None):


Rename is_casual to is_causal to fix the typo and prevent potential TypeError when calling the function with keyword arguments.

Suggested change

def sparse_mla_bwd(q, kv, o, do, indices, lse, sm_scale=None, is_casual=True, return_kernel=False, delta=None):

def sparse_mla_bwd(q, kv, o, do, indices, lse, sm_scale=None, is_causal=True, return_kernel=False, delta=None):

gemini-code-assist · 2026-07-10T15:08:19Z

+    assert q.is_contiguous()
+    assert kv.is_contiguous()
+    assert indices.is_contiguous()
+    assert lse.is_contiguous()
+    B, S, H, dim_plus_tail_dim = q.shape
+    _, S_kv, kv_group, _ = kv.shape
+    assert kv.shape[-1] == dim_plus_tail_dim
+    assert kv.shape[0] == B
+    D = 512
+    D_tail = dim_plus_tail_dim - D
+    topk = indices.shape[-1]
+    assert indices.shape == (B, S, kv_group, topk)
+    assert lse.shape == (B, S, H)


Use ValueError instead of assert for validating function arguments. Additionally, avoid hardcoding the model dimension D = 512 and instead derive it dynamically from the input tensor shape (e.g., o.shape[-1]).

if not q.is_contiguous(): raise ValueError("q must be contiguous") if not kv.is_contiguous(): raise ValueError("kv must be contiguous") if not indices.is_contiguous(): raise ValueError("indices must be contiguous") if not lse.is_contiguous(): raise ValueError("lse must be contiguous") B, S, H, dim_plus_tail_dim = q.shape _, S_kv, kv_group, _ = kv.shape if kv.shape[-1] != dim_plus_tail_dim: raise ValueError("kv dimension mismatch") if kv.shape[0] != B: raise ValueError("kv batch size mismatch") D = o.shape[-1] D_tail = dim_plus_tail_dim - D topk = indices.shape[-1] if indices.shape != (B, S, kv_group, topk): raise ValueError("indices shape mismatch") if lse.shape != (B, S, H): raise ValueError("lse shape mismatch")

References

Use ValueError instead of assert for validating function or constructor arguments (such as checking for positive or non-negative values).

Avoid hardcoding model dimensions like head dimension D; derive them from the input tensor shape or configuration instead.

Updated as requested.

gemini-code-assist · 2026-07-10T15:08:19Z

+    assert lse.shape == (B, S, H)
+
+    preprocess_kernel = preprocess(B, S, H, D)
+    bwd_kernel = bwd(B, S, S_kv, H, D, D_tail, topk, kv_group, sm_scale, is_casual)


Update the parameter name from is_casual to is_causal to match the renamed parameter in the signature.

Suggested change

bwd_kernel = bwd(B, S, S_kv, H, D, D_tail, topk, kv_group, sm_scale, is_casual)

bwd_kernel = bwd(B, S, S_kv, H, D, D_tail, topk, kv_group, sm_scale, is_causal)

Updated as requested.

Butterfingrz · 2026-07-10T15:13:57Z

Would appreciate a review from @fzyzcjy @yueming-yuan @guapisolo @maocheng23 ， Thank you! This PR only makes a small change to the TileLang DSA backward kernel used by GLM-5.

Add an optimized backward kernel for the GLM-5 series (5/5.1/5.2) DSA sparse-MLA, alongside its parity test and benchmark: - tilelang_sparse_mla_bwd_opt: pipelined split-D + shared-memory swizzle + GEMM-operand-reuse variant; the sparse_mla.py autograd path now uses it (fp32-accumulated, parity cosine ~1.0 vs the vendored reference). - Adaptive threads in both kernels: 256 once block_H>=64 (per-rank heads >= 33), else fall back to 128 so small/high-TP shapes still build. - Parity test under tests/manual/models/glm5 and a seq-len-sweep benchmark comparing the reference and optimized kernels. Signed-off-by: Butterfingrz <13524387014@163.com>

gemini-code-assist Bot reviewed Jul 10, 2026

View reviewed changes

Butterfingrz force-pushed the feat/glm5_dsa_bwd_kernel branch from aa0f0b9 to cd26a22 Compare July 10, 2026 15:28

Butterfingrz marked this pull request as draft July 10, 2026 15:54

Butterfingrz marked this pull request as ready for review July 13, 2026 05:24

-    assert dtype == T.bfloat16
-    assert accum_dtype == T.float32
+    if dtype != T.bfloat16:
+        raise ValueError("dtype must be T.bfloat16")
+    if accum_dtype != T.float32:
+        raise ValueError("accum_dtype must be T.float32")

		return sparse_mla_bwd_kernel


		def sparse_mla_bwd(q, kv, o, do, indices, lse, sm_scale=None, is_casual=True, return_kernel=False, delta=None):

	bwd_kernel = bwd(B, S, S_kv, H, D, D_tail, topk, kv_group, sm_scale, is_casual)
	bwd_kernel = bwd(B, S, S_kv, H, D, D_tail, topk, kv_group, sm_scale, is_causal)

Uh oh!

Conversation

Butterfingrz commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Validation

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Butterfingrz Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Butterfingrz Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Butterfingrz Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Butterfingrz Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Butterfingrz Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Butterfingrz commented Jul 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Butterfingrz commented Jul 10, 2026 •

edited

Loading