Skip to content

[ATOM SGL] Minimax M3 FP8 & EAGLE #1564

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

Draft
ZhiweiYan-96 wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[ATOM SGL] Minimax M3 FP8 & EAGLE #1564

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
277 changes: 260 additions & 17 deletions atom/plugin/sglang/attention_backend/full_attention/full_attention_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,14 @@
# be handled by ATOM's native backend, making sglang-specific overrides
# unnecessary.

import logging
import math
from typing import TYPE_CHECKING, Optional

import torch

import sglang.srt.layers.attention.aiter_backend as _sglang_aiter
from aiter.ops.triton.gluon.pa_decode_gluon import get_recommended_splits
from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
from sglang.srt.layers.attention.utils import (
create_flashinfer_kv_indices_triton,
Expand All @@ -26,6 +28,7 @@
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.utils import get_bool_env_var

from atom.model_ops.base_attention import run_pa_decode_gluon
from atom.plugin.sglang.attention_backend.full_attention.kv_cache import (
set_kv_buffer_with_layout_shuffle as _set_kv_buffer_with_layout_shuffle,
)
Expand All @@ -35,6 +38,9 @@
build_pa_metadata_for_decode as _build_pa_metadata_for_decode,
build_pa_metadata_for_prefill as _build_pa_metadata_for_prefill,
)
from atom.utils import envs

logger = logging.getLogger("atom.plugin.sglang.attention_backend.full_attention")

if TYPE_CHECKING:
from sglang.srt.layers.radix_attention import RadixAttention
Expand Down Expand Up @@ -191,8 +197,14 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
self._init_draft_extend_v2_mla(forward_batch.batch_size, forward_batch)
elif self.use_mla and forward_batch.forward_mode.is_draft_extend():
self._init_draft_extend_mla(forward_batch.batch_size, forward_batch)
elif (not self.use_mla) and forward_batch.forward_mode.is_draft_extend_v2():
self._init_draft_extend_v2_mha(forward_batch.batch_size, forward_batch)
elif (not self.use_mla) and forward_batch.forward_mode.is_draft_extend():
self._init_draft_extend_mha(forward_batch.batch_size, forward_batch)
elif self.use_mla and forward_batch.forward_mode.is_target_verify():
self._init_target_verify_mla(forward_batch.batch_size, forward_batch)
elif (not self.use_mla) and forward_batch.forward_mode.is_target_verify():
self._init_target_verify_mha(forward_batch.batch_size, forward_batch)
else:
self._init_forward_metadata_extend(forward_batch)
self._fixup_page_table(forward_batch)
Expand Down Expand Up @@ -517,6 +529,66 @@ def _init_draft_extend_v2_mla(self, bs, forward_batch):
run_graph=False,
)

def _init_draft_extend_v2_mha(self, bs, forward_batch):
spec_info = forward_batch.spec_info
if spec_info is None:
raise RuntimeError("MHA draft_extend_v2 requires speculative metadata")
if forward_batch.extend_prefix_lens is None:
raise RuntimeError("MHA draft_extend_v2 requires extend_prefix_lens")
self.indices_updater_prefill.update(
forward_batch.req_pool_indices,
forward_batch.seq_lens,
forward_batch.seq_lens_sum,
prefix_lens=forward_batch.extend_prefix_lens,
encoder_lens=forward_batch.encoder_lens,
spec_info=None,
)
max_q_len = self._max_len(
forward_batch.extend_seq_lens_cpu,
forward_batch.extend_seq_lens,
)
max_kv_len = self._max_len(forward_batch.seq_lens_cpu, forward_batch.seq_lens)
self.forward_metadata = ForwardMetadata(
self.indices_updater_prefill.kv_indptr,
self.indices_updater_prefill.kv_indices,
self.qo_indptr[: bs + 1],
None,
max_q_len,
max_kv_len,
None,
None,
custom_mask=None,
mask_indptr=None,
max_extend_len=max_q_len,
)
def _init_draft_extend_mha(self, bs, forward_batch):
spec_info = forward_batch.spec_info
if spec_info is None:
raise RuntimeError("MHA draft_extend requires speculative metadata")
kv_indices, kv_indptr, qo_indptr, custom_mask = (
spec_info.generate_attn_arg_prefill(
forward_batch.req_pool_indices,
forward_batch.seq_lens,
forward_batch.seq_lens_sum,
self.req_to_token,
)
)
kv_indices = kv_indices.to(torch.int64)
draft_max_extend_len = int(torch.max(spec_info.num_accept_tokens).item())
self.forward_metadata = ForwardMetadata(
kv_indptr,
kv_indices,
qo_indptr,
None,
draft_max_extend_len,
None,
None,
None,
custom_mask=custom_mask,
mask_indptr=None,
max_extend_len=draft_max_extend_len,
)

