# okc: problem=muse_attention_decode # Experiment: limit sliding-window split count using available batch/head parallelism. """SGLang's attention for Muse Spark Open on H200 (FA3 backend, FP8 E4M3 KV cache), at SGLang commit a55ab7d5ae1bba59ac351e55026fe7f7644ecc5f: the kernels between the QKV and output projections (MuseSpark1xAttention.gated_attention), in order: FusedQKNorm (Triton) per-head RMSNorm of q and k from the fused projection output, with interleaved RoPE fused in on sliding layers set_kv_buffer k and v cast to E4M3 (divided by the cache scales when set), stored at the step's cache rows (store_cache) q.to(float8_e4m3fn) as prepare_paged_mha_query does for an FP8 cache flash_attn_with_kvcache (FA3) paged, causal, window (2047, 0) on sliding layers, 5.473/8 softmax scale rms_norm_gate (JIT CUDA) per-head RMSNorm x sigmoid(gate) Differs from SGLang: the step's positions, cache rows, FA3 lengths and FA3 scheduler metadata come in with the inputs (SGLang builds them once per step for all layers).""" from typing import TypedDict import torch from sgl_kernel.flash_attn import flash_attn_with_kvcache from sglang.kernels.ops.kvcache.kvcache import store_cache from sglang.kernels.ops.muse_spark_v12.fused_qk_norm import FusedQKNorm from sglang.kernels.ops.muse_spark_v12.norm_gate import rms_norm_gate class Config(TypedDict): batch: int # B q_len: int # L: 1 for decode, 8 for a speculative verify block (causal inside it) num_q_heads: int # Hq: 16 per GPU under TP8 num_kv_heads: int # Hkv: 2; each serves Hq / Hkv query heads head_dim: int # 64 page_size: int # 1 in the scored cases (as SGLang serves Muse on H200), 64 in some tests window: int # 2048 on sliding-window layers (a token and the 2,047 before it); 0 = all rope_theta: float # 500000 on sliding-window layers; 0 = no RoPE qk_scale_factor: float sm_scale: float # 1/8 kv_scale: float # the cache scales' value: 1 in the scored cases (SGLang leaves them unset) class Inputs(TypedDict): qkv: torch.Tensor # bfloat16 [B, L, (2*Hq + 2*Hkv)*64], 2304 under tp8: q (Hq*64), keys (Hkv*64), # values (Hkv*64), the output gate's logits (Hq*64) positions: torch.Tensor # int64 [B*L], each new token's position cache_loc: torch.Tensor # int64 [B*L], each new token's flat cache row cache_seqlens: torch.Tensor # int32 [B], seq_lens + L cu_seqlens_q: torch.Tensor # int32 [B+1], 0, L, 2L, ... scheduler_metadata: torch.Tensor | None # int32, FA3's tile schedule on full-attention decode steps k_cache: torch.Tensor # float8_e4m3fn [pages, page_size, Hkv, 64], written in place v_cache: torch.Tensor # float8_e4m3fn [pages, page_size, Hkv, 64], written in place k_scale: torch.Tensor # float32 [], key = e4m3 * k_scale v_scale: torch.Tensor # float32 [], value = e4m3 * v_scale page_table: torch.Tensor # int32 [B, max_pages], physical page of each logical page (-1 unused) seq_lens: torch.Tensor # int32 [B], cached tokens per request; the block starts at this position config: Config RMS_EPS = 1e-5 FP8 = torch.float8_e4m3fn _NORMS = {} _ROPE = {} ROPE_POSITIONS = 1 << 21 def _qk_norm(hq, hkv, d): key = (hq, hkv, d) if key not in _NORMS: _NORMS[key] = FusedQKNorm(head_dim=d, num_q_heads=hq, num_kv_heads=hkv, eps=RMS_EPS) return _NORMS[key] def _cos_sin(theta, d, device): """[2^21, d] float32 cos | sin, as SGLang's RotaryEmbedding builds it, cached per theta.""" cache = _ROPE.get(theta) if cache is None: inv_freq = 1.0 / (theta ** (torch.arange(0, d, 2, dtype=torch.float, device=device) / d)) freqs = torch.einsum("i,j -> ij", torch.arange(ROPE_POSITIONS, dtype=torch.float, device=device), inv_freq) cache = _ROPE[theta] = torch.cat((freqs.cos(), freqs.sin()), dim=-1) return cache import triton import triton.language as tl @triton.jit def _prepare_qkv(QKV, Q, KC, VC, LOC, POS, CS, KS, VS, T: tl.constexpr, HQ: tl.constexpr, HK: tl.constexpr, D: tl.constexpr, WIDTH: tl.constexpr, ROPE: tl.constexpr, SCALED: tl.constexpr, BM: tl.constexpr): rows = tl.program_id(0) * BM + tl.arange(0, BM) heads: tl.constexpr = HQ + 2 * HK token, head = rows // heads, rows % heads col = tl.arange(0, D) mask = token[:, None] < T raw = tl.load(QKV + token[:, None] * WIDTH + head[:, None] * D + col[None, :], mask, 0) x = raw.to(tl.float32) inv = tl.rsqrt(tl.sum(x * x, 1) / D + 1.0e-5) y = (x * inv[:, None]).to(tl.bfloat16).to(tl.float32) if ROPE: partner = tl.load(QKV + token[:, None] * WIDTH + head[:, None] * D + (col[None, :] ^ 1), mask, 0).to(tl.float32) partner = (partner * inv[:, None]).to(tl.bfloat16).to(tl.float32) pos = tl.load(POS + token, token < T, 0).to(tl.int64) c = tl.load(CS + pos[:, None] * D + col[None, :] // 2, mask, 1) s = tl.load(CS + pos[:, None] * D + D // 2 + col[None, :] // 2, mask, 0) y = tl.where(col[None, :] % 2 == 0, y*c-partner*s, y*c+partner*s) y = y.to(tl.bfloat16).to(tl.float32) isq = head < HQ isk = (head >= HQ) & (head < HQ+HK) isv = head >= HQ+HK tl.store(Q + token[:, None]*HQ*D + head[:, None]*D + col[None, :], y, mask & isq[:, None]) key = y value = raw.to(tl.float32) if SCALED: key = (key / tl.load(KS)).to(tl.bfloat16).to(tl.float32) value = (value / tl.load(VS)).to(tl.bfloat16).to(tl.float32) slot = tl.load(LOC + token, token < T, 0).to(tl.int64) tl.store(KC + slot[:, None]*HK*D + (head[:, None]-HQ)*D + col[None, :], key, mask & isk[:, None]) tl.store(VC + slot[:, None]*HK*D + (head[:, None]-HQ-HK)*D + col[None, :], value, mask & isv[:, None]) def kernel_fn(data): cfg = data["config"] batch, length = cfg["batch"], cfg["q_len"] hq, hkv, d = cfg["num_q_heads"], cfg["num_kv_heads"], cfg["head_dim"] tokens = batch * length qkv = data["qkv"].view(tokens, -1) device = qkv.device k_cache, v_cache = data["k_cache"], data["v_cache"] scaled = cfg["kv_scale"] != 1.0 q = torch.empty((tokens, hq, d), device=device, dtype=FP8) cs = _cos_sin(cfg["rope_theta"], d, device) if cfg["rope_theta"] else q _prepare_qkv[(triton.cdiv(tokens * (hq + 2*hkv), 4),)]( qkv, q, k_cache, v_cache, data["cache_loc"], data["positions"], cs, data["k_scale"], data["v_scale"], tokens, hq, hkv, d, qkv.shape[1], bool(cfg["rope_theta"]), scaled, 4, num_warps=4) window = cfg["window"] # On H200, bound local-attention splits by batch/head parallelism. # Sliding layers have no incoming schedule; FA3 computes a compatible one. local_splits = max(1, min(16, triton.cdiv(132, batch * hkv))) if window else 0 descale = {} if scaled: descale = dict(k_descale=data["k_scale"].reshape(1, 1).expand(batch, hkv), v_descale=data["v_scale"].reshape(1, 1).expand(batch, hkv)) out = flash_attn_with_kvcache( q=q.view(tokens, hq, d).to(FP8), k_cache=k_cache, v_cache=v_cache, page_table=data["page_table"], cache_seqlens=data["cache_seqlens"], cu_seqlens_q=data["cu_seqlens_q"], max_seqlen_q=length, # The 5.473 query factor is folded into the softmax scale, as SGLang does. softmax_scale=cfg["sm_scale"] * cfg["qk_scale_factor"], causal=True, window_size=(window - 1, 0) if window else (-1, -1), scheduler_metadata=data["scheduler_metadata"], num_splits=local_splits, **descale, ) if isinstance(out, tuple): out = out[0] gate = qkv[:, (hq + 2 * hkv) * d:] o = rms_norm_gate(out.reshape(tokens, hq * d).to(torch.bfloat16), gate, None, head_dim=d, eps=RMS_EPS, has_norm=True) return o.view(batch, length, hq * d)