Skip to content

[AMD] Skip redundant CatArrayBatchedCopy in GLM-5 NSA TileLang decode #24125

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
HaiShaw merged 3 commits into from
May 13, 2026
Merged

[AMD] Skip redundant CatArrayBatchedCopy in GLM-5 NSA TileLang decode #24125

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8d4b571
perf(rocm-nsa): skip redundant concat in NSA TileLang decode (cat-skip)
Jacob0226 Apr 30, 2026
1f2b7c4
fixup! address review: keep cat-skip strictly on AMD side
Jacob0226 May 5, 2026
35d60b5
Merge branch 'main' into jacob/glm5-rocm-nsa-cat-skip
HaiShaw May 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 14 additions & 4 deletions python/sglang/srt/layers/attention/nsa_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1548,7 +1548,13 @@ def forward_decode(
q_rope = q_rope.view(
-1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
)
# Caller passed split q_nope / q_rope; we'll need to concat below if
# the chosen impl wants q_all.
q_all = None
else:
# Caller passed already-concatenated q (q_all = q). Reuse it directly
# via a zero-copy view; the impl-specific blocks below will skip the
# otherwise redundant concat_mla_absorb_q_general call.
q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
q_nope = q_all[:, :, : layer.v_head_dim]
q_rope = q_all[:, :, layer.v_head_dim :]
Expand All @@ -1573,8 +1579,12 @@ def forward_decode(
page_size=1,
)

# Cat-skip is HIP-only: when caller passes q_rope=None on HIP, q_all has
# already been set to a zero-copy view of q in the else branch above and
# we can reuse it directly. Non-HIP backends fall through to the original
# "always cat" behavior to keep CUDA / MUSA paths byte-identical.
if self.nsa_decode_impl == "flashmla_sparse":
if q_rope is not None:
if q_all is None or not _is_hip:

@1am9trash 1am9trash May 5, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I think we can only make the change in amd side (e.g, tilelang/aiter backend).
This change may be never reached in nv code path (e.g, flashmla_sparse/flashmla_kv).

@Jacob0226 Jacob0226 May 5, 2026

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.

Good catch, thanks! Done in 1f2b7c4.

Reverted the gate change for flashmla_sparse and flashmla_kv back to the pre-patch if q_rope is not None: pattern. Both impls import from sgl_kernel.flashmla_ops which requires CUDA driver, so on HIP they are unreachable and the new not _is_hip clause was a strict no-op there. On CUDA the behavior is byte-identical (gate is always True either way).

The cat-skip itself stays on the HIP-only tilelang and aiter blocks, which is what this PR is actually solving. Also moved the "Cat-skip is HIP-only ..." comment from above flashmla_sparse to above tilelang to match the new (narrower) scope.

q_all = concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_flashmla_sparse(
q_all=q_all,
Expand All @@ -1584,7 +1594,7 @@ def forward_decode(
v_head_dim=layer.v_head_dim,
)
elif self.nsa_decode_impl == "flashmla_kv":
if q_rope is not None:
if q_all is None or not _is_hip:
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_flashmla_kv(
q_all=q_all,
Expand All @@ -1597,7 +1607,7 @@ def forward_decode(
page_table_1=page_table_1,
)
elif self.nsa_decode_impl == "tilelang":
if q_rope is not None:
if q_all is None or not _is_hip:
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
return self._forward_tilelang(
q_all=q_all,
Expand All @@ -1622,7 +1632,7 @@ def forward_decode(
page_size=1,
)
elif self.nsa_decode_impl == "aiter":
if q_rope is not None:
if q_all is None or not _is_hip:
q_all = torch.cat([q_nope, q_rope], dim=-1)
return self._forward_aiter(
q_all=q_all,
Expand Down
66 changes: 48 additions & 18 deletions python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -376,25 +376,55 @@ def forward_absorb_core(
self.rotary_emb.is_neox_style,
q_out_dtype=kv_cache_dtype,
)
q_nope_fused = q_cat[..., : self.kv_lora_rank]
q_pe_fused = q_cat[..., self.kv_lora_rank :]
save_kv_cache = False
if llama_4_scaling is not None:
q_nope_fused *= llama_4_scaling
attn_output = self.attn_mqa(
q_nope_fused,
None,
None,
forward_batch,
q_rope=q_pe_fused,
k_rope=k_pe_fused,
save_kv_cache=save_kv_cache,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
# On decode, pass q_cat directly to attn_mqa with q_rope=None so
# nsa_backend.forward_decode reuses q_cat as a zero-copy view
# (`q.contiguous().view(...)` fast-path) instead of running the
# redundant `concat_mla_absorb_q_general(q_nope_fused, q_pe_fused)`
# that would otherwise rebuild a tensor byte-identical to q_cat.
# On ROCm tilelang decode, this eliminates the
# `CatArrayBatchedCopy<OpaqueType<1u>, ...>` kernel that used to
# fire once per layer per decode step (~2.6 us / layer saved).
# Prefill keeps the split form because nsa_backend.forward_extend
# asserts `q_rope is not None`.
if forward_batch.forward_mode.is_decode_or_idle():
if llama_4_scaling is not None:
# llama_4_scaling applies only to the q_nope portion;
# mutate in place via the slice view of q_cat.
q_cat[..., : self.kv_lora_rank] *= llama_4_scaling
attn_output = self.attn_mqa(
q_cat,
None,
None,
forward_batch,
q_rope=None,
k_rope=k_pe_fused,
save_kv_cache=save_kv_cache,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
else:
q_nope_fused = q_cat[..., : self.kv_lora_rank]
q_pe_fused = q_cat[..., self.kv_lora_rank :]
if llama_4_scaling is not None:
q_nope_fused *= llama_4_scaling
attn_output = self.attn_mqa(
q_nope_fused,
None,
None,
forward_batch,
q_rope=q_pe_fused,
k_rope=k_pe_fused,
save_kv_cache=save_kv_cache,
**(
dict(topk_indices=topk_indices)
if topk_indices is not None
else {}
),
)
else:
extra_args = {}
if self._fuse_rope_for_trtllm_mla(forward_batch):
Expand Down
Loading