[AMD] gfx950 assembly attention: length-aware split-KV for dynamic workload

python/sglang/kernels/ops/attention/unified_attention_3d_mtp.py

@@ -52,15 +52,26 @@
     from sglang.kernels.ops.attention.vattn_asm_gfx950 import (
         mtp_verify_attn_fwd_asm as _mtp_verify_attn_fwd_asm,
     )
+    from sglang.kernels.ops.attention.vattn_asm_gfx950 import (
+        reset_seg_plan_cache as _reset_seg_plan_cache,
+    )
 except ImportError:
     _mtp_verify_attn_fwd_asm = None
     _AsmKernelUnavailable = RuntimeError
+    _reset_seg_plan_cache = None
 
 import os as _os
 
 from sglang.srt.utils import get_hip_version, is_gfx95_supported
 
 
+def reset_verify_attn_plan_cache() -> None:
+    """Drop the per-forward split-KV segment plan of the asm kernel;
+    the attention backend calls this at the start of every forward."""
+    if _reset_seg_plan_cache is not None:
+        _reset_seg_plan_cache()
+
+
 def asm_verify_attn_enabled() -> bool:
     """The in-tree gfx950 assembly attention kernel (vattn_asm_gfx950) is used
     by default on gfx950 with ROCm 7.2 or newer when ROCm clang is available and

python/sglang/kernels/ops/attention/vattn_asm_gfx950/__init__.py

@@ -61,6 +61,7 @@ class VattnKernelArgs(ctypes.Structure):
         ("_pad", ctypes.c_uint32),
         ("k_descale_ptr", ctypes.c_int64),
         ("v_descale_ptr", ctypes.c_int64),
+        ("seg_plan_ptr", ctypes.c_int64),  # 0 = legacy fixed-SEGS split
     ]
 
 
@@ -77,11 +78,12 @@ class VredKernelArgs(ctypes.Structure):
         ("out_stride1", ctypes.c_uint32),
         ("magic_m", ctypes.c_uint32),
         ("magic_sh", ctypes.c_uint32),
+        ("tok_nseg_ptr", ctypes.c_int64),  # 0 = legacy: reduce all num_segments
     ]
 
 
-assert ctypes.sizeof(VattnKernelArgs) == 128
-assert ctypes.sizeof(VredKernelArgs) == 56
+assert ctypes.sizeof(VattnKernelArgs) == 136
+assert ctypes.sizeof(VredKernelArgs) == 64
 
 
 def _declared_kernarg_size(source_file):
@@ -296,6 +298,174 @@ def mtp_verify_attn_num_segments(num_seqs: int, num_kv_heads: int) -> int:
     return max(1, min(64, segs))
 
 
+_SEG_PLAN_TARGET_WGS = None
+
+
+def _seg_plan_target_wgs() -> int:
+    global _SEG_PLAN_TARGET_WGS
+    if _SEG_PLAN_TARGET_WGS is None:
+        _SEG_PLAN_TARGET_WGS = max(
+            1,
+            torch.cuda.get_device_properties(
+                torch.cuda.current_device()
+            ).multi_processor_count,
+        )
+    return _SEG_PLAN_TARGET_WGS
+
+
+def mtp_verify_attn_seg_max(num_seqs: int, num_kv_heads: int) -> int:
+    """Static grid.x for the planned split: 2x the legacy per-seq count, clamped to 16..64."""
+    return max(16, min(64, 2 * mtp_verify_attn_num_segments(num_seqs, num_kv_heads)))
+
+
+def _get_plan_kernel():
+    kern = _kernels.get("plan")
+    if kern is None:
+        import triton
+        import triton.language as tl
+
+        @triton.jit
+        def _vattn_seg_plan_kernel(
+            seq_lens_ptr,
+            cu_q_ptr,
+            plan_ptr,
+            tok_nseg_ptr,
+            num_seqs,
+            num_work,
+            target_wgs,
+            seg_max,
+            BLOCK_B: tl.constexpr,
+            BLOCK_W: tl.constexpr,
+            BLOCK_Q: tl.constexpr,
+        ):
+            # one program per sequence; every program recomputes the (cheap) batch-wide plan
+            pid = tl.program_id(0)
+            b = tl.arange(0, BLOCK_B)
+            bm = b < num_seqs
+            slen = tl.load(seq_lens_ptr + b, mask=bm, other=0).to(tl.int32)
+            nt = (slen + 15) // 16
+            total = tl.sum(nt, axis=0)
+            mx = tl.max(nt, axis=0)
+            # smallest T (tiles per segment) with sum_b ceil(nt_b / T) <= target_wgs and max_b ceil(nt_b / T) <= seg_max
+            lo = tl.maximum(
+                tl.maximum(
+                    (total + target_wgs - 1) // target_wgs,
+                    (mx + seg_max - 1) // seg_max,
+                ),
+                1,
+            )
+            slack = target_wgs - num_seqs
+            hi = tl.where(
+                slack > 0,
+                (total + tl.maximum(slack, 1) - 1) // tl.maximum(slack, 1),
+                lo,
+            )
+            hi = tl.maximum(hi, lo)
+            for _ in range(16):
+                mid = (lo + hi) // 2
+                fits = tl.sum((nt + mid - 1) // mid, axis=0) <= target_wgs
+                hi = tl.where(fits, mid, hi)
+                lo = tl.where(fits, lo, mid + 1)
+            T = hi
+            nseg = (nt + T - 1) // T
+            ends = tl.cumsum(nseg, axis=0)
+            tot = tl.sum(nseg, axis=0)
+            my_n = tl.sum(tl.where(b == pid, nseg, 0), axis=0)
+            my_start = tl.sum(tl.where(b == pid, ends, 0), axis=0) - my_n
+            if pid == 0:
+                tl.store(plan_ptr, T)
+            w = tl.arange(0, BLOCK_W)
+            tl.store(plan_ptr + 1 + my_start + w, (pid << 16) | w, mask=w < my_n)
+            idle = tot + pid + w * num_seqs  # idle tail, strided over programs
+            tl.store(
+                plan_ptr + 1 + idle,
+                tl.full((BLOCK_W,), -1, tl.int32),
+                mask=idle < num_work,
+            )
+            q0 = tl.load(cu_q_ptr + pid).to(tl.int32)
+            q1 = tl.load(cu_q_ptr + pid + 1).to(tl.int32)
+            for t0 in range(q0, q1, BLOCK_Q):
+                t = t0 + tl.arange(0, BLOCK_Q)
+                tl.store(
+                    tok_nseg_ptr + t,
+                    tl.full((BLOCK_Q,), 0, tl.int32) + my_n,
+                    mask=t < q1,
+                )
+
+        kern = _vattn_seg_plan_kernel
+        _kernels["plan"] = kern
+    return kern
+
+
+def seg_plan_target_wgs(num_kv_heads: int) -> int:
+    """Working WGs to aim for: one per CU, shared over the kv-head grid dim."""
+    return max(1, _seg_plan_target_wgs() // max(1, num_kv_heads))
+
+
+def build_seg_plan(seq_lens, cu_seqlens_q, num_tokens, seg_max, num_kv_heads=1):
+    """plan int32[1 + seg_max*num_seqs] = (T tiles/segment, work list seq<<16|seg, -1 past the end),
+    tok_nseg int32[num_tokens] = segment count of the sequence owning each query token. One Triton launch,
+    static shapes, graph-capture safe. T is the smallest segment length whose total WG count fits the CU
+    budget, so uniform batches reproduce the legacy split exactly and skewed batches get per-length counts.
+    """
+    import triton
+
+    num_seqs = seq_lens.shape[0]
+    num_work = seg_max * num_seqs
+    plan = torch.empty(1 + num_work, dtype=torch.int32, device=seq_lens.device)
+    tok_nseg = torch.empty(
+        max(num_tokens, 1), dtype=torch.int32, device=seq_lens.device
+    )
+    _get_plan_kernel()[(num_seqs,)](
+        seq_lens,
+        cu_seqlens_q,
+        plan,
+        tok_nseg,
+        num_seqs,
+        num_work,
+        seg_plan_target_wgs(num_kv_heads),
+        seg_max,
+        BLOCK_B=max(16, triton.next_power_of_2(num_seqs)),
+        BLOCK_W=64,
+        BLOCK_Q=16,
+        num_warps=4,
+    )
+    return plan, tok_nseg
+
+
+_PLAN_CACHE = {}
+
+
+def reset_seg_plan_cache():
+    """Called by the attention backend at the start of every forward (eager and graph capture)."""
+    _PLAN_CACHE.clear()
+
+
+def _cached_seg_plan(seq_lens, cu_seqlens_q, num_tokens, seg_max, num_kv_heads):
+    # torch.cuda.is_current_stream_capturing() is part of the key: graph capture warms
+    # up and then records on the same tensors, and a plan built during warmup must not
+    # be reused while recording (its kernel would be missing from the graph).
+    key = (
+        seq_lens.data_ptr(),
+        cu_seqlens_q.data_ptr(),
+        seq_lens._version,
+        cu_seqlens_q._version,
+        num_tokens,
+        seq_lens.shape[0],
+        seg_max,
+        num_kv_heads,
+        torch.cuda.is_current_stream_capturing(),
+    )
+    hit = _PLAN_CACHE.get(key)
+    if hit is None:
+        plan, tok_nseg = build_seg_plan(
+            seq_lens, cu_seqlens_q, num_tokens, seg_max, num_kv_heads
+        )
+        # keep the key tensors alive so their storage cannot be reused under the same address while cached
+        hit = _PLAN_CACHE[key] = (plan, tok_nseg, seq_lens, cu_seqlens_q)
+    return hit[0], hit[1]
+
+
 def mtp_verify_attn_fwd_asm(
     q,
     k_cache,
@@ -308,13 +478,26 @@ def mtp_verify_attn_fwd_asm(
     softmax_scale,
     num_segments=None,
     out=None,
+    use_seg_plan=True,
 ):
-    """Same contract as aiter.mtp_verify_attn_fwd_asm (see that docstring)."""
+    """Same contract as aiter.mtp_verify_attn_fwd_asm (see that docstring).
+
+    use_seg_plan=False forces the fixed per-sequence split of #37465 (used by the tests as the
+    reference split); production callers leave it on."""
     num_tokens, num_q_heads, head_size = q.shape
     num_seqs = seq_lens.shape[0]
     num_kv_heads = k_cache.shape[2]
+    plan = tok_nseg = None
     if num_segments is None:
-        num_segments = mtp_verify_attn_num_segments(num_seqs, num_kv_heads)
+        if (
+            use_seg_plan and num_seqs > 1
+        ):  # bs=1: nothing to balance, the fixed split already uses 64 segments
+            num_segments = mtp_verify_attn_seg_max(num_seqs, num_kv_heads)
+            plan, tok_nseg = _cached_seg_plan(
+                seq_lens, cu_seqlens_q, num_tokens, num_segments, num_kv_heads
+            )
+        else:
+            num_segments = mtp_verify_attn_num_segments(num_seqs, num_kv_heads)
     segm_out = torch.empty(
         num_tokens,
         num_q_heads,
@@ -356,8 +539,15 @@ def mtp_verify_attn_fwd_asm(
         magic_sh=sh,
         k_descale_ptr=k_descale.data_ptr(),
         v_descale_ptr=v_descale.data_ptr(),
+        seg_plan_ptr=plan.data_ptr() if plan is not None else 0,
     )
-    kern.launch((num_segments, num_seqs, num_kv_heads), (512, 1, 1), args, stream)
+    if plan is not None:
+        # 1-D work list: working WGs first, idle tail exits in the prologue
+        kern.launch(
+            (num_segments * num_seqs, 1, num_kv_heads), (512, 1, 1), args, stream
+        )
+    else:
+        kern.launch((num_segments, num_seqs, num_kv_heads), (512, 1, 1), args, stream)
 
     assert out.stride(2) == 1 and out.dtype == torch.bfloat16
     rm, rsh = _magic_u32(num_q_heads)
@@ -372,6 +562,7 @@ def mtp_verify_attn_fwd_asm(
         out_stride1=out.stride(1),
         magic_m=rm,
         magic_sh=rsh,
+        tok_nseg_ptr=tok_nseg.data_ptr() if tok_nseg is not None else 0,
     )
     _get_reduce().launch((num_tokens * num_q_heads, 1, 1), (256, 1, 1), rargs, stream)
     return out

python/sglang/kernels/ops/attention/vattn_asm_gfx950/vattn3_core.s

@@ -133,6 +133,7 @@
 .set sPhysC,   71
 .set sPBrow,   72       // hkv*256 (token stride inside a page)
 .set sKvhOff,  73       // kvh*256
+.set sPlan,    74       // s74:75 segment plan ptr (0 = legacy fixed-SEGS split)
 
 // ---------------- macros ----------------
 
@@ -401,7 +402,22 @@ vattn_asm:
     s_load_dwordx8  s[24:31], s[0:1], 0x40
     s_load_dwordx4  s[32:35], s[0:1], 0x60
     s_load_dwordx4  s[56:59], s[0:1], 0x70
+    s_load_dwordx2  s[sPlan:sPlan+1], s[0:1], 0x80
     s_waitcnt lgkmcnt(0)
+    s_cmp_eq_u64 s[sPlan:sPlan+1], 0
+    s_cbranch_scc1 L_PLAN_DONE
+    // segment plan (1-D grid): plan[0] = tiles per segment, plan[1+wg_x] = seq<<16 | seg, or -1 past the
+    // end of the work list. Idle WGs sit at the tail of the grid and exit before touching memory.
+    s_load_dword s[sTps], s[sPlan:sPlan+1], 0x0
+    s_lshl_b32 s[sT2], s[sSeg], 2
+    s_add_i32 s[sT2], s[sT2], 4
+    s_load_dword s[sT3], s[sPlan:sPlan+1], s[sT2]
+    s_waitcnt lgkmcnt(0)
+    s_cmp_lt_i32 s[sT3], 0
+    s_cbranch_scc1 L_EXIT
+    s_lshr_b32 s[sSeq], s[sT3], 16
+    s_and_b32 s[sSeg], s[sT3], 0xffff
+L_PLAN_DONE:
     s_load_dword s[sT2], s[sKd:sKd+1], 0x0
     s_load_dword s[sT], s[sVd:sVd+1], 0x0
     s_waitcnt lgkmcnt(0)
@@ -428,11 +444,14 @@ vattn_asm:
     // num_tiles, tps, pg0/pg1
     s_add_i32 s[sNt], s[sSlen], 15
     s_lshr_b32 s[sNt], s[sNt], 4
+    s_cmp_lg_u64 s[sPlan:sPlan+1], 0
+    s_cbranch_scc1 L_TPS_DONE                   // planned: sTps already holds T
     s_lshl_b32 s[sT2], s[sSEGS], 4
     s_add_i32 s[sT2], s[sT2], -1
     s_add_i32 s[sT2], s[sSlen], s[sT2]
     s_mul_hi_u32 s[sTps], s[sT2], s[sMagic]
     s_lshr_b32 s[sTps], s[sTps], s[sShift]
+L_TPS_DONE:
     s_mul_i32 s[sPg0], s[sSeg], s[sTps]
     s_add_i32 s[sPg1], s[sPg0], s[sTps]
     s_min_i32 s[sPg1], s[sPg1], s[sNt]
@@ -725,7 +744,7 @@ L_EXIT:
 .amdhsa_kernel vattn_asm
     .amdhsa_group_segment_fixed_size LDS_TOTAL
     .amdhsa_private_segment_fixed_size 0
-    .amdhsa_kernarg_size 128
+    .amdhsa_kernarg_size 136
     .amdhsa_user_sgpr_count 2
     .amdhsa_user_sgpr_kernarg_segment_ptr 1
     .amdhsa_system_sgpr_workgroup_id_x 1
@@ -751,7 +770,7 @@ amdhsa.target: amdgcn-amd-amdhsa--gfx950
 amdhsa.kernels:
   - .name: vattn_asm
     .symbol: vattn_asm.kd
-    .kernarg_segment_size: 128
+    .kernarg_segment_size: 136
     .kernarg_segment_align: 8
     .group_segment_fixed_size: 148480
     .private_segment_fixed_size: 0
@@ -781,5 +800,6 @@ amdhsa.kernels:
       - {.offset: 104, .size: 4, .value_kind: by_value}
       - {.address_space: global, .offset: 112, .size: 8, .value_kind: global_buffer}
       - {.address_space: global, .offset: 120, .size: 8, .value_kind: global_buffer}
+      - {.address_space: global, .offset: 128, .size: 8, .value_kind: global_buffer}
 ...
 .end_amdgpu_metadata