def _init_target_verify_mla(self, bs, forward_batch):
"""Init MLA metadata for speculative target_verify."""
spec_info = forward_batch.spec_info
Expand Down Expand Up @@ -709,6 +781,55 @@ def _init_extend_mha(self, bs, forward_batch):
forward_batch.seq_lens,
)

def _init_target_verify_mha(self, bs, forward_batch):
spec_info = forward_batch.spec_info
if spec_info is None:
raise RuntimeError("MHA target_verify requires speculative metadata")

draft_num = spec_info.draft_token_num
qo_indptr = torch.arange(
0,
(1 + bs) * draft_num,
step=draft_num,
dtype=torch.int32,
device=self.device,
)

kv_indptr = self.kv_indptr[: bs + 1]
kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]

kv_indices_len = int(kv_indptr[-1].item())
kv_indices = torch.empty(kv_indices_len, dtype=torch.int64, device=self.device)
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
forward_batch.req_pool_indices,
forward_batch.seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)

seq_mask_len = draft_num * (forward_batch.seq_lens + draft_num)
mask_indptr = self.mask_indptr
mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len[:bs], dim=0)
mask_indptr = mask_indptr[: bs + 1]

self.forward_metadata = ForwardMetadata(
kv_indptr,
kv_indices,
qo_indptr,
None,
draft_num,
None,
None,
None,
custom_mask=spec_info.custom_mask,
mask_indptr=mask_indptr,
max_extend_len=draft_num,
)

