diff --git a/README.md b/README.md index 613f10e3..6684d7a8 100644 --- a/README.md +++ b/README.md @@ -59,9 +59,7 @@ SGLANG_USE_AITER=1 AITER_BF16_FP8_MOE_BOUND=0 python -m sglang.launch_server \ --monokernel-backend amd ``` -On gfx950 the decode graph capture fails in AITER's MoE without -`AITER_BF16_FP8_MOE_BOUND=0`. The flags are those of the serving tests -(`test/e2e/deepseek_v41_server.py`). +The flags are those of the serving tests (`test/e2e/deepseek_v41_server.py`). GLM-5.3, with MTP (its NextN layer drafts one chain; a verify step of 4 rows a request runs the kernels at batch sizes 1 and 2 at TP8): @@ -130,7 +128,8 @@ MXFP8 activations, as the kernels do. For MoRI PD, enable monokernels only on the decode server (`--disaggregation-mode decode --disaggregation-transfer-backend mori`): TP8 for DeepSeek-V4-Pro, TP2 or TP4 for DeepSeek-V4.1-Flash. Leave the flag off on prefill; -both peers must use the same KV layout. +both peers must use the same KV layout, `SGLANG_HIP_DSV41_INDEX_K_PAGE8` included +(SGLang's PD handshake does not compare it). If the model, GPU or configuration is not supported, the server refuses to start and lists every reason. diff --git a/sgl_amd_mono/deepseek_v41/gfx950/indexer/k2i.py b/sgl_amd_mono/deepseek_v41/gfx950/indexer/k2i.py index 3c67f160..e9cd7fbc 100644 --- a/sgl_amd_mono/deepseek_v41/gfx950/indexer/k2i.py +++ b/sgl_amd_mono/deepseek_v41/gfx950/indexer/k2i.py @@ -25,11 +25,11 @@ MAX_ROWS = 48 # a step's rows at most (8 DSpark requests of 6) -def build_k2i(tokens: int): - """K2i for a step of ``tokens`` rows.""" +def build_k2i(tokens: int, page8: bool = False): + """K2i for a step of ``tokens`` rows; ``page8``: on a page8 pool's plane (``score_fp4``).""" assert 0 < tokens <= MAX_ROWS, tokens - name = kernel_symbol("dsv41_mono_k2i", s=tokens) - keyed = (("tokens", tokens), ("sources", SOURCES)) + name = kernel_symbol("dsv41_mono_k2i", s=tokens, **({"p8": 1} if page8 else {})) + keyed = (("tokens", tokens), ("page8", page8), ("sources", SOURCES)) @fx.struct class RowsLds: @@ -55,7 +55,7 @@ def k2i( } # fmt: skip c = { "tid": tid, "lane": tid % 64, "wave": tid // 64, "args": args, - "rows_lds": lds.rows.ptr, + "rows_lds": lds.rows.ptr, "page8": page8, } # fmt: skip run_scores4(c, tokens, fx.block_idx.x) diff --git a/sgl_amd_mono/deepseek_v41/gfx950/indexer/runner.py b/sgl_amd_mono/deepseek_v41/gfx950/indexer/runner.py index 7bb5d925..df95b530 100644 --- a/sgl_amd_mono/deepseek_v41/gfx950/indexer/runner.py +++ b/sgl_amd_mono/deepseek_v41/gfx950/indexer/runner.py @@ -92,13 +92,14 @@ def __init__(self, device: torch.device, bound_max: int): ) self._kernels: dict = {} - def kernel(self, rows: int): + def kernel(self, rows: int, page8: bool = False): from ..runner import no_capture - if rows not in self._kernels: - no_capture(f"K2i width {rows}") - self._kernels[rows] = build_k2i(rows) - return self._kernels[rows] + key = (rows, page8) + if key not in self._kernels: + no_capture(f"K2i width {rows} (page8 {page8})") + self._kernels[key] = build_k2i(rows, page8) + return self._kernels[key] def walk_plan( self, *, ratio: int, bounds: torch.Tensor, rows: int, per_req: int @@ -121,7 +122,7 @@ def select( per_req: int, ratio: int, query: tuple, # [rows, 32, 64] E2M1 pairs, [rows, 1, 4, 16, 4] E8M0, [rows, 32] bf16 - plane: tuple, # the layer's index plane (``check_plane``) + plane: tuple, # the layer's index plane (``check_plane``), whether a page8 pool's page_table: torch.Tensor, # int32 [rows, pages] of 64-row pages bounds: torch.Tensor, # int32 [rows]: the columns each row sees raw_out: torch.Tensor, # int32 [rows, 512]: sparse_raw_indices @@ -235,9 +236,9 @@ def _score( ibound, lstride, reindex, ) -> None: # fmt: skip q_fp4, q_scale, weights = query - payload, scale = plane + payload, scale, page8 = plane _run_compiled( - self.kernel(rows), + self.kernel(rows, page8), q_fp4.data_ptr(), q_scale.data_ptr(), weights.data_ptr(), payload.data_ptr(), scale.data_ptr(), plan.data_ptr(), itab.data_ptr(), itab_stride, itab_len, ibatch.data_ptr(), ishift, ibound.data_ptr(), self.ilog.data_ptr(), lstride, reindex, diff --git a/sgl_amd_mono/deepseek_v41/gfx950/indexer/score_fp4.py b/sgl_amd_mono/deepseek_v41/gfx950/indexer/score_fp4.py index 42c517c8..56e84d44 100644 --- a/sgl_amd_mono/deepseek_v41/gfx950/indexer/score_fp4.py +++ b/sgl_amd_mono/deepseek_v41/gfx950/indexer/score_fp4.py @@ -11,6 +11,9 @@ iq [S, HEADS, DIM / 2] E2M1, iqs [S, 4 (K chunk), 16 (head % 16), 4 (head // 16, padded)] E8M0, iw [S, HEADS] bf16 (widening to f32 is exact) +A page8 pool (``SGLANG_HIP_DSV41_INDEX_K_PAGE8``) holds kv and ks in aiter's 8-row pages instead, +which are the virtual pages: only ``_keys`` reads the two layouts differently. + Lane (g, l) of tile nt holds key 4 l + nt, so a step's page is one lookup and its four tiles' scales one dword. The logits match aiter ``pa_mqa_logits_fp4_rowgroup`` only in its arithmetic order: relu as x + |x| against half the weight, the lane's heads in one fma chain, then the @@ -40,6 +43,7 @@ from . import index_plan as ip PAGE = ip.BLOCK_ROWS +# a page8 pool's page, in aiter's ``pack_kv_cache`` layout PAGE_BYTES = PAGE * ip.DIM // 2 PAGE_SCALES = PAGE * ip.DIM // 32 # SGLang's FP4 index plane (``deepseek_v4_memory_pool``): 64-row pages, payload and scales apart @@ -216,6 +220,13 @@ def _keys(c, page, nt): REINDEX row, whose request's other rows read the same candidate pages.""" a, lane = c["args"], c["lane"] l, g = lane % 16, lane // 16 + if const_expr(c["page8"]): + # payload [key % 4][chunk][key // 4][16 B], scales [chunk][key]: this lane's 4 (l % 2) + nt + page = fx.Int64(page) + base = a["plane"] + page * PAGE_BYTES + fx.Int64(g * 32 + (l % 2) * 16) + kv = [global_load(base + t * 128, 4, nt) for t in range(4)] + at = a["pscale"] + page * PAGE_SCALES + fx.Int64(g * 8 + (l % 2) * 4) + return kv, fx.Int32(global_load(at)) # ``page`` is virtual: 8 x SGLang's 64-row page + the 8-row sub page sub = page & 7 phys = fx.Int64(page >> 3) diff --git a/sgl_amd_mono/deepseek_v41/gfx950/runner.py b/sgl_amd_mono/deepseek_v41/gfx950/runner.py index 55610444..14d78b37 100644 --- a/sgl_amd_mono/deepseek_v41/gfx950/runner.py +++ b/sgl_amd_mono/deepseek_v41/gfx950/runner.py @@ -456,20 +456,19 @@ def precompile( k1=(), k2=(), ffn: bool = False, - k2i: bool = False, + k2i=(), tok_par: bool = False, ) -> None: """Allocate this width's scratch and compile its launches without running them. K1 for each (fold, index, iq) in ``k1``, K2 for each (compressed, index) in - ``k2``, the FFN launch, and K2i (``enable_k2i`` first).""" + ``k2``, the FFN launch, and K2i for each page8 in ``k2i`` (``enable_k2i`` first).""" self.scratch(tokens) kernels = [self.k1_kernel(tokens, *v, tok_par=tok_par) for v in sorted(k1)] kernels += [self.k2_kernel(tokens, *v, tok_par=tok_par) for v in sorted(k2)] if ffn: kernels.append(self.ffn_kernel(tokens)) - if k2i: - kernels.append(self.k2i.kernel(tokens)) + kernels += [self.k2i.kernel(tokens, page8) for page8 in sorted(k2i)] for launch in kernels: compile_only(launch) diff --git a/sgl_amd_mono/deepseek_v41/spec.py b/sgl_amd_mono/deepseek_v41/spec.py index 0853ad0d..4c4432ba 100644 --- a/sgl_amd_mono/deepseek_v41/spec.py +++ b/sgl_amd_mono/deepseek_v41/spec.py @@ -140,6 +140,7 @@ "disaggregation_transfer_backend", ), "srt/environ.py": ( + "envs", "SGLANG_DSV4_FP4_EXPERTS", "SGLANG_DSV4_FP4_DEQUANT", "SGLANG_USE_AITER_MOE_GU_ITLV", @@ -224,6 +225,7 @@ "get_index_k_fp4_payload_buffer", "get_index_k_fp4_scale_buffer", "low_ratio_index_k_is_split", + "low_ratio_index_k_page8", "get_index_k_page_size", ), "srt/model_executor/forward_context.py": ( @@ -564,7 +566,7 @@ class _Layer(msgspec.Struct): swa_block: int = 0 comp: Any = None # the compressed cache, the same layout comp_block: int = 0 - plane: Any = None # an indexer's FP4 index plane: (payload, scales) + plane: Any = None # an indexer's FP4 index plane: (payload, scales, page8) def _keep(copies: dict, t: torch.Tensor, *key) -> torch.Tensor: @@ -923,7 +925,7 @@ def _index_plane(*, L, M: int, meta, raw_idx) -> tuple: payload = pool.get_index_k_fp4_payload_buffer(L.i) scale = pool.get_index_k_fp4_scale_buffer(L.i) check_plane(payload=payload, scale=scale) - return payload.view(torch.uint8), scale + return payload.view(torch.uint8), scale, pool.low_ratio_index_k_page8(L.i) class _Built: @@ -1161,6 +1163,7 @@ def _longest_context(hf_config) -> int: def _build(model, hf_config, ctx) -> _Built: + from sglang.srt.environ import envs from sglang.srt.runtime_context import get_parallel from sgl_amd_mono.deepseek_v41.gfx950.runner import ( @@ -1185,7 +1188,14 @@ def _build(model, hf_config, ctx) -> _Built: shuffled=bool(shuffled_layout(L.module.self_attn for L in layers.values())), ) runner.set_deadline(ctx.deadline_s) - k2i = any(L.module.self_attn.indexer is not None for L in layers.values()) + # the plane layouts K2i reads: SGLang's page8 pool is ratio 1's unless aiter lacks the + # row-group kernel (each layer's plane picks its build, ``_index_plane``) + page8 = envs.SGLANG_HIP_DSV41_INDEX_K_PAGE8.get() + k2i = { + page8 and L.ratio == 1 + for L in layers.values() + if L.module.self_attn.indexer is not None + } if k2i: # the indexer's buffers, sized here: before SGLang sizes the KV cache runner.enable_k2i(_longest_context(hf_config)) @@ -1240,8 +1250,6 @@ def _build(model, hf_config, ctx) -> _Built: "SGLANG_DSV4_KV_LAYOUT": "v4", # K1 quantizes the SWA nope payload from f32, as SGLang's fused writer does "SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE": False, - # the scorer reads the ratio-1 index-K pages in the page-64 byte layout - "SGLANG_HIP_DSV41_INDEX_K_PAGE8": False, # benchmark routing overrides: K2 routes with the model's own router "SGLANG_SIMULATE_UNIFORM_EXPERTS": False, "SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS": False, diff --git a/test/kernels/test_deepseek_v41_k2i.py b/test/kernels/test_deepseek_v41_k2i.py index 897d55d7..faec5c22 100644 --- a/test/kernels/test_deepseek_v41_k2i.py +++ b/test/kernels/test_deepseek_v41_k2i.py @@ -1,9 +1,9 @@ """K2i, the index sources' FP4 scorer (RFC #43949), against SGLang's indexer ops on one gfx950. -On synthetic FP4 planes in SGLang's layout: the device walk plan against the host plan, the -logits against SGLang's and fp64, a candidate consumer's selection against SGLang's top-k within -the same candidate blocks, the two-level pipeline against SGLang's by per-row overlap, and a -graph replay against eager runs. +On synthetic FP4 planes in SGLang's layouts (64-row pages, and a page8 pool's aiter pages): the +device walk plan against the host plan, the logits against SGLang's and fp64, a candidate +consumer's selection against SGLang's top-k within the same candidate blocks, the two-level +pipeline against SGLang's by per-row overlap, and a graph replay against eager runs. Adapted from sgl-project/sglang#43984's test_dsv41_k2a_indexer.py.""" @@ -127,7 +127,24 @@ def fp4_case(*, ctxs, seed, per_req=PER_REQ, pages_per_req=None, dev="cuda"): device=dev, ) return dict(q=q, qscale=qscale, w=w, payload=payload, kscale=kscale, table=table, - bounds=bounds, per_req=per_req) # fmt: skip + bounds=bounds, per_req=per_req, page8=False) # fmt: skip + + +def as_page8(case): + """The case on a page8 pool: its keys in aiter's 8-row pages (``pack_kv_cache``, which + SGLang's page8 writer matches), in the same buffer shapes.""" + from aiter.ops.flydsl.kernels.mqa_logits.pa_mqa_logits_fp4_rowgroup import ( + pack_kv_cache, + ) + + payload, kscale = case["payload"], case["kscale"] + pages = payload.shape[0] + slot = torch.arange(64, device=payload.device) + rows = payload.view(pages, 4, 64, 16).transpose(1, 2).reshape(pages * 64, 64) + e8m0 = kscale.view(pages, 4, 64)[:, :, (slot % 16) * 4 + slot // 16].transpose(1, 2) + kv, ks = pack_kv_cache(rows, e8m0.reshape(pages * 64, 4), page_size=8) + return dict(case, payload=kv.reshape(payload.shape), kscale=ks.reshape(kscale.shape), + page8=True) # fmt: skip def _other_query(*, case, seed): @@ -147,7 +164,9 @@ def _unpack(b): def ref_logits(*, case, t, cols): - """fp64 logits of row t at columns ``cols``: sum_h w_h relu(sum_c 2^qe 2^ke q.k).""" + """fp64 logits of row t at columns ``cols``: sum_h w_h relu(sum_c 2^qe 2^ke q.k). + + ``case`` is on 64-row pages.""" page = case["table"][t][cols // 64].long() slot = cols % 64 k = _unpack(case["payload"][page, 0, :, slot, :]) # [n, 4, 32] @@ -178,7 +197,7 @@ def _inputs(case): rows=rows, per_req=case["per_req"], query=(case["q"], case["qscale"], case["w"]), - plane=(case["payload"], case["kscale"]), + plane=(case["payload"], case["kscale"], case["page8"]), page_table=case["table"], bounds=case["bounds"], ) @@ -214,21 +233,32 @@ def sglang_logits(case): q_fp4=case["q"], q_scale=case["qscale"], k_payload=case["payload"], k_scale=case["kscale"], weights=case["w"], page_table=case["table"], c4_seq_lens=case["bounds"], weight_scale=1.0, page_table_bucket=64, - is_decode=True, + is_decode=True, page8=case["page8"], ) # fmt: skip def sglang_select(*, case, logits, candidates=None): """SGLang's decode selection on ``logits``: the plain paged top-k and the - producer's candidate blocks, or a consumer's top-k inside ``candidates``. + producer's candidate blocks, or a consumer's top-k inside ``candidates`` (on a + page8 pool, from its own scores of the candidate blocks alone). -> (raw, page, the candidate blocks).""" from sglang.kernels.ops.attention.dsv4.candidate_blocks_hip import ( select_candidate_blocks_hip, topk_transform_paged_hip, topk_within_candidate_blocks_hip, + topk_within_candidate_blocks_page8, + with_candidate_page8_table, ) raw, page = _outs(case["q"].shape[0]) + if candidates is not None and case["page8"]: + topk_within_candidate_blocks_page8( + q_fp4=case["q"], q_scale=case["qscale"], weights=case["w"], + k_payload=case["payload"], k_scale=case["kscale"], candidates=candidates, + seq_lens=case["bounds"], page_table=case["table"], page_size=64, + page_indices=page, raw_indices=raw, + ) # fmt: skip + return raw, page, candidates if candidates is not None: topk_within_candidate_blocks_hip( logits=logits, seq_lens=case["bounds"], candidates=candidates, @@ -239,6 +269,10 @@ def sglang_select(*, case, logits, candidates=None): logits=logits, seq_lens=case["bounds"], topk_blocks=ip.TOPK_BLOCKS, block_size=ip.BLOCK_ROWS, ) # fmt: skip + if case["page8"]: + published = with_candidate_page8_table( + published, seq_lens=case["bounds"], page_table=case["table"], page_size=64 + ) topk_transform_paged_hip( scores=logits, seq_lens=case["bounds"], page_table=case["table"], page_indices=page, page_size=64, raw_indices=raw, @@ -264,17 +298,24 @@ def _jaccard(*, got, want): # ---------------------------------------------------------------- the launch +PAGE8 = pytest.mark.parametrize("page8", [False, True], ids=["page64", "page8"]) + + +@PAGE8 @pytest.mark.parametrize( "ctxs", [[8192], [300, 8192], [131072], [524288]], ids=["8k", "mix", "128k", "512k"] ) -def test_logits(ctxs): - """K2i's FULL logits of every scored row stay within ``k2i.logits.rel`` of SGLang's and fp64.""" +def test_logits(ctxs, page8): + """K2i's FULL logits of every scored row stay within ``k2i.logits.rel`` of SGLang's and fp64; + a page8 plane's are the same keys' on 64-row pages, bit for bit.""" k2i = _k2i() - case = fp4_case(ctxs=ctxs, seed=len(ctxs) + max(ctxs)) + case64 = fp4_case(ctxs=ctxs, seed=len(ctxs) + max(ctxs)) + case = as_page8(case64) if page8 else case64 k2i.walk_plan( ratio=1, bounds=case["bounds"], rows=case["q"].shape[0], per_req=PER_REQ ) - got = k2i.score(ratio=1, **_inputs(case)) + got = k2i.score(ratio=1, **_inputs(case)).clone() + same = k2i.score(ratio=1, **_inputs(case64)) torch.cuda.synchronize() want = sglang_logits(case) bound = dsv41.NUMERICS.bound("k2i.logits.rel") @@ -284,6 +325,7 @@ def test_logits(ctxs): continue a, s = got[t, :b], want[t, :b] assert torch.isfinite(a).all(), t + assert torch.equal(a, same[t, :b]), t err = float((a - s).abs().max() / s.abs().max().clamp_min(1e-30)) assert err <= bound, (t, err) cols = ( @@ -291,17 +333,20 @@ def test_logits(ctxs): if b > 8192 else torch.arange(b, device="cuda", dtype=torch.int32) ) - ref = ref_logits(case=case, t=t, cols=cols) + ref = ref_logits(case=case64, t=t, cols=cols) err64 = float((a[cols.long()].double() - ref).abs().max() / ref.abs().max()) assert err64 <= bound, (t, err64) +@PAGE8 @pytest.mark.parametrize("ctxs", [[20000, 9000], [131072]], ids=["20k9k", "128k"]) -def test_reindex_selects_sglangs_rows_within_the_candidates(ctxs): +def test_reindex_selects_sglangs_rows_within_the_candidates(ctxs, page8): """A REINDEX consumer, scoring its candidate columns alone, selects the rows SGLang's top-k - within the candidate blocks selects from the same layer's full logits.""" + within the candidate blocks selects from the same layer's full logits (on a page8 pool, from + aiter's scores of those columns).""" k2i = _k2i() case = fp4_case(ctxs=ctxs, seed=21 + max(ctxs)) + case = as_page8(case) if page8 else case rows = case["q"].shape[0] _, _, cands = run_k2i(k2i=k2i, case=case, kind="producer") consumer = _other_query(case=case, seed=6) @@ -315,14 +360,16 @@ def test_reindex_selects_sglangs_rows_within_the_candidates(ctxs): assert torch.equal(_sorted_sets(page), _sorted_sets(want_page)) +@PAGE8 @pytest.mark.parametrize( "ctxs", [[8192, 20000], [131072], [300, 524288]], ids=["8k20k", "128k", "512k"] ) -def test_pipeline_against_sglang(ctxs): +def test_pipeline_against_sglang(ctxs, page8): """A two-level step's producer rows, candidate blocks and consumer rows overlap SGLang's whole pipeline (on its own logits) per row.""" k2i = _k2i() case = fp4_case(ctxs=ctxs, seed=11 + max(ctxs)) + case = as_page8(case) if page8 else case raw, _, cands = run_k2i(k2i=k2i, case=case, kind="producer") want_raw, _, want_cands = sglang_select(case=case, logits=sglang_logits(case)) consumer = _other_query(case=case, seed=5) @@ -341,14 +388,18 @@ def test_pipeline_against_sglang(ctxs): assert med >= JACCARD_MEDIAN and p1 >= JACCARD_P1, (name, med, p1) -def test_graph_replay_selects_as_eager(): - """A captured step of the three kinds, replayed on other contexts and tables, selects as eager. +@PAGE8 +def test_graph_replay_selects_as_eager(page8): + """A captured step of the three kinds, replayed on other contexts and tables, selects as eager; + with page8 its ratio-1 sources (producer, reindex) read a page8 plane, as SGLang's pools do. Compared as sets: SGLang's top-k leaves a row's order to its atomics; eager runs differ too.""" k2i = _k2i() ppr = _cdiv(40000, 64) static = fp4_case(ctxs=[20000, 9000], seed=3, pages_per_req=ppr) - static2 = _other_query(case=static, seed=4) + # its own plane, the rest shared with ``static`` + ratio1 = as_page8(static) if page8 else static + static2 = _other_query(case=ratio1, seed=4) rows = static["q"].shape[0] outs = [_outs(rows) for _ in range(3)] held = {} @@ -356,7 +407,7 @@ def test_graph_replay_selects_as_eager(): def step(): run_k2i(k2i=k2i, case=static, kind="full", outs=outs[0], ratio=2) _, _, held["cands"] = run_k2i( - k2i=k2i, case=static, kind="producer", outs=outs[1] + k2i=k2i, case=ratio1, kind="producer", outs=outs[1] ) run_k2i( k2i=k2i, @@ -382,6 +433,10 @@ def results(): new = fp4_case(ctxs=ctxs, seed=seed, pages_per_req=ppr) for name in ("q", "qscale", "w", "payload", "kscale", "table", "bounds"): static[name].copy_(new[name]) + if page8: + new8 = as_page8(new) + ratio1["payload"].copy_(new8["payload"]) + ratio1["kscale"].copy_(new8["kscale"]) other = _other_query(case=new, seed=seed + 100) static2["q"].copy_(other["q"]) static2["w"].copy_(other["w"])