minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[FullCG] Preserve attention LSE through the custom-op boundary (#31050)

Browse Source
This commit is contained in:
paulzhang-tm
2026-07-21 09:01:40 +08:00
committed by GitHub
parent 9462c303a5
commit d093c6a4bb
3 changed files with 499 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/radix_attention.py
+167 -26
View File
@@ -147,6 +147,7 @@ class RadixAttention(nn.Module):
v,
forward_batch: ForwardBatch,
save_kv_cache: bool = True,
key_value_num_tokens: Optional[int] = None,
**kwargs,
):
if k is not None:
@@ -158,9 +159,10 @@ class RadixAttention(nn.Module):
else:
k = k.view(-1, self.tp_k_head_num, self.v_head_dim)
context = get_tc_piecewise_forward_context()
if (
forward_batch.forward_mode.is_extend()
and get_tc_piecewise_forward_context() is not None
and context is not None
# ``_force_eager_attn`` is only set inside Inkling's eager
# norm+attn+sconv region, never during tc-piecewise capture. Reading
# the ContextVar under the fullgraph torch.compile trace is
@@ -232,14 +234,39 @@ class RadixAttention(nn.Module):
q, k, v, output, save_kv_cache, self.layer_id, kwargs
)
return output
# Chunked-prefix MHA needs LSE to merge independently normalized
# suffix and cached-prefix attention states.
return_lse = bool(forward_batch.mha_return_lse)
mha_companion_layers = context.mha_companion_layers
use_mha_companion = (
mha_companion_layers is not None
and mha_companion_layers[self.layer_id] is self
)
if is_in_breakable_cuda_graph():
breakable_unified_attention_with_output(
q, k, v, output, save_kv_cache, self.layer_id, **kwargs
op = (
breakable_unified_attention_with_output_and_lse
if return_lse
else breakable_unified_attention_with_output
)
else:
unified_attention_with_output(
q, k, v, output, save_kv_cache, self.layer_id, **kwargs
op = (
unified_attention_with_output_and_lse
if return_lse
else unified_attention_with_output
)
lse = op(
q,
k,
v,
output,
save_kv_cache,
self.layer_id,
use_mha_companion=use_mha_companion,
key_value_num_tokens=key_value_num_tokens,
**kwargs,
)
if return_lse:
return output.view(-1, self.tp_q_head_num, self.v_head_dim), lse
return output
else:
return get_attn_backend().forward(
@@ -253,16 +280,17 @@ class RadixAttention(nn.Module):
)
@register_custom_op(mutates_args=["output"])
@register_split_op()
def unified_attention_with_output(
def _unified_attention_with_output_impl(
query: torch.Tensor,
key: Optional[torch.Tensor],
value: Optional[torch.Tensor],
output: torch.Tensor,
save_kv_cache: bool,
layer_id: int,
use_mha_companion: bool,
return_lse: bool,
*,
key_value_num_tokens: Optional[int] = None,
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
sinks: Optional[torch.Tensor] = None,
@@ -272,29 +300,38 @@ def unified_attention_with_output(
is_neox: Optional[bool] = None,
llama_4_scaling: Optional[torch.Tensor] = None,
topk_indices: Optional[torch.Tensor] = None,
) -> None:
) -> Optional[torch.Tensor]:
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layers = context.attention_layers
attention_layer = attention_layers[layer_id]
real_num_tokens = forward_batch.num_token_non_padded_cpu
real_query_num_tokens = forward_batch.num_token_non_padded_cpu
# Ordinary PCG attention pads Q/K/V to the same token bucket. Prefix MHA
# instead supplies a fixed-capacity K/V chunk whose extent is independent
# of the suffix queries, so its caller must preserve that separate extent.
if key_value_num_tokens is None:
key_value_num_tokens = real_query_num_tokens
query = query[:real_num_tokens]
query = query[:real_query_num_tokens]
if key is not None:
key = key[:real_num_tokens]
key = key[:key_value_num_tokens]
if value is not None:
value = value[:real_num_tokens]
value = value[:key_value_num_tokens]
if not save_kv_cache and context.mha_companion_layers is not None:
mha_companion_layer = context.mha_companion_layers[layer_id]
if mha_companion_layer is not None:
attention_layer = mha_companion_layer
# DeepSeek MLA has two RadixAttention instances per layer (attn_mqa and
# attn_mha) that share the same layer_id. Preserve the calling instance's
# identity through the custom-op boundary; save_kv_cache is not an identity
# signal because absorbed MLA can also disable a redundant cache store.
if use_mha_companion:
assert context.mha_companion_layers is not None
attention_layer = context.mha_companion_layers[layer_id]
assert attention_layer is not None
kwargs = {}
if q_rope is not None:
kwargs["q_rope"] = q_rope[:real_num_tokens]
kwargs["q_rope"] = q_rope[:real_query_num_tokens]
if k_rope is not None:
kwargs["k_rope"] = k_rope[:real_num_tokens]
kwargs["k_rope"] = k_rope[:key_value_num_tokens]
if sinks is not None:
kwargs["sinks"] = sinks
if cos_sin_cache is not None:
@@ -304,17 +341,17 @@ def unified_attention_with_output(
if llama_4_scaling is not None:
kwargs["llama_4_scaling"] = llama_4_scaling
if topk_indices is not None:
kwargs["topk_indices"] = topk_indices[:real_num_tokens]
kwargs["topk_indices"] = topk_indices[:real_query_num_tokens]
original_out_cache_loc = forward_batch.out_cache_loc
# Keep the original ForwardBatch object and only narrow cache locations for
# this backend call so model/backend state is still written to the same batch.
forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens]
forward_batch.out_cache_loc = original_out_cache_loc[:real_query_num_tokens]
# Store pre-allocated output for FA backend to write directly into.
# Must slice to real_num_tokens to match the narrowed query shape —
# Must slice to real_query_num_tokens to match the narrowed query shape —
# the FA kernel validates out.size(0) == q.size(0).
forward_batch._attn_output = output[:real_num_tokens]
forward_batch._attn_output = output[:real_query_num_tokens]
ret = get_attn_backend().forward(
query,
@@ -327,18 +364,119 @@ def unified_attention_with_output(
)
forward_batch.out_cache_loc = original_out_cache_loc
lse = None
if return_lse:
assert isinstance(ret, tuple)
ret, lse, *_ = ret
else:
assert isinstance(ret, torch.Tensor)
if ret.data_ptr() != output.data_ptr():
output[:real_num_tokens].view(ret.shape).copy_(ret)
output[:real_query_num_tokens].view(ret.shape).copy_(ret)
# During PCG replay the attention backend writes only the narrowed
# real-token slice (output[:real_num_tokens]) and leaves padded positions
# real-token slice (output[:real_query_num_tokens]) and leaves padded positions
# as uninitialized torch.empty garbage. Zero them so garbage (NaN/Inf) does
# not propagate through residual connections, MoE routing, and allreduce.
# This affects every backend that varlen-writes under PCG, not just ROCm.
# Use context.raw_num_tokens (pre-padding count from PCG runner) instead of
# forward_batch.extend_num_tokens, which is None for TARGET_VERIFY batches.
_zero_padded_pcg_tail(output, context)
return
if lse is not None and lse.shape[0] != output.shape[0]:
padded_lse = lse.new_zeros((output.shape[0], *lse.shape[1:]))
padded_lse[:real_query_num_tokens].copy_(lse)
lse = padded_lse
return lse
@register_custom_op(mutates_args=["output"])
@register_split_op()
def unified_attention_with_output(
query: torch.Tensor,
key: Optional[torch.Tensor],
value: Optional[torch.Tensor],
output: torch.Tensor,
save_kv_cache: bool,
layer_id: int,
*,
use_mha_companion: bool = False,
key_value_num_tokens: Optional[int] = None,
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
sinks: Optional[torch.Tensor] = None,
cos_sin_cache: Optional[torch.Tensor] = None,
is_neox: Optional[bool] = None,
llama_4_scaling: Optional[torch.Tensor] = None,
topk_indices: Optional[torch.Tensor] = None,
) -> None:
_unified_attention_with_output_impl(
query,
key,
value,
output,
save_kv_cache,
layer_id,
use_mha_companion,
False,
key_value_num_tokens=key_value_num_tokens,
q_rope=q_rope,
k_rope=k_rope,
sinks=sinks,
cos_sin_cache=cos_sin_cache,
is_neox=is_neox,
llama_4_scaling=llama_4_scaling,
topk_indices=topk_indices,
)
def _unified_attention_with_output_and_lse_fake(
query: torch.Tensor, *args, **kwargs
) -> torch.Tensor:
return query.new_empty((query.shape[0], query.shape[1]), dtype=torch.float32)
@register_custom_op(
mutates_args=["output"], fake_impl=_unified_attention_with_output_and_lse_fake
)
@register_split_op()
def unified_attention_with_output_and_lse(
query: torch.Tensor,
key: Optional[torch.Tensor],
value: Optional[torch.Tensor],
output: torch.Tensor,
save_kv_cache: bool,
layer_id: int,
*,
use_mha_companion: bool = False,
key_value_num_tokens: Optional[int] = None,
q_rope: Optional[torch.Tensor] = None,
k_rope: Optional[torch.Tensor] = None,
sinks: Optional[torch.Tensor] = None,
cos_sin_cache: Optional[torch.Tensor] = None,
is_neox: Optional[bool] = None,
llama_4_scaling: Optional[torch.Tensor] = None,
topk_indices: Optional[torch.Tensor] = None,
) -> torch.Tensor:
lse = _unified_attention_with_output_impl(
query,
key,
value,
output,
save_kv_cache,
layer_id,
use_mha_companion,
True,
key_value_num_tokens=key_value_num_tokens,
q_rope=q_rope,
k_rope=k_rope,
sinks=sinks,
cos_sin_cache=cos_sin_cache,
is_neox=is_neox,
llama_4_scaling=llama_4_scaling,
topk_indices=topk_indices,
)
assert lse is not None
return lse
@register_custom_op(mutates_args=["attn_out", "idx_out"])
@@ -401,6 +539,9 @@ def unified_sparse_attention_with_output(
breakable_unified_attention_with_output = eager_on_graph(True)(
unified_attention_with_output
)
breakable_unified_attention_with_output_and_lse = eager_on_graph(True)(
unified_attention_with_output_and_lse
)
def attention_with_output_extra_kwargs(
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
+11 -5
View File
@@ -141,7 +141,6 @@ def _forward_dsa_indexer_for_mha(
class DeepseekMHAForwardMixin:
def init_mha_forward(self: DeepseekV2AttentionMLA):
self.disable_chunked_prefix_cache = (
get_server_args().disable_chunked_prefix_cache
@@ -224,7 +223,6 @@ class DeepseekMHAForwardMixin:
)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
q, _, _, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
@@ -256,7 +254,6 @@ class DeepseekMHAForwardMixin:
latent_cache = latent_cache.unsqueeze(1)
if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant(
kv_a,
self.kv_a_layernorm.weight,
@@ -464,7 +461,6 @@ class DeepseekMHAForwardMixin:
accum_lse: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
# kv_b_proj needs BF16 input, but legacy q.dtype was BF16 by accident.
backend = _resolve_attn_backend(forward_batch)
pack_fn = getattr(backend, "pack_prefix_chunk_kv", None)
@@ -507,7 +503,17 @@ class DeepseekMHAForwardMixin:
k[..., : self.qk_nope_head_dim] = k_nope
k[..., self.qk_nope_head_dim :] = k_pe
output, lse = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
output, lse = self.attn_mha(
q,
k,
v,
forward_batch,
save_kv_cache=False,
# Prefix K/V is independent of the suffix query length. Under
# FullCG this is the fixed captured chunk extent; per-request
# active lengths remain encoded in the backend metadata.
key_value_num_tokens=k.shape[0],
)
tmp_output = torch.empty_like(accum_output)
tmp_lse = torch.empty_like(accum_lse)
merge_state_v2(output, lse, accum_output, accum_lse, tmp_output, tmp_lse)