def _fixup_page_table(self, forward_batch: ForwardBatch):
"""Post-process page_table for non-MLA extend mode."""
if (
Expand Down Expand Up @@ -764,6 +885,14 @@ def init_cuda_graph_state(
else:
self.cuda_graph_kv_indices = kv_indices_buf

draft_tokens = int(self.num_draft_tokens or 1)
custom_mask_size = max_bs * draft_tokens * (self.max_context_len + draft_tokens)
self.cuda_graph_custom_mask = torch.ones(
custom_mask_size,
dtype=torch.bool,
device=self.device,
)

# Always use preshuffle layout for pa_fwd_asm
self.page_table = torch.zeros(
(max_bs, self.max_context_len // self.page_size),
Expand Down Expand Up @@ -1584,8 +1713,9 @@ def _update_decode_page_table(
def _should_use_native_dense_mha(self, layer) -> bool:
sliding_window_size = getattr(layer, "sliding_window_size", None)
is_minimax_m3 = bool(getattr(layer, "_atom_minimax_m3_dense_mha", False))
is_eagle3 = bool(getattr(layer, "_atom_eagle3_dense_mha", False))
if (
is_minimax_m3
(is_minimax_m3 or is_eagle3)
and not self.use_mla
and not layer.is_cross_attention
and layer.head_dim == 128
Expand Down Expand Up @@ -1740,25 +1870,16 @@ def forward_extend(
elif self.use_mla:
forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
else:
k_buffer, v_buffer = forward_batch.token_to_kv_pool.get_kv_buffer(
layer.layer_id
)
k_scale, v_scale = self._ensure_fp8_kv_scales(
layer, k_buffer, self.page_size
)
self.set_kv_buffer_with_layout_shuffle(
cache_loc,
k,
v,
k_buffer,
v_buffer,
k_scale,
v_scale,
self.page_size,
forward_batch.token_to_kv_pool.set_kv_buffer(
layer, cache_loc, k, v
)

if self.use_mla:
return self._forward_extend_mla(q, k, v, layer, forward_batch)
if forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend(
include_v2=True
):
return self._forward_extend_mha_speculative(q, k, v, layer, forward_batch)
if bool(getattr(layer, "_atom_minimax_m3_dense_mha", False)):
# M3 dense decode benefits from the native ragged path, but batched
# SGLang prefill is safer through the standard varlen extend path.
Expand Down Expand Up @@ -1801,6 +1922,42 @@ def _forward_extend_native_dense_mha(self, q, layer, forward_batch):
)
return o.view(-1, layer.tp_q_head_num * layer.head_dim)

def _forward_extend_mha_speculative(self, q, k, v, layer, forward_batch):
"""Non-MLA EAGLE verify/draft-extend path.

Mirrors /app/sglang's AiterAttnBackend: speculative MHA uses the
extend-attention kernel with custom masks, not plain flash_attn_varlen.
"""

if k is None or v is None:
raise RuntimeError("MHA speculative extend requires explicit k/v tensors")

if layer.qk_head_dim != layer.v_head_dim:
o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
else:
o = torch.empty_like(q)

self.extend_attention_fwd(
q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
k.contiguous(),
v.contiguous(),
o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
self.forward_metadata.qo_indptr,
self.forward_metadata.kv_indptr,
self.forward_metadata.kv_indices,
self.forward_metadata.custom_mask,
True,
self.forward_metadata.mask_indptr,
self.forward_metadata.max_extend_len,
1.0,
1.0,
layer.scaling,
logit_cap=layer.logit_cap,
)
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)

def _forward_extend_mha(self, q, k, v, layer, forward_batch):
"""Non-MLA extend path: standard MHA with flash_attn_varlen_func."""
seqlens_in_batch = forward_batch.seq_lens
Expand Down Expand Up @@ -2389,6 +2546,18 @@ def forward_decode(
)
return self._forward_decode_native_dense_mha(q, layer, forward_batch)

if (
envs.ATOM_FORCE_ATTN_TRITON
and not layer.is_cross_attention
and layer.qk_head_dim == layer.v_head_dim
and layer.head_dim == layer.qk_head_dim
):
if save_kv_cache:
self._set_kv_buffer_native_dense(
layer, forward_batch.out_cache_loc, k, v, forward_batch
)
return self._forward_decode_native_dense_mha(q, layer, forward_batch)

o = q.new_empty((batch_size, layer.tp_q_head_num, head_dim_out))

if save_kv_cache:
Expand Down Expand Up @@ -2485,7 +2654,81 @@ def forward_decode(
return o.view(-1, layer.tp_q_head_num * head_dim_out)

def _forward_decode_native_dense_mha(self, q, layer, forward_batch):
k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
k_cache, v_cache = forward_batch.token_to_kv_pool.get_kv_buffer(
layer.layer_id
)
if envs.ATOM_FORCE_ATTN_TRITON:
block_size = self.page_size
num_slots, num_kv_heads, head_size = k_cache.shape
num_blocks = num_slots // block_size
k_cache = k_cache[: num_blocks * block_size].view(
num_blocks, block_size, num_kv_heads, head_size
)
v_cache = v_cache[: num_blocks * block_size].view(
num_blocks, block_size, num_kv_heads, layer.v_head_dim
)
x = 16 // k_cache.element_size()
k_cache = (
k_cache.view(num_blocks, block_size, num_kv_heads, head_size // x, x)
.permute(0, 2, 3, 1, 4)
.contiguous()
)
v_cache = (
v_cache.view(num_blocks, block_size // x, x, num_kv_heads, layer.v_head_dim)
.permute(0, 3, 1, 4, 2)
.contiguous()
)
q_3d = q.view(-1, layer.tp_q_head_num, layer.head_dim)
out = torch.empty_like(q_3d, dtype=self.input_dtype)
num_seqs = self.forward_metadata.kv_lens.shape[0]
query_group_size = 1 * (layer.tp_q_head_num // layer.tp_k_head_num)
max_context_partition_num = get_recommended_splits(
num_seqs, layer.tp_k_head_num
)
context_partition_size = 256
intermediate_shape = (
num_seqs,
layer.tp_k_head_num,
max_context_partition_num,
query_group_size,
)
exp_sums = torch.empty(
intermediate_shape, dtype=torch.float32, device=q.device
)
max_logits = torch.empty(
intermediate_shape, dtype=torch.float32, device=q.device
)
temporary_output = torch.empty(
*intermediate_shape,
layer.head_dim,
dtype=q.dtype,
device=q.device,
)
run_pa_decode_gluon(
output=out,
q=q_3d,
k_cache=k_cache,
v_cache=v_cache,
context_lens=self.forward_metadata.kv_lens,
block_tables=self.forward_metadata.page_table,
softmax_scale=layer.scaling,
max_seqlen_q=1,
max_context_partition_num=max_context_partition_num,
context_partition_size=context_partition_size,
compute_type=torch.bfloat16,
q_scale=None,
k_scale=None,
v_scale=None,
exp_sums=exp_sums,
max_logits=max_logits,
temporary_output=temporary_output,
alibi_slopes=None,
sinks=None,
sliding_window=-1,
ps=True,
)
return out.reshape_as(q)

aiter_kv_str = self._get_aiter_paged_ragged_kv_cache_dtype()

o = torch.empty_like(q, dtype=self.input_dtype)
Expand Down
4 changes: 4 additions & 0 deletions atom/plugin/sglang/attention_backend/full_attention/metadata.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,10 @@ class ForwardMetadata:
fp8_prefill_kv_indices: Optional[torch.Tensor] = None
num_kv_splits: Optional[int] = None
run_graph: Optional[bool] = True
custom_mask: Optional[torch.Tensor] = None
mask_indptr: Optional[torch.Tensor] = None
max_extend_len: Optional[int] = None
swa_page_table: Optional[torch.Tensor] = None
# PA metadata for pa_persistent_fwd (only used in decode mode, non-MLA)
pa_metadata_qo_indptr: Optional[torch.Tensor] = None
pa_metadata_pages_kv_indptr: Optional[torch.Tensor] = None
Expand Down
Loading
Loading