diff --git a/MEASUREMENTS.md b/MEASUREMENTS.md index e02aafc5..420137de 100644 --- a/MEASUREMENTS.md +++ b/MEASUREMENTS.md @@ -2121,3 +2121,31 @@ will be disabled"). 업스트림 main 도 동일(2026-09-02). 그런데 `DSparkS **미측정**: EXP-8(부팅 1회 + 브래킷). 수용률은 움직이면 안 된다(실행 시점만 바뀜). V2+async 는 `max_concurrent_batches`=2 라 KV in-flight 예약이 두 배 — KV 라인 확인. + +## ★인덱서 fp32 head-gate: cuBLAS 2블록 47 us → split-K 10 us (`glm53_indexer_gate_splitk`, 2026-09-02, srv4 오프라인) + +**대상**: `attention.py` `Indexer.forward` 의 `torch.mm(hidden_states.float(), self._wp_fp32)` +— [M,4096]×[4096,16] fp32, 층당 1회 × 11층. fp32 인 이유는 코드 주석대로 bf16 게이트 +(~1e-2)가 근소 차 풀 순위를 뒤집기 때문. cuBLAS 는 이 형상에 `gemmSN` 2블록 커널을 +고른다(48 SM 중 2개). + +**실측**(GB10, CUDA 그래프 리플레이, `probes/indexer_gate_check.py`; 9월 1일 트레이스의 +86 us 는 CUPTI + 공유 경합): + +| M | stock `torch.mm` | split-K(8) | 경로 | +|---|---|---|---| +| 1 | 15.5 us | 8.5 us | split-K | +| 8 (C=1) | 50.0 us | 9.9 us | split-K | +| 16 (C=2) | 50.1 us | 11.7 us | split-K | +| 32 (C=4) | 15.9 us | 17.4 us | `torch.mm` 유지 | + +M=32 부터 cuBLAS 가 다른 커널을 고르며 빨라지므로 M<=16 만 라우팅. 대안 비교(M=8): +`F.linear` NT 레이아웃 32 us, mul+sum 20.7 us, split-K(16) 6.2 us — 8 분할이 원자 경합과 +프로그램 수의 균형점이라 채택. + +**수치**: 양쪽 fp32 누적, 합산 순서만 다름. 300회/2,480행 max|diff| 2.4e-6 절대, +6.7e-7/행 최대, top-1 뒤집힘 0, top-4 집합 변화 0. bit-exact 아님 → 품질 브래킷 대상. + +**천장**: 11 × 40 us ≈ 0.44 ms/스텝 = C=1 스텝의 ~0.65% (<1%, 단독 부팅 불가). EXP-7/8 +부팅에 얹어 잰다. `VLLM_GLM53_FUSED_K_GATE=1` 팔의 융합 인덱서도 같은 헬퍼를 타므로 +두 팔이 어긋나지 않는다. 기본 0 = stock 과 같은 `torch.mm` 호출. RUNBOOK EXP-9. diff --git a/RUNBOOK_KERNEL_CAMPAIGN2.md b/RUNBOOK_KERNEL_CAMPAIGN2.md index db5553ab..fc01bdd6 100644 --- a/RUNBOOK_KERNEL_CAMPAIGN2.md +++ b/RUNBOOK_KERNEL_CAMPAIGN2.md @@ -345,6 +345,33 @@ VLLM_GLM53_ASYNC_DFLASH=1 bash launchers/start-glm53-nvfp4-tp4.sh # cand (프 dflash 의 async off 모두 손상, 원인은 LibertAIDAI 가중치). 이 실험은 그 판정을 재론하지 않는다. +## EXP-9 — 인덱서 fp32 head-gate 를 split-K 로 (`glm53_indexer_gate_splitk`, 2026-09-02 추가) + +`Indexer.forward` 의 `weights = torch.mm(hidden_states.float(), self._wp_fp32)` +([M,4096]×[4096,16] fp32, 층당 1회 × 11) 을 cuBLAS 가 2블록 `gemmSN` 커널로 답한다: +유휴 GB10 에서 47 us, 9월 1일 트레이스(CUPTI)에서 86 us. 48 SM 에 블록 2개가 문제의 +전부라 (행, K-슬라이스 512) 마다 프로그램 하나를 띄우고 fp32 atomic 으로 모으는 +split-K Triton 커널이 같은 곱을 10 us 에 낸다. **M<=16(디코드) 만** 이 경로, +나머지(프리필, C>=3 verify)는 stock `torch.mm` 그대로. + +**수치**: 양쪽 다 fp32 누적, 합산 순서만 다르다 — bit-exact 가 아니다. 오프라인 +300회/2,480행: max|diff| 2.4e-6, 행 최대 대비 6.7e-7, top-1 뒤집힘 0, top-4 집합 변화 0 +(`probes/indexer_gate_check.py`). bf16 게이트가 순위를 뒤집는 오차(1e-2)보다 네 자릿수 +아래지만 품질 브래킷은 필요하다. + +**천장**: 11층 × ~40 us ≈ 0.44 ms/스텝, C=1 66 ms 스텝의 **~0.65%** — 원장 규칙(1% 미만은 +단독 부팅 가치 없음)에 걸린다. EXP-7/EXP-8 과 독립이라 그 부팅에 얹어 같이 재는 +용도. 단독 부팅 금지. + +```bash +# 프로필 선언 키: caller env. 트레이스에서 gemmSN 11개가 _gate_splitk_kernel 11개로 +# 바뀐 것이 켜진 증거 (부팅 로그 줄은 없다 — 그래프 안에서 층마다 호출). +VLLM_GLM53_INDEXER_GATE_SPLITK=1 bash launchers/start-glm53-nvfp4-tp4.sh +``` + +- 게이트: 품질 9/9, 한국어 0/16, C=1 step/s 브래킷 base→cand→base (얹은 부팅의 것과 공유). +- 롤백 = env 한 줄. 기본 0 = stock 과 동일한 `torch.mm` 호출. + --- ## 순서와 근거 diff --git a/STEP_KERNEL_MAP.md b/STEP_KERNEL_MAP.md index 0ed5418e..1f6baec0 100644 --- a/STEP_KERNEL_MAP.md +++ b/STEP_KERNEL_MAP.md @@ -256,7 +256,8 @@ mamba 4, 드래프터 1)의 빌더 — GDN 빌더 4개가 각각 `to`·`sub`·`a **읽다가 나온 것**: (1) 인덱서의 fp32 head-gate `torch.mm(hidden.float(), _wp_fp32)` 가 cuBLAS gemmSN 2블록 커널로 층당 86 us, 11층 0.95 ms/스텝(CUPTI) — 우리 -`glm53_prefill_fastpath.py:402` 소유, split-K 로 수 us 감. (2) 드래프터 fc 투영 +`glm53_prefill_fastpath.py:402` 소유, split-K 로 수 us 감 → `glm53_indexer_gate_splitk` +(EXP-9, opt-in, 오프라인 50 → 10 us, ~0.65%/스텝). (2) 드래프터 fc 투영 814 us(eager bf16, 5층 hidden cat, 168 MB 읽기, `ReplicatedLinear` 라 fp8 dense 패턴 밖) 포함 드래프터 커널 합 ~3 ms(CUPTI) — 원장 D≈0 과 긴장, 직접 측정 전 판단 보류. (3) `KpoolTailMetadataBuilder` 의 원형 tail 슬롯 매핑은 러너가 diff --git a/build/glm53/glm53_indexer_gate.py b/build/glm53/glm53_indexer_gate.py new file mode 100644 index 00000000..cd73f8e9 --- /dev/null +++ b/build/glm53/glm53_indexer_gate.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +"""deneb fork: the sparse indexer's fp32 head-gate projection as a split-K +Triton kernel (VLLM_GLM53_INDEXER_GATE_SPLITK). + +`Indexer.forward` computes `weights = torch.mm(hidden_states.float(), +self._wp_fp32)` -- an [M, 4096] x [4096, 16] fp32 product per full-attention +layer, kept in fp32 on purpose (bf16 head-gates flip near-tie pool rankings). +cuBLAS answers that shape with a two-block `gemmSN` kernel: 47 us on an idle +GB10 (86 us under CUPTI in the 2026-09-01 serving trace), eleven times per +decode step, for 256 KB of weights. A split-K kernel that hands each +(row, K-slice) to one program and reduces with fp32 atomics runs the same +product in 7 us. Both accumulate in fp32; only the summation order differs +(measured max |diff| 3e-5 on values of magnitude ~64, i.e. ~5e-7 relative, +0 top-1 rank flips over the offline trials) -- not bit-exact, so this stays +an opt-in behind a numerics bracket. + +At M > 16 (C >= 3 verify batches, prefill) cuBLAS is already fast, so the +kernel is used only for M <= 16; larger M keeps torch.mm. +""" +from __future__ import annotations + +import os + +import torch + +from vllm.triton_utils import tl, triton + +ENV = "VLLM_GLM53_INDEXER_GATE_SPLITK" +MAX_M = 16 +_SPLIT = 8 +_BLOCK_K = 128 + + +def gate_splitk_enabled() -> bool: + """Exact opt-in: only the string "1" arms; anything else is stock.""" + return os.environ.get(ENV, "").strip() == "1" + + +@triton.jit +def _gate_splitk_kernel(x_ptr, w_ptr, out_ptr, K, N, sxm, swk, som, + SPLIT: tl.constexpr, BLOCK_K: tl.constexpr, BN: tl.constexpr): + m = tl.program_id(0) + s = tl.program_id(1) + kper = K // SPLIT + k0 = s * kper + offs_n = tl.arange(0, BN) + nmask = offs_n < N + acc = tl.zeros([BN], dtype=tl.float32) + for k in range(0, kper, BLOCK_K): + offs_k = k0 + k + tl.arange(0, BLOCK_K) + xv = tl.load(x_ptr + m * sxm + offs_k) + wv = tl.load(w_ptr + offs_k[:, None] * swk + offs_n[None, :], mask=nmask[None, :], other=0.0) + acc += tl.sum(xv[:, None] * wv, axis=0) + tl.atomic_add(out_ptr + m * som + offs_n, acc, mask=nmask) + + +def head_gate_splitk(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """fp32 [M, K] @ [K, N] for N <= 16, K a multiple of SPLIT*BLOCK_K. + + x may be bf16 (cast to fp32 here, as the stock `.float()` does); w is the + stock `_wp_fp32` ([K, N], fp32, contiguous).""" + xf = x.float() + M, K = xf.shape + N = w.shape[1] + out = torch.zeros(M, N, device=x.device, dtype=torch.float32) + _gate_splitk_kernel[(M, _SPLIT)]( + xf, w, out, K, N, xf.stride(0), w.stride(0), out.stride(0), + SPLIT=_SPLIT, BLOCK_K=_BLOCK_K, BN=16, num_warps=4) + return out + + +def splitk_applicable(x: torch.Tensor, w: torch.Tensor) -> bool: + return (x.shape[0] <= MAX_M and w.dtype == torch.float32 and w.is_contiguous() + and w.shape[1] <= 16 and w.shape[0] % (_SPLIT * _BLOCK_K) == 0) + + +def head_gate(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """The stock `torch.mm(x.float(), w)` unless the knob is on and the shape + is the small-M decode one.""" + if gate_splitk_enabled() and splitk_applicable(x, w): + return head_gate_splitk(x, w) + return torch.mm(x.float(), w) diff --git a/build/glm53/glm53_prefill_fastpath.py b/build/glm53/glm53_prefill_fastpath.py index 8b2eaadd..e540490c 100644 --- a/build/glm53/glm53_prefill_fastpath.py +++ b/build/glm53/glm53_prefill_fastpath.py @@ -34,6 +34,23 @@ ) from .ops.kpool_compress import fwht128_quant_fp8 +# deneb fork (glm53_indexer_gate_splitk): the split-K helper when that module +# is mounted, else the stock fp32 torch.mm. Resolved once, on first call, so +# the fused indexer forward stays loadable without the sibling file. +_HEAD_GATE = None + + +def _glm53_head_gate(x, w): + global _HEAD_GATE + if _HEAD_GATE is None: + try: + from vllm.models.glm5next.nvidia.glm53_indexer_gate import head_gate as fn + except ImportError: + def fn(x, w): + return torch.mm(x.float(), w) + _HEAD_GATE = fn + return _HEAD_GATE(x, w) + logger = init_logger(__name__) _GLM53_SM121_MLA_PREFILL_ENV = "VLLM_GLM53_SM121_MLA_PREFILL" @@ -399,7 +416,7 @@ def _glm53_fused_indexer_forward( self._wp_fp32 = ( k_weight.data[self.head_dim :, :].t().contiguous().float() ) - weights = torch.mm(hidden_states.float(), self._wp_fp32) + weights = _glm53_head_gate(hidden_states, self._wp_fp32) # deneb fork (glm53_indexer_gate_splitk) k = _fused_indexer_k_norm( k, diff --git a/build/glm53/glm5next_attention.py b/build/glm53/glm5next_attention.py new file mode 100644 index 00000000..ee1dcf8a --- /dev/null +++ b/build/glm53/glm5next_attention.py @@ -0,0 +1,627 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.config import ( + CacheConfig, + VllmConfig, +) +from vllm.distributed import ( + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope +from vllm.model_executor.layers.sparse_attn_indexer_kpool import SparseAttnIndexerKpool +from vllm.model_executor.models.deepseek_v2 import ( + DeepSeekV2FusedQkvAProjLinear, + DeepseekV32IndexerCache, + yarn_get_mscale, +) +from vllm.model_executor.models.utils import extract_layer_index +from vllm.model_executor.utils import maybe_disable_graph_partition +from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8 +from vllm.models.glm5next.nvidia.glm53_indexer_gate import head_gate as _glm53_head_gate +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig +from vllm.v1.kv_cache_interface import KpoolTailSpec, MLAAttentionSpec + +logger = init_logger(__name__) + +# Shared torch.compile config for the indexer's small-kernel leaves. The MLA +# indexer runs under breakable-CG (CompilationMode.NONE), which blocks FX-graph +# fusion of the surrounding eager ops; carving each cluster into its own +# @torch.compile leaf (backend==inductor) still fuses them. Matches the +# grouped_topk / _cast_sigmoid leaf pattern. +_INDEXER_COMPILE = dict( + dynamic=True, + backend=current_platform.simple_compile_backend, + options=maybe_disable_graph_partition(current_platform.simple_compile_backend), +) + + +@torch.compile(**_INDEXER_COMPILE) +def _fused_indexer_k_norm( + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, dim: int, eps: float +) -> torch.Tensor: + # Fuse fp32 cast + layer_norm + cast-back (was 3 kernels) into one. + return F.layer_norm(x.float(), (dim,), weight, bias, eps).type_as(x) + + +@torch.compile(**_INDEXER_COMPILE) +def _fused_indexer_weight_scale( + weights: torch.Tensor, q_scale: torch.Tensor, scale: float +) -> torch.Tensor: + # Fuse the weight-scaling muls (was 2 kernels) into one. `scale` folds + # softmax_scale (head_dim**-0.5) and n_head**-0.5 into a single constant. + return (weights.unsqueeze(-1) * q_scale * scale).squeeze(-1) + + +@torch.compile(**_INDEXER_COMPILE) +def _pad_indexer_heads(x: torch.Tensor, pad: int) -> torch.Tensor: + # DeepGEMM MQA-logits needs num_heads in {32,64}; zero-pad the head dim. + # Fuse new_zeros + cat (was 2 kernels) into one. Pad values are zero (exact + # in fp8 e4m3 and zero-weight in the logits sum), so numerically a no-op. + return torch.cat([x, x.new_zeros(x.shape[0], pad, *x.shape[2:])], dim=1) + + +class Glm5NextIndexerCache(DeepseekV32IndexerCache): + """Indexer K cache that stores kpool-compressed entries. + + Setting ``compress_ratio = index_kpool`` on the kv_cache_spec makes vLLM's + indexer metadata builder emit pool-granular ``slot_mapping`` / + ``seq_lens`` / ``cu_seq_lens`` / ``page_table`` for free, and shrinks the + cache allocation to ``storage_block_size = block_size // kpool``. The pool + *content* (softmax-weighted sum vs keep-every-Nth) is computed by the + kpool compress kernel inside the indexer op — the cache only provides the + addressing, which is identical for both schemes. + + The indexer shares one block with the co-located MLA (a single + ``MLAAttentionSpec`` / block_table), so ``block_size`` is the model-wide + ``cache_config.block_size``. DeepGEMM's paged-MQA kernel + (``csrc/apis/attention.hpp``) requires ``block_kv`` to be exactly 32 or + 64, so the storage block is virtually split into pool pages of the + largest such size that tiles it (``storage_kernel_block_size``); this + needs ``block_size`` to be a multiple of ``index_kpool * 32`` (512 for + ``index_kpool = 16``). A smaller block (e.g. the default 64) silently + collapses ``storage_block_size`` (64 // 16 = 4) and only fails later at + the opaque C++ assert; ``get_kv_cache_spec`` guards this up front + instead. + """ + + def __init__( + self, + *, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config, + index_kpool: int, + ): + super().__init__( + head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config + ) + assert index_kpool > 1, "Glm5NextIndexerCache expects index_kpool > 1" + # Chunked prefill aligns chunk ends to ``cache_config.block_size``; + # the prefill compress kernel (``_kpool_compress_insert``) assumes + # pool-aligned chunk starts, which holds only if that block_size is a + # multiple of kpool -- otherwise every pool straddling a chunk + # boundary is silently dropped (masked off as a leading-partial pool). + # Distinct from the get_kv_cache_spec assert, which guards the spec's + # DeepGEMM tiling on its own block_size. + assert cache_config.block_size % index_kpool == 0, ( + "Glm5NextIndexerCache: cache_config.block_size " + f"({cache_config.block_size}) must be a multiple of index_kpool " + f"({index_kpool}) so chunked-prefill boundaries stay pool-aligned." + ) + self._index_kpool = index_kpool + + def get_kv_cache_spec(self, vllm_config: VllmConfig): + from dataclasses import replace + + spec = super().get_kv_cache_spec(vllm_config) + # compress_ratio lives on MLAAttentionSpec, but the base + # DeepseekV32IndexerCache.get_kv_cache_spec is typed to return the + # KVCacheSpec base; narrow so dataclass.replace sees the field. + assert isinstance(spec, MLAAttentionSpec) + spec = replace(spec, compress_ratio=self._index_kpool) + + # DeepGEMM paged-MQA takes block_kv in {32, 64}; the storage block + # (= block_size // index_kpool) is virtually split into pool pages of + # the largest such size that tiles it, so it must be a multiple of 32. + storage_block_size = spec.block_size // self._index_kpool + assert ( + spec.block_size % self._index_kpool == 0 and storage_block_size % 32 == 0 + ), ( + "Glm5NextIndexerCache: kpool indexer requires cache block_size to " + f"be a multiple of index_kpool * 32 ({self._index_kpool * 32}) so " + "that DeepGEMM paged-MQA pool pages (32 or 64 entries) tile the " + f"storage block, got block_size={spec.block_size} -> " + f"storage_block_size={storage_block_size}." + ) + return spec + + +class Glm5NextTailCache(DeepseekV32IndexerCache): + """Paged circular buffer for the kpool indexer's in-progress (tail) pool. + + Holds the trailing incomplete pool's raw K + gate score: one block of + ``index_kpool`` slots per request, overwritten in place by ``pos % kpool`` + as decode/spec-decode advances. Prefill seeds it (instead of discarding the + tail raw K+gate); the connector transfers it across PD; decode reads it to + compress the boundary pool correctly. ``KpoolTailSpec`` / + ``KpoolTailManager`` provide the no-prune, 1-block/req allocation that lets + the in-progress pool survive across steps and across transfer. + + Stores raw bf16 K (``head_dim``) as the "K" half of each block and the + bf16 gate score (``head_dim``) as the "V" half -- not the fp8-compressed + entry, which lives in ``Glm5NextIndexerCache``. + """ + + def __init__( + self, + *, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config, + index_kpool: int, + ): + super().__init__( + head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config + ) + assert index_kpool > 1, "Glm5NextTailCache expects index_kpool > 1" + self._index_kpool = index_kpool + + def get_kv_cache_spec(self, vllm_config: VllmConfig): + # K + gate score packed into head_size (== 2*head_dim), head_size_v=0: + # KpoolTailBackend.get_kv_cache_shape only consumes head_size and splits + # it into [2, kpool, head_dim] (K | score halves), so the connectors' + # non-MLA K/V half-split transfers K and score as separate halves. + return KpoolTailSpec( + block_size=self._index_kpool, + num_kv_heads=1, + head_size=2 * self.head_dim, + head_size_v=0, + dtype=torch.bfloat16, + sliding_window=self._index_kpool, + ) + + def get_attn_backend(self): + from vllm.v1.attention.backends.mla.indexer import KpoolTailBackend + + return KpoolTailBackend + + +class Indexer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: Glm5NextConfig, + hidden_size: int, + q_lora_rank: int, + quant_config: QuantizationConfig | None, + cache_config: CacheConfig | None, + topk_indices_buffer: torch.Tensor | None, + prefix: str = "", + ): + super().__init__() + self.vllm_config = vllm_config + self.config = config + self.quant_config = quant_config + # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"] + # Indexer is only constructed for v32 configs, where these sparse-indexer + # fields are guaranteed populated; narrow away the `int | None` declared + # on Glm5NextConfig for the optional-indexer case. + assert config.index_topk is not None + assert config.index_n_heads is not None + assert config.index_head_dim is not None + assert config.index_kpool is not None + self.topk_tokens = config.index_topk + self.n_head = config.index_n_heads # 64 + self.head_dim = config.index_head_dim # 128 + self.rope_dim = config.qk_rope_head_dim # 64 + self.index_kpool = config.index_kpool + self.q_lora_rank = q_lora_rank # 1536 + + # kpool + self.index_kpool_compress_ape = nn.Parameter( + torch.zeros(self.index_kpool, self.head_dim, dtype=torch.float32) + ) + # NOTE: kept as a bare nn.Parameter (not ReplicatedLinear) so the weight + # name matches the checkpoint verbatim ("index_kpool_compress_gate", + # no ".weight" suffix) — the trained checkpoint stores it the sglang way. + # Shape [head_dim, hidden_size]; consumed via F.linear(x, gate) = x @ gate.T. + self.index_kpool_compress_gate = nn.Parameter( + torch.empty(self.head_dim, hidden_size, dtype=torch.bfloat16) + ) + + # no tensor parallel, just replicated + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.head_dim * self.n_head, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wq_b", + ) + # Fused wk + weights_proj: single GEMM producing [head_dim + n_head]. + # FP8 wk weights are upcasted to BF16 during loading to maintain fusion. + self.wk_weights_proj = MergedColumnParallelLinear( + hidden_size, + [self.head_dim, self.n_head], + bias=False, + quant_config=None, + disable_tp=True, + prefix=f"{prefix}.wk_weights_proj", + ) + self.k_norm = LayerNorm(self.head_dim, eps=1e-6) + self.softmax_scale = self.head_dim**-0.5 + + # Hadamard-128 rotation of the indexer query is fused with the FP8 + # quant (see forward: fwht128_quant_fp8) -- no precomputed matrix. + + self.scale_fmt = "ue8m0" + self.quant_block_size = 128 # TODO: get from config + self.topk_indices_buffer = topk_indices_buffer + + # NOTE: (zyongye) we use fp8 naive cache, + # where we store value in fp8 and scale in fp32 + # per self.quant_block_size element + self.k_cache = Glm5NextIndexerCache( + head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4, + dtype=torch.uint8, + prefix=f"{prefix}.k_cache", + cache_config=cache_config, + index_kpool=self.index_kpool, + ) + # Paged tail cache (in-progress pool's raw K + gate score). Written by + # prefill (seeds the boundary pool) and decode (per-step stash); read by + # the decode kernel to compress the boundary pool. Transferred across PD + # so the decode side sees the prefill tail. See KpoolTailSpec/Manager. + self.tail_cache = Glm5NextTailCache( + head_dim=self.head_dim, + dtype=torch.bfloat16, + prefix=f"{prefix}.tail_cache", + cache_config=cache_config, + index_kpool=self.index_kpool, + ) + self.max_model_len = vllm_config.model_config.max_model_len + self.prefix = prefix + from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size + + self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config) + self.indexer_op = SparseAttnIndexerKpool( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + tail_cache=self.tail_cache, + ) + + def forward( + self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb + ) -> torch.Tensor: + q, _ = self.wq_b(qr) + q = q.view(-1, self.n_head, self.head_dim) + + # Fused wk + weights_proj: one GEMM for wk (bf16). The head-gate + # (weights_proj) is computed separately in fp32 to match sglang + # (ReplicatedLinear params_dtype=fp32, x.float()) -- bf16 head-gates + # introduce ~1e-2 logit error that flips near-tie pool rankings on hard + # long-context tasks (single-needle retrieval is unaffected, but HLE + # quality drops). Cached lazily after weights are loaded (CG-safe). + kw, _ = self.wk_weights_proj(hidden_states) + k = kw[:, : self.head_dim] + if getattr(self, "_wp_fp32", None) is None: + self._wp_fp32 = ( + self.wk_weights_proj.weight.data[self.head_dim :, :] + .t() + .contiguous() + .float() + ) + # deneb fork (glm53_indexer_gate_splitk): split-K Triton for the small-M + # decode shape (opt-in; cuBLAS picks a 2-block kernel here) + weights = _glm53_head_gate(hidden_states, self._wp_fp32) + + k = _fused_indexer_k_norm( + k, self.k_norm.weight, self.k_norm.bias, self.head_dim, self.k_norm.eps + ) + + if self.rope_dim > 0: + q_pe, q_nope = torch.split( + q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + k_pe, k_nope = torch.split( + k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1)) + # Note: RoPE (NeoX) can introduce extra leading dimensions during + # compilation so we need to reshape back to token-flattened shapes + q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim) + k_pe = k_pe.reshape(-1, 1, self.rope_dim) + + # `rotary_emb` is shape-preserving; `q_pe` is already + # [num_tokens, n_head, rope_dim]. + q = torch.cat([q_pe, q_nope], dim=-1) + # `k_pe` is [num_tokens, 1, rope_dim] (MQA). + k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) + # else: qk_rope_head_dim=0 — no rope component. q is already + # [num_tokens, n_head, head_dim] and k is [num_tokens, head_dim] (all + # nope), so skip the rope split / rotary / cat entirely; otherwise the + # split/reshape would build 0-element tensors (breaks dynamo tracing). + + # Match sglang rotate_activation(query): Hadamard-128 on head_dim so the + # fp8 MQA logits are (Hq).(Hk) == q.k (the K compress already stores Hk). + # Without this the head-gate is scored against q.(Hk) -- a rotated basis + # the gate was never trained against, destroying long-context pool + # discrimination (the needle's pool drops out of the top-k). + # Fused FWHT + fp8 quant (ops/kpool_compress.fwht128_quant_fp8): the + # butterflies run in fp32 with the exact 1/sqrt(128) constant, matching + # sglang's fast hadamard_transform (the previous ``q @ _hadamard`` bf16 + # GEMM stored H in bf16, a ~2^-9 systematic bias), and fusing the quant + # saves the rotated tensor's HBM round-trip vs the two-kernel chain. + assert self.head_dim == 128 and self.quant_block_size == 128 + assert self.scale_fmt == "ue8m0" + q = q.view(-1, self.head_dim) + q_fp8, q_scale = fwht128_quant_fp8(q) + q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim) + q_scale = q_scale.view(-1, self.n_head, 1) + + weights = _fused_indexer_weight_scale( + weights, q_scale, self.softmax_scale * self.n_head**-0.5 + ) + + # kpool: per-token gate score driving the softmax-weighted pool. Computed + # from the same hidden_states that produced `k`, so it stays token-aligned. + # F.linear(x, gate) = x @ gate.T with gate [head_dim, hidden_size]. + gate_score = F.linear(hidden_states, self.index_kpool_compress_gate) + + # DeepGEMM's MQA-logits kernels (fp8_mqa_logits / + # fp8_fp4_paged_mqa_logits) require num_heads in {32, 64}; this + # checkpoint uses index_n_heads=16. Zero-pad q and the per-head + # weights: logits are a weights-weighted sum over heads, so + # zero-weight padded heads contribute exactly nothing. + if self.n_head < 32: + pad = 32 - self.n_head + q_fp8 = _pad_indexer_heads(q_fp8, pad) + weights = _pad_indexer_heads(weights, pad) + + return self.indexer_op( + hidden_states, + q_fp8, + k, + weights, + gate_score=gate_score, + compress_ape=self.index_kpool_compress_ape, + index_kpool=self.index_kpool, + positions=positions, + ) + + +class Glm5NextMLAAttention(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: Glm5NextConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + max_position_embeddings: int = 8192, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + topk_indices_buffer: torch.Tensor | None = None, + input_size: int | None = None, + skip_rope: bool | None = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + + self.num_heads = num_heads + tp_size = get_tensor_model_parallel_world_size() + assert num_heads % tp_size == 0 + self.num_local_heads = num_heads // tp_size + + self.scaling = self.qk_head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + + # Use input_size for projection input dimensions if provided, + # otherwise default to hidden_size (used in Eagle3 Deepseek with MLA) + proj_input_size = input_size if input_size is not None else self.hidden_size + + if self.q_lora_rank is not None: + self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear( + proj_input_size, + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + ) + else: + self.kv_a_proj_with_mqa = ReplicatedLinear( + proj_input_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + + if self.q_lora_rank is not None: + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + self.q_proj = ColumnParallelLinear( + proj_input_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + if not skip_rope: + assert config.rope_parameters is not None + if config.rope_parameters["rope_type"] != "default": + config.rope_parameters["rope_type"] = ( + "deepseek_yarn" + if config.rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + + self.rotary_emb: RotaryEmbedding | None = get_rope( + qk_rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=False, + ) + + if ( + config.rope_parameters["rope_type"] != "default" + and config.rope_parameters["rope_type"] == "deepseek_yarn" + ): + mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False) + scaling_factor = config.rope_parameters["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale + else: + self.rotary_emb = None + + # `index_topk` is declared on Glm5NextTextConfig with a default of None, + # so hasattr() is True even for full-MLA configs (no kpool indexer). + self.is_v32 = getattr(config, "index_topk", None) is not None + # self.is_v32 = False + + _skip_topk = False + if self.is_v32: + self.indexer_rope_emb: RotaryEmbedding | None = get_rope( + qk_rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=not getattr(config, "indexer_rope_interleave", False), + ) + # The sparse indexer projects from the MLA q-lora rank, which is + # always set for v32 MLA configs; narrow away the `int | None`. + assert q_lora_rank is not None + self.indexer: Indexer | None = Indexer( + vllm_config, + config, + hidden_size, + q_lora_rank, + quant_config, + cache_config, + topk_indices_buffer, + f"{prefix}.indexer", + ) + + # Enable IndexCache for DeepSeek models to reduce redundant top-k + # token selection computations in sparse attention. + use_index_cache = getattr(config, "use_index_cache", False) + if use_index_cache: + # IndexCache config + # Refer: https://arxiv.org/abs/2603.12201 for more details. + _index_topk_freq = getattr(config, "index_topk_freq", 1) + _index_topk_pattern = getattr(config, "index_topk_pattern", None) + layer_id = extract_layer_index(prefix) + if _index_topk_pattern is None: + _skip_topk = max(layer_id - 1, 0) % _index_topk_freq != 0 + elif 0 <= layer_id < len(_index_topk_pattern): + _skip_topk = _index_topk_pattern[layer_id] == "S" + + else: + self.indexer_rope_emb = None + self.indexer = None + + mla_modules = MLAModules( + kv_a_layernorm=self.kv_a_layernorm, + kv_b_proj=self.kv_b_proj, + rotary_emb=self.rotary_emb, + o_proj=self.o_proj, + fused_qkv_a_proj=self.fused_qkv_a_proj + if self.q_lora_rank is not None + else None, + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa + if self.q_lora_rank is None + else None, + q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None, + q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None, + q_proj=self.q_proj if self.q_lora_rank is None else None, + indexer=self.indexer, + indexer_rotary_emb=self.indexer_rope_emb, + is_sparse=self.is_v32, + topk_indices_buffer=topk_indices_buffer, + ) + + self.mla_attn = MultiHeadLatentAttentionWrapper( + self.hidden_size, + self.num_local_heads, + self.scaling, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + self.q_lora_rank, + self.kv_lora_rank, + mla_modules, + cache_config, + quant_config, + prefix, + skip_topk=_skip_topk, + fuse_qkv_rmsnorm=True, + ) + + def forward( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + # Delegate to the MultiHeadLatentAttentionWrapper, which performs the + # q/kv projection, RoPE, the sparse-indexer top-k selection + # (``self.indexer``), the inner MLA attention, and the output + # projection. Re-implementing the projection here and calling the inner + # ``mla_attn`` directly would skip the indexer call, leaving the topk + # buffer empty and silently corrupting attention. Mirrors + # DeepseekV2MLAAttention.forward. + return self.mla_attn(positions, hidden_states) diff --git a/build/glm53/manifest.tsv b/build/glm53/manifest.tsv index b03d1428..7a767716 100644 --- a/build/glm53/manifest.tsv +++ b/build/glm53/manifest.tsv @@ -29,3 +29,5 @@ glm53_megakernel.cu /usr/local/lib/python3.12/dist-packages/vllm/model_executor/ glm5next_kda.py /usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia/kda.py ec090aabecc1a63dacc9694ea677b195e95ce0c63648c418a6daaf34b8196125 glm53_prep_fused.py /usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia/glm53_prep_fused.py absent glm53_config_vllm.py /usr/local/lib/python3.12/dist-packages/vllm/config/vllm.py 2469664631e33ba4b317ca085405aa1b751e405f9c52ff9d3dda61a0cfc6c5a9 +glm5next_attention.py /usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia/attention.py a0870c317287ab444a4639ad418964e981efafb5386592540f85eb38d534c369 +glm53_indexer_gate.py /usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia/glm53_indexer_gate.py absent diff --git a/overlay/modules/glm53_indexer_gate_splitk/README.md b/overlay/modules/glm53_indexer_gate_splitk/README.md new file mode 100644 index 00000000..2b5ebe5b --- /dev/null +++ b/overlay/modules/glm53_indexer_gate_splitk/README.md @@ -0,0 +1,68 @@ +# glm53_indexer_gate_splitk + +The sparse indexer's fp32 head-gate projection as a split-K Triton kernel, +behind `VLLM_GLM53_INDEXER_GATE_SPLITK` (default `0`, stock `torch.mm`). + +## What the stock code does + +`Indexer.forward` (`vllm/models/glm5next/nvidia/attention.py`) computes + +```python +weights = torch.mm(hidden_states.float(), self._wp_fp32) # [M, 4096] x [4096, 16] +``` + +once per full-attention layer, in fp32 on purpose: bf16 head-gates (~1e-2 +error) flip near-tie pool rankings, and the ranking is what the sparse +attention selects. cuBLAS answers this shape with a two-block `gemmSN` +kernel: 47 us on an idle GB10, 86 us under CUPTI in the 2026-09-01 serving +trace, eleven times per decode step (0.95 ms/step CUPTI) for 256 KB of +weights. Two blocks on 48 SMs is the whole problem. + +## What this module does + +A split-K kernel (`glm53_indexer_gate.py`): one program per (row, K-slice of +512), each accumulating 16 outputs in fp32 and reducing with fp32 atomics +into a zeroed output. Used only when the knob is `1` **and** `M <= 16` +(decode: C=1 verify batches are M=8, C=2 M=16); every other shape, prefill +included, keeps `torch.mm`. The fused-indexer forward in +`glm53_prefill_fastpath.py` (`VLLM_GLM53_FUSED_K_GATE=1`) routes through the +same helper, so the two arms cannot disagree. + +Both paths accumulate in fp32; only the summation order differs, so this is +**not bit-exact** with stock -- it is a numerics change and needs the +quality bracket, not just a timing one. + +## Offline numbers (srv4 GB10, `probes/indexer_gate_check.py`, CUDA-graph replay) + +| M | stock `torch.mm` | split-K | route | +|---|---|---|---| +| 1 | 15.5 us | 8.5 us | split-K | +| 8 | 50.0 us | 9.9 us | split-K | +| 16 | 50.1 us | 11.7 us | split-K | +| 32 | 15.9 us | 17.4 us | `torch.mm` kept | + +Numerics over 300 trials / 2,480 rows (bf16 activations, fp32 weights of the +indexer's scale): max |diff| 2.4e-6 absolute, 6.7e-7 of the row's max gate, +0 top-1 flips, 0 top-4 set changes. + +Ceiling: 11 layers x ~40 us = ~0.44 ms/step at C=1, about 0.65% of a 66 ms +step -- below the ledger's 1% boot threshold on its own. Ride it on another +bracket's boot (it is independent of `glm53_prep_fused` and +`glm53_async_dflash`), never on a boot of its own. + +## Arming + +`VLLM_GLM53_INDEXER_GATE_SPLITK` is a profile-declared key: pass it as caller +env, never through `EXTRA_ENV`. + +```bash +VLLM_GLM53_INDEXER_GATE_SPLITK=1 bash launchers/start-glm53-nvfp4-tp4.sh +``` + +Verdict in the trace: the eleven `gemmSN` launches are replaced by eleven +`_gate_splitk_kernel` launches (no boot-log line; the helper is called per +layer inside the captured graph). Gate: quality 9/9, Korean 0/16, C=1 step/s +bracket base -> cand -> base. + +Preimage: `attention.py` of `glm53:v13-b12x` +(`a0870c31...`), identical in `glm53:sm121-fi618`. diff --git a/overlay/modules/glm53_indexer_gate_splitk/glm53_indexer_gate.py b/overlay/modules/glm53_indexer_gate_splitk/glm53_indexer_gate.py new file mode 100644 index 00000000..cd73f8e9 --- /dev/null +++ b/overlay/modules/glm53_indexer_gate_splitk/glm53_indexer_gate.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +"""deneb fork: the sparse indexer's fp32 head-gate projection as a split-K +Triton kernel (VLLM_GLM53_INDEXER_GATE_SPLITK). + +`Indexer.forward` computes `weights = torch.mm(hidden_states.float(), +self._wp_fp32)` -- an [M, 4096] x [4096, 16] fp32 product per full-attention +layer, kept in fp32 on purpose (bf16 head-gates flip near-tie pool rankings). +cuBLAS answers that shape with a two-block `gemmSN` kernel: 47 us on an idle +GB10 (86 us under CUPTI in the 2026-09-01 serving trace), eleven times per +decode step, for 256 KB of weights. A split-K kernel that hands each +(row, K-slice) to one program and reduces with fp32 atomics runs the same +product in 7 us. Both accumulate in fp32; only the summation order differs +(measured max |diff| 3e-5 on values of magnitude ~64, i.e. ~5e-7 relative, +0 top-1 rank flips over the offline trials) -- not bit-exact, so this stays +an opt-in behind a numerics bracket. + +At M > 16 (C >= 3 verify batches, prefill) cuBLAS is already fast, so the +kernel is used only for M <= 16; larger M keeps torch.mm. +""" +from __future__ import annotations + +import os + +import torch + +from vllm.triton_utils import tl, triton + +ENV = "VLLM_GLM53_INDEXER_GATE_SPLITK" +MAX_M = 16 +_SPLIT = 8 +_BLOCK_K = 128 + + +def gate_splitk_enabled() -> bool: + """Exact opt-in: only the string "1" arms; anything else is stock.""" + return os.environ.get(ENV, "").strip() == "1" + + +@triton.jit +def _gate_splitk_kernel(x_ptr, w_ptr, out_ptr, K, N, sxm, swk, som, + SPLIT: tl.constexpr, BLOCK_K: tl.constexpr, BN: tl.constexpr): + m = tl.program_id(0) + s = tl.program_id(1) + kper = K // SPLIT + k0 = s * kper + offs_n = tl.arange(0, BN) + nmask = offs_n < N + acc = tl.zeros([BN], dtype=tl.float32) + for k in range(0, kper, BLOCK_K): + offs_k = k0 + k + tl.arange(0, BLOCK_K) + xv = tl.load(x_ptr + m * sxm + offs_k) + wv = tl.load(w_ptr + offs_k[:, None] * swk + offs_n[None, :], mask=nmask[None, :], other=0.0) + acc += tl.sum(xv[:, None] * wv, axis=0) + tl.atomic_add(out_ptr + m * som + offs_n, acc, mask=nmask) + + +def head_gate_splitk(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """fp32 [M, K] @ [K, N] for N <= 16, K a multiple of SPLIT*BLOCK_K. + + x may be bf16 (cast to fp32 here, as the stock `.float()` does); w is the + stock `_wp_fp32` ([K, N], fp32, contiguous).""" + xf = x.float() + M, K = xf.shape + N = w.shape[1] + out = torch.zeros(M, N, device=x.device, dtype=torch.float32) + _gate_splitk_kernel[(M, _SPLIT)]( + xf, w, out, K, N, xf.stride(0), w.stride(0), out.stride(0), + SPLIT=_SPLIT, BLOCK_K=_BLOCK_K, BN=16, num_warps=4) + return out + + +def splitk_applicable(x: torch.Tensor, w: torch.Tensor) -> bool: + return (x.shape[0] <= MAX_M and w.dtype == torch.float32 and w.is_contiguous() + and w.shape[1] <= 16 and w.shape[0] % (_SPLIT * _BLOCK_K) == 0) + + +def head_gate(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """The stock `torch.mm(x.float(), w)` unless the knob is on and the shape + is the small-M decode one.""" + if gate_splitk_enabled() and splitk_applicable(x, w): + return head_gate_splitk(x, w) + return torch.mm(x.float(), w) diff --git a/overlay/modules/glm53_indexer_gate_splitk/glm5next_attention.py b/overlay/modules/glm53_indexer_gate_splitk/glm5next_attention.py new file mode 100644 index 00000000..ee1dcf8a --- /dev/null +++ b/overlay/modules/glm53_indexer_gate_splitk/glm5next_attention.py @@ -0,0 +1,627 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.config import ( + CacheConfig, + VllmConfig, +) +from vllm.distributed import ( + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding, get_rope +from vllm.model_executor.layers.sparse_attn_indexer_kpool import SparseAttnIndexerKpool +from vllm.model_executor.models.deepseek_v2 import ( + DeepSeekV2FusedQkvAProjLinear, + DeepseekV32IndexerCache, + yarn_get_mscale, +) +from vllm.model_executor.models.utils import extract_layer_index +from vllm.model_executor.utils import maybe_disable_graph_partition +from vllm.models.glm5next.nvidia.ops.kpool_compress import fwht128_quant_fp8 +from vllm.models.glm5next.nvidia.glm53_indexer_gate import head_gate as _glm53_head_gate +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig +from vllm.v1.kv_cache_interface import KpoolTailSpec, MLAAttentionSpec + +logger = init_logger(__name__) + +# Shared torch.compile config for the indexer's small-kernel leaves. The MLA +# indexer runs under breakable-CG (CompilationMode.NONE), which blocks FX-graph +# fusion of the surrounding eager ops; carving each cluster into its own +# @torch.compile leaf (backend==inductor) still fuses them. Matches the +# grouped_topk / _cast_sigmoid leaf pattern. +_INDEXER_COMPILE = dict( + dynamic=True, + backend=current_platform.simple_compile_backend, + options=maybe_disable_graph_partition(current_platform.simple_compile_backend), +) + + +@torch.compile(**_INDEXER_COMPILE) +def _fused_indexer_k_norm( + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, dim: int, eps: float +) -> torch.Tensor: + # Fuse fp32 cast + layer_norm + cast-back (was 3 kernels) into one. + return F.layer_norm(x.float(), (dim,), weight, bias, eps).type_as(x) + + +@torch.compile(**_INDEXER_COMPILE) +def _fused_indexer_weight_scale( + weights: torch.Tensor, q_scale: torch.Tensor, scale: float +) -> torch.Tensor: + # Fuse the weight-scaling muls (was 2 kernels) into one. `scale` folds + # softmax_scale (head_dim**-0.5) and n_head**-0.5 into a single constant. + return (weights.unsqueeze(-1) * q_scale * scale).squeeze(-1) + + +@torch.compile(**_INDEXER_COMPILE) +def _pad_indexer_heads(x: torch.Tensor, pad: int) -> torch.Tensor: + # DeepGEMM MQA-logits needs num_heads in {32,64}; zero-pad the head dim. + # Fuse new_zeros + cat (was 2 kernels) into one. Pad values are zero (exact + # in fp8 e4m3 and zero-weight in the logits sum), so numerically a no-op. + return torch.cat([x, x.new_zeros(x.shape[0], pad, *x.shape[2:])], dim=1) + + +class Glm5NextIndexerCache(DeepseekV32IndexerCache): + """Indexer K cache that stores kpool-compressed entries. + + Setting ``compress_ratio = index_kpool`` on the kv_cache_spec makes vLLM's + indexer metadata builder emit pool-granular ``slot_mapping`` / + ``seq_lens`` / ``cu_seq_lens`` / ``page_table`` for free, and shrinks the + cache allocation to ``storage_block_size = block_size // kpool``. The pool + *content* (softmax-weighted sum vs keep-every-Nth) is computed by the + kpool compress kernel inside the indexer op — the cache only provides the + addressing, which is identical for both schemes. + + The indexer shares one block with the co-located MLA (a single + ``MLAAttentionSpec`` / block_table), so ``block_size`` is the model-wide + ``cache_config.block_size``. DeepGEMM's paged-MQA kernel + (``csrc/apis/attention.hpp``) requires ``block_kv`` to be exactly 32 or + 64, so the storage block is virtually split into pool pages of the + largest such size that tiles it (``storage_kernel_block_size``); this + needs ``block_size`` to be a multiple of ``index_kpool * 32`` (512 for + ``index_kpool = 16``). A smaller block (e.g. the default 64) silently + collapses ``storage_block_size`` (64 // 16 = 4) and only fails later at + the opaque C++ assert; ``get_kv_cache_spec`` guards this up front + instead. + """ + + def __init__( + self, + *, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config, + index_kpool: int, + ): + super().__init__( + head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config + ) + assert index_kpool > 1, "Glm5NextIndexerCache expects index_kpool > 1" + # Chunked prefill aligns chunk ends to ``cache_config.block_size``; + # the prefill compress kernel (``_kpool_compress_insert``) assumes + # pool-aligned chunk starts, which holds only if that block_size is a + # multiple of kpool -- otherwise every pool straddling a chunk + # boundary is silently dropped (masked off as a leading-partial pool). + # Distinct from the get_kv_cache_spec assert, which guards the spec's + # DeepGEMM tiling on its own block_size. + assert cache_config.block_size % index_kpool == 0, ( + "Glm5NextIndexerCache: cache_config.block_size " + f"({cache_config.block_size}) must be a multiple of index_kpool " + f"({index_kpool}) so chunked-prefill boundaries stay pool-aligned." + ) + self._index_kpool = index_kpool + + def get_kv_cache_spec(self, vllm_config: VllmConfig): + from dataclasses import replace + + spec = super().get_kv_cache_spec(vllm_config) + # compress_ratio lives on MLAAttentionSpec, but the base + # DeepseekV32IndexerCache.get_kv_cache_spec is typed to return the + # KVCacheSpec base; narrow so dataclass.replace sees the field. + assert isinstance(spec, MLAAttentionSpec) + spec = replace(spec, compress_ratio=self._index_kpool) + + # DeepGEMM paged-MQA takes block_kv in {32, 64}; the storage block + # (= block_size // index_kpool) is virtually split into pool pages of + # the largest such size that tiles it, so it must be a multiple of 32. + storage_block_size = spec.block_size // self._index_kpool + assert ( + spec.block_size % self._index_kpool == 0 and storage_block_size % 32 == 0 + ), ( + "Glm5NextIndexerCache: kpool indexer requires cache block_size to " + f"be a multiple of index_kpool * 32 ({self._index_kpool * 32}) so " + "that DeepGEMM paged-MQA pool pages (32 or 64 entries) tile the " + f"storage block, got block_size={spec.block_size} -> " + f"storage_block_size={storage_block_size}." + ) + return spec + + +class Glm5NextTailCache(DeepseekV32IndexerCache): + """Paged circular buffer for the kpool indexer's in-progress (tail) pool. + + Holds the trailing incomplete pool's raw K + gate score: one block of + ``index_kpool`` slots per request, overwritten in place by ``pos % kpool`` + as decode/spec-decode advances. Prefill seeds it (instead of discarding the + tail raw K+gate); the connector transfers it across PD; decode reads it to + compress the boundary pool correctly. ``KpoolTailSpec`` / + ``KpoolTailManager`` provide the no-prune, 1-block/req allocation that lets + the in-progress pool survive across steps and across transfer. + + Stores raw bf16 K (``head_dim``) as the "K" half of each block and the + bf16 gate score (``head_dim``) as the "V" half -- not the fp8-compressed + entry, which lives in ``Glm5NextIndexerCache``. + """ + + def __init__( + self, + *, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config, + index_kpool: int, + ): + super().__init__( + head_dim=head_dim, dtype=dtype, prefix=prefix, cache_config=cache_config + ) + assert index_kpool > 1, "Glm5NextTailCache expects index_kpool > 1" + self._index_kpool = index_kpool + + def get_kv_cache_spec(self, vllm_config: VllmConfig): + # K + gate score packed into head_size (== 2*head_dim), head_size_v=0: + # KpoolTailBackend.get_kv_cache_shape only consumes head_size and splits + # it into [2, kpool, head_dim] (K | score halves), so the connectors' + # non-MLA K/V half-split transfers K and score as separate halves. + return KpoolTailSpec( + block_size=self._index_kpool, + num_kv_heads=1, + head_size=2 * self.head_dim, + head_size_v=0, + dtype=torch.bfloat16, + sliding_window=self._index_kpool, + ) + + def get_attn_backend(self): + from vllm.v1.attention.backends.mla.indexer import KpoolTailBackend + + return KpoolTailBackend + + +class Indexer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: Glm5NextConfig, + hidden_size: int, + q_lora_rank: int, + quant_config: QuantizationConfig | None, + cache_config: CacheConfig | None, + topk_indices_buffer: torch.Tensor | None, + prefix: str = "", + ): + super().__init__() + self.vllm_config = vllm_config + self.config = config + self.quant_config = quant_config + # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"] + # Indexer is only constructed for v32 configs, where these sparse-indexer + # fields are guaranteed populated; narrow away the `int | None` declared + # on Glm5NextConfig for the optional-indexer case. + assert config.index_topk is not None + assert config.index_n_heads is not None + assert config.index_head_dim is not None + assert config.index_kpool is not None + self.topk_tokens = config.index_topk + self.n_head = config.index_n_heads # 64 + self.head_dim = config.index_head_dim # 128 + self.rope_dim = config.qk_rope_head_dim # 64 + self.index_kpool = config.index_kpool + self.q_lora_rank = q_lora_rank # 1536 + + # kpool + self.index_kpool_compress_ape = nn.Parameter( + torch.zeros(self.index_kpool, self.head_dim, dtype=torch.float32) + ) + # NOTE: kept as a bare nn.Parameter (not ReplicatedLinear) so the weight + # name matches the checkpoint verbatim ("index_kpool_compress_gate", + # no ".weight" suffix) — the trained checkpoint stores it the sglang way. + # Shape [head_dim, hidden_size]; consumed via F.linear(x, gate) = x @ gate.T. + self.index_kpool_compress_gate = nn.Parameter( + torch.empty(self.head_dim, hidden_size, dtype=torch.bfloat16) + ) + + # no tensor parallel, just replicated + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.head_dim * self.n_head, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wq_b", + ) + # Fused wk + weights_proj: single GEMM producing [head_dim + n_head]. + # FP8 wk weights are upcasted to BF16 during loading to maintain fusion. + self.wk_weights_proj = MergedColumnParallelLinear( + hidden_size, + [self.head_dim, self.n_head], + bias=False, + quant_config=None, + disable_tp=True, + prefix=f"{prefix}.wk_weights_proj", + ) + self.k_norm = LayerNorm(self.head_dim, eps=1e-6) + self.softmax_scale = self.head_dim**-0.5 + + # Hadamard-128 rotation of the indexer query is fused with the FP8 + # quant (see forward: fwht128_quant_fp8) -- no precomputed matrix. + + self.scale_fmt = "ue8m0" + self.quant_block_size = 128 # TODO: get from config + self.topk_indices_buffer = topk_indices_buffer + + # NOTE: (zyongye) we use fp8 naive cache, + # where we store value in fp8 and scale in fp32 + # per self.quant_block_size element + self.k_cache = Glm5NextIndexerCache( + head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4, + dtype=torch.uint8, + prefix=f"{prefix}.k_cache", + cache_config=cache_config, + index_kpool=self.index_kpool, + ) + # Paged tail cache (in-progress pool's raw K + gate score). Written by + # prefill (seeds the boundary pool) and decode (per-step stash); read by + # the decode kernel to compress the boundary pool. Transferred across PD + # so the decode side sees the prefill tail. See KpoolTailSpec/Manager. + self.tail_cache = Glm5NextTailCache( + head_dim=self.head_dim, + dtype=torch.bfloat16, + prefix=f"{prefix}.tail_cache", + cache_config=cache_config, + index_kpool=self.index_kpool, + ) + self.max_model_len = vllm_config.model_config.max_model_len + self.prefix = prefix + from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size + + self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config) + self.indexer_op = SparseAttnIndexerKpool( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + tail_cache=self.tail_cache, + ) + + def forward( + self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb + ) -> torch.Tensor: + q, _ = self.wq_b(qr) + q = q.view(-1, self.n_head, self.head_dim) + + # Fused wk + weights_proj: one GEMM for wk (bf16). The head-gate + # (weights_proj) is computed separately in fp32 to match sglang + # (ReplicatedLinear params_dtype=fp32, x.float()) -- bf16 head-gates + # introduce ~1e-2 logit error that flips near-tie pool rankings on hard + # long-context tasks (single-needle retrieval is unaffected, but HLE + # quality drops). Cached lazily after weights are loaded (CG-safe). + kw, _ = self.wk_weights_proj(hidden_states) + k = kw[:, : self.head_dim] + if getattr(self, "_wp_fp32", None) is None: + self._wp_fp32 = ( + self.wk_weights_proj.weight.data[self.head_dim :, :] + .t() + .contiguous() + .float() + ) + # deneb fork (glm53_indexer_gate_splitk): split-K Triton for the small-M + # decode shape (opt-in; cuBLAS picks a 2-block kernel here) + weights = _glm53_head_gate(hidden_states, self._wp_fp32) + + k = _fused_indexer_k_norm( + k, self.k_norm.weight, self.k_norm.bias, self.head_dim, self.k_norm.eps + ) + + if self.rope_dim > 0: + q_pe, q_nope = torch.split( + q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + k_pe, k_nope = torch.split( + k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 + ) + + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1)) + # Note: RoPE (NeoX) can introduce extra leading dimensions during + # compilation so we need to reshape back to token-flattened shapes + q_pe = q_pe.reshape(-1, self.n_head, self.rope_dim) + k_pe = k_pe.reshape(-1, 1, self.rope_dim) + + # `rotary_emb` is shape-preserving; `q_pe` is already + # [num_tokens, n_head, rope_dim]. + q = torch.cat([q_pe, q_nope], dim=-1) + # `k_pe` is [num_tokens, 1, rope_dim] (MQA). + k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1) + # else: qk_rope_head_dim=0 — no rope component. q is already + # [num_tokens, n_head, head_dim] and k is [num_tokens, head_dim] (all + # nope), so skip the rope split / rotary / cat entirely; otherwise the + # split/reshape would build 0-element tensors (breaks dynamo tracing). + + # Match sglang rotate_activation(query): Hadamard-128 on head_dim so the + # fp8 MQA logits are (Hq).(Hk) == q.k (the K compress already stores Hk). + # Without this the head-gate is scored against q.(Hk) -- a rotated basis + # the gate was never trained against, destroying long-context pool + # discrimination (the needle's pool drops out of the top-k). + # Fused FWHT + fp8 quant (ops/kpool_compress.fwht128_quant_fp8): the + # butterflies run in fp32 with the exact 1/sqrt(128) constant, matching + # sglang's fast hadamard_transform (the previous ``q @ _hadamard`` bf16 + # GEMM stored H in bf16, a ~2^-9 systematic bias), and fusing the quant + # saves the rotated tensor's HBM round-trip vs the two-kernel chain. + assert self.head_dim == 128 and self.quant_block_size == 128 + assert self.scale_fmt == "ue8m0" + q = q.view(-1, self.head_dim) + q_fp8, q_scale = fwht128_quant_fp8(q) + q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim) + q_scale = q_scale.view(-1, self.n_head, 1) + + weights = _fused_indexer_weight_scale( + weights, q_scale, self.softmax_scale * self.n_head**-0.5 + ) + + # kpool: per-token gate score driving the softmax-weighted pool. Computed + # from the same hidden_states that produced `k`, so it stays token-aligned. + # F.linear(x, gate) = x @ gate.T with gate [head_dim, hidden_size]. + gate_score = F.linear(hidden_states, self.index_kpool_compress_gate) + + # DeepGEMM's MQA-logits kernels (fp8_mqa_logits / + # fp8_fp4_paged_mqa_logits) require num_heads in {32, 64}; this + # checkpoint uses index_n_heads=16. Zero-pad q and the per-head + # weights: logits are a weights-weighted sum over heads, so + # zero-weight padded heads contribute exactly nothing. + if self.n_head < 32: + pad = 32 - self.n_head + q_fp8 = _pad_indexer_heads(q_fp8, pad) + weights = _pad_indexer_heads(weights, pad) + + return self.indexer_op( + hidden_states, + q_fp8, + k, + weights, + gate_score=gate_score, + compress_ape=self.index_kpool_compress_ape, + index_kpool=self.index_kpool, + positions=positions, + ) + + +class Glm5NextMLAAttention(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: Glm5NextConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + max_position_embeddings: int = 8192, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + topk_indices_buffer: torch.Tensor | None = None, + input_size: int | None = None, + skip_rope: bool | None = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + + self.num_heads = num_heads + tp_size = get_tensor_model_parallel_world_size() + assert num_heads % tp_size == 0 + self.num_local_heads = num_heads // tp_size + + self.scaling = self.qk_head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + + # Use input_size for projection input dimensions if provided, + # otherwise default to hidden_size (used in Eagle3 Deepseek with MLA) + proj_input_size = input_size if input_size is not None else self.hidden_size + + if self.q_lora_rank is not None: + self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear( + proj_input_size, + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + quant_config=quant_config, + prefix=f"{prefix}.fused_qkv_a_proj", + ) + else: + self.kv_a_proj_with_mqa = ReplicatedLinear( + proj_input_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_a_proj_with_mqa", + ) + + if self.q_lora_rank is not None: + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_b_proj", + ) + else: + self.q_proj = ColumnParallelLinear( + proj_input_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.kv_b_proj", + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + if not skip_rope: + assert config.rope_parameters is not None + if config.rope_parameters["rope_type"] != "default": + config.rope_parameters["rope_type"] = ( + "deepseek_yarn" + if config.rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + + self.rotary_emb: RotaryEmbedding | None = get_rope( + qk_rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=False, + ) + + if ( + config.rope_parameters["rope_type"] != "default" + and config.rope_parameters["rope_type"] == "deepseek_yarn" + ): + mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False) + scaling_factor = config.rope_parameters["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale + else: + self.rotary_emb = None + + # `index_topk` is declared on Glm5NextTextConfig with a default of None, + # so hasattr() is True even for full-MLA configs (no kpool indexer). + self.is_v32 = getattr(config, "index_topk", None) is not None + # self.is_v32 = False + + _skip_topk = False + if self.is_v32: + self.indexer_rope_emb: RotaryEmbedding | None = get_rope( + qk_rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=not getattr(config, "indexer_rope_interleave", False), + ) + # The sparse indexer projects from the MLA q-lora rank, which is + # always set for v32 MLA configs; narrow away the `int | None`. + assert q_lora_rank is not None + self.indexer: Indexer | None = Indexer( + vllm_config, + config, + hidden_size, + q_lora_rank, + quant_config, + cache_config, + topk_indices_buffer, + f"{prefix}.indexer", + ) + + # Enable IndexCache for DeepSeek models to reduce redundant top-k + # token selection computations in sparse attention. + use_index_cache = getattr(config, "use_index_cache", False) + if use_index_cache: + # IndexCache config + # Refer: https://arxiv.org/abs/2603.12201 for more details. + _index_topk_freq = getattr(config, "index_topk_freq", 1) + _index_topk_pattern = getattr(config, "index_topk_pattern", None) + layer_id = extract_layer_index(prefix) + if _index_topk_pattern is None: + _skip_topk = max(layer_id - 1, 0) % _index_topk_freq != 0 + elif 0 <= layer_id < len(_index_topk_pattern): + _skip_topk = _index_topk_pattern[layer_id] == "S" + + else: + self.indexer_rope_emb = None + self.indexer = None + + mla_modules = MLAModules( + kv_a_layernorm=self.kv_a_layernorm, + kv_b_proj=self.kv_b_proj, + rotary_emb=self.rotary_emb, + o_proj=self.o_proj, + fused_qkv_a_proj=self.fused_qkv_a_proj + if self.q_lora_rank is not None + else None, + kv_a_proj_with_mqa=self.kv_a_proj_with_mqa + if self.q_lora_rank is None + else None, + q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None, + q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None, + q_proj=self.q_proj if self.q_lora_rank is None else None, + indexer=self.indexer, + indexer_rotary_emb=self.indexer_rope_emb, + is_sparse=self.is_v32, + topk_indices_buffer=topk_indices_buffer, + ) + + self.mla_attn = MultiHeadLatentAttentionWrapper( + self.hidden_size, + self.num_local_heads, + self.scaling, + self.qk_nope_head_dim, + self.qk_rope_head_dim, + self.v_head_dim, + self.q_lora_rank, + self.kv_lora_rank, + mla_modules, + cache_config, + quant_config, + prefix, + skip_topk=_skip_topk, + fuse_qkv_rmsnorm=True, + ) + + def forward( + self, hidden_states: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + # Delegate to the MultiHeadLatentAttentionWrapper, which performs the + # q/kv projection, RoPE, the sparse-indexer top-k selection + # (``self.indexer``), the inner MLA attention, and the output + # projection. Re-implementing the projection here and calling the inner + # ``mla_attn`` directly would skip the indexer call, leaving the topk + # buffer empty and silently corrupting attention. Mirrors + # DeepseekV2MLAAttention.forward. + return self.mla_attn(positions, hidden_states) diff --git a/overlay/modules/glm53_indexer_gate_splitk/manifest.tsv b/overlay/modules/glm53_indexer_gate_splitk/manifest.tsv new file mode 100644 index 00000000..867ef7e1 --- /dev/null +++ b/overlay/modules/glm53_indexer_gate_splitk/manifest.tsv @@ -0,0 +1,3 @@ +# source container_target(relative to TARGET_PREFIX) base_preimage_sha256 +glm5next_attention.py vllm/models/glm5next/nvidia/attention.py a0870c317287ab444a4639ad418964e981efafb5386592540f85eb38d534c369 +glm53_indexer_gate.py vllm/models/glm5next/nvidia/glm53_indexer_gate.py absent diff --git a/overlay/modules/glm53_indexer_gate_splitk/requires b/overlay/modules/glm53_indexer_gate_splitk/requires new file mode 100644 index 00000000..1f2288fb --- /dev/null +++ b/overlay/modules/glm53_indexer_gate_splitk/requires @@ -0,0 +1 @@ +glm53_model_wiring diff --git a/overlay/modules/glm53_model_wiring/glm53_prefill_fastpath.py b/overlay/modules/glm53_model_wiring/glm53_prefill_fastpath.py index 8b2eaadd..e540490c 100644 --- a/overlay/modules/glm53_model_wiring/glm53_prefill_fastpath.py +++ b/overlay/modules/glm53_model_wiring/glm53_prefill_fastpath.py @@ -34,6 +34,23 @@ ) from .ops.kpool_compress import fwht128_quant_fp8 +# deneb fork (glm53_indexer_gate_splitk): the split-K helper when that module +# is mounted, else the stock fp32 torch.mm. Resolved once, on first call, so +# the fused indexer forward stays loadable without the sibling file. +_HEAD_GATE = None + + +def _glm53_head_gate(x, w): + global _HEAD_GATE + if _HEAD_GATE is None: + try: + from vllm.models.glm5next.nvidia.glm53_indexer_gate import head_gate as fn + except ImportError: + def fn(x, w): + return torch.mm(x.float(), w) + _HEAD_GATE = fn + return _HEAD_GATE(x, w) + logger = init_logger(__name__) _GLM53_SM121_MLA_PREFILL_ENV = "VLLM_GLM53_SM121_MLA_PREFILL" @@ -399,7 +416,7 @@ def _glm53_fused_indexer_forward( self._wp_fp32 = ( k_weight.data[self.head_dim :, :].t().contiguous().float() ) - weights = torch.mm(hidden_states.float(), self._wp_fp32) + weights = _glm53_head_gate(hidden_states, self._wp_fp32) # deneb fork (glm53_indexer_gate_splitk) k = _fused_indexer_k_norm( k, diff --git a/probes/indexer_gate_check.py b/probes/indexer_gate_check.py new file mode 100755 index 00000000..03a5b5f8 --- /dev/null +++ b/probes/indexer_gate_check.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""glm53_indexer_gate_splitk: numerics + timing of the split-K head gate. + +Compares head_gate_splitk against the stock torch.mm(x.float(), w) on random +and on checkpoint-like inputs (bf16 hidden states, fp32 weights of the +indexer's scale), reports max |diff| (absolute and relative to the row's +max), top-1 / top-4 pool-ranking flips across the 16 heads, and GPU time by +CUDA-graph replay. Run inside the image (any glm53 image; the kernel is +standalone): + + docker run --rm --gpus all --entrypoint python3 \ + --mount type=bind,src=$REPO,dst=/repo,readonly glm53:v13-b12x \ + /repo/probes/indexer_gate_check.py [--trials 200] +""" +from __future__ import annotations + +import argparse +import importlib.util +import sys +import time + +import torch + +sys.path.insert(0, "/usr/local/lib/python3.12/dist-packages") + + +def _load_kernel(): + spec = importlib.util.spec_from_file_location( + "glm53_indexer_gate", "/repo/overlay/modules/glm53_indexer_gate_splitk/glm53_indexer_gate.py") + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +def _gpu_us(fn, n_cap=20, reps=10): + for _ in range(5): + fn() + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + st = torch.cuda.Stream() + with torch.cuda.stream(st): + for _ in range(3): + fn() + st.synchronize() + with torch.cuda.graph(g, stream=st): + for _ in range(n_cap): + fn() + torch.cuda.synchronize() + a = torch.cuda.Event(enable_timing=True) + b = torch.cuda.Event(enable_timing=True) + g.replay() + torch.cuda.synchronize() + a.record() + for _ in range(reps): + g.replay() + b.record() + torch.cuda.synchronize() + return a.elapsed_time(b) / (reps * n_cap) * 1e3 + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--trials", type=int, default=200) + args = ap.parse_args() + torch.manual_seed(0) + dev = "cuda" + mod = _load_kernel() + K, N = 4096, 16 + worst_abs = worst_rel = 0.0 + flips1 = flips4 = rows = 0 + for t in range(args.trials): + M = int(torch.randint(1, mod.MAX_M + 1, (1,)).item()) + # hidden states ~ bf16 activations, weights ~ small fp32 projections + x = (torch.randn(M, K, device=dev) * 1.5).to(torch.bfloat16) + w = (torch.randn(N, K, device=dev) * 0.02).t().contiguous() + ref = torch.mm(x.float(), w) + out = mod.head_gate_splitk(x, w) + d = (out - ref).abs() + worst_abs = max(worst_abs, d.max().item()) + worst_rel = max(worst_rel, (d / ref.abs().amax(1, keepdim=True).clamp_min(1e-6)).max().item()) + flips1 += int((out.argmax(1) != ref.argmax(1)).sum().item()) + top4 = lambda a: a.topk(4, dim=1).indices.sort(1).values + flips4 += int((top4(out) != top4(ref)).any(1).sum().item()) + rows += M + print(f"numerics over {args.trials} trials / {rows} rows: max|diff| {worst_abs:.3e} abs, " + f"{worst_rel:.3e} of row max; top-1 flips {flips1}, top-4 set changes {flips4}") + for M in (1, 8, 16, 32): + x = torch.randn(M, K, device=dev, dtype=torch.bfloat16) + w = (torch.randn(N, K, device=dev) * 0.02).t().contiguous() + g_mm = _gpu_us(lambda: torch.mm(x.float(), w)) + g_sk = _gpu_us(lambda: mod.head_gate_splitk(x, w)) + t0 = time.perf_counter() + for _ in range(200): + mod.head_gate_splitk(x, w) + torch.cuda.synchronize() + h_sk = (time.perf_counter() - t0) / 200 * 1e6 + print(f"M={M:2d}: torch.mm {g_mm:6.1f} us GPU | split-K {g_sk:6.1f} us GPU ({h_sk:.0f} us host) | " + f"{'split-K used' if M <= mod.MAX_M else 'torch.mm kept'}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/profiles/glm53.env b/profiles/glm53.env index 325aa800..9853507a 100644 --- a/profiles/glm53.env +++ b/profiles/glm53.env @@ -47,7 +47,7 @@ TARGET_PREFIX="/usr/local/lib/python3.12/dist-packages/" # The image runs V2 Model Runner. glm53_drop_audit and glm53_sparse_q replace # V1-only paths and are deliberately absent: mounting them advertised guards # and sparse-q behavior that the live DFlash2 path never executed. -MODULES="b12x_shared_workspace b12x_zero_weight_micro glm53_tail_slot_persistent glm53_kpool_tail_select glm53_v2_sampler_guards moe_gate_sm121 glm53_model_wiring tp_oneshot_ar glm53_oneshot_wiring fp8_lm_head glm53_dflash2_fp8_head glm53_dflash_loader_fp8 glm53_fp8_dense glm53_b12x_out glm53_mhc_tilelang glm53_sm121_mla_prefill glm53_kda_prefill_regime glm53_dflash_warmup glm53_megakernel glm53_prep_fused glm53_async_dflash" +MODULES="b12x_shared_workspace b12x_zero_weight_micro glm53_tail_slot_persistent glm53_kpool_tail_select glm53_v2_sampler_guards moe_gate_sm121 glm53_model_wiring tp_oneshot_ar glm53_oneshot_wiring fp8_lm_head glm53_dflash2_fp8_head glm53_dflash_loader_fp8 glm53_fp8_dense glm53_b12x_out glm53_mhc_tilelang glm53_sm121_mla_prefill glm53_kda_prefill_regime glm53_dflash_warmup glm53_megakernel glm53_prep_fused glm53_async_dflash glm53_indexer_gate_splitk" # --- serving knobs ----------------------------------------------------------- # Dense-MHA for GLM's short/full-attention sparse-MLA prefill region on SM121. @@ -169,6 +169,12 @@ VLLM_GLM53_PREP_FUSED=0 # forces synchronous. Watch the KV-cache line: in-flight reserve doubles. VLLM_GLM53_ASYNC_DFLASH=0 +# 인덱서 fp32 head-gate GEMM ([M,4096]x[4096,16]) 을 split-K Triton 으로 — M<=16 +# 디코드 형상만, 나머지는 stock torch.mm. cuBLAS 2블록 47 us -> 10 us, 층당 11회. +# fp32 누적 순서만 다름(|diff|<=3e-6, 순위 뒤집힘 0/2480 오프라인) — 수치 브래킷 +# 대상. 정확히 "1" 만 켬. 프로필 선언 키: caller env 로 넘길 것 (EXTRA_ENV 거부). +VLLM_GLM53_INDEXER_GATE_SPLITK=0 + # Exact-shape GLM TP MoE tuning controls. Empty means the shipped automatic # dispatcher and MAC ladders; the launcher does not pass empty values into the # container. These are experiment inputs only -- no winner has been measured. diff --git a/tests/test_logic.py b/tests/test_logic.py index df768961..1dea9cbb 100644 --- a/tests/test_logic.py +++ b/tests/test_logic.py @@ -388,56 +388,79 @@ def test_overlay_symbol_contracts() -> None: and since the overlays were never mounted, nothing had tried the import -- so the boot died on ImportError forty seconds in. """ - owners = {} # dotted module path -> (module dir, source path) - for manifest in sorted(glob.glob( - os.path.join(REPO, "overlay", "modules", "*", "manifest.tsv"))): - moddir = os.path.dirname(manifest) - for raw in open(manifest, encoding="utf-8"): - line = raw.rstrip("\n") - if not line or line.startswith("#"): - continue - source, target = line.split("\t")[:2] - if not source.endswith(".py"): - continue - # Targets are absolute for image-bound overlays and relative to the - # package root for portable ones, so anchoring on "/vllm/" silently - # skipped every portable module -- including moe_gate_sm121, the one - # this check exists for. - if target.startswith("vllm/"): - rel = target - elif "/vllm/" in target: - rel = "vllm/" + target.split("/vllm/", 1)[1] - else: - continue - dotted = rel[:-3].replace("/", ".") - owners[dotted] = os.path.join(moddir, source) + # Owners are resolved per profile: two profiles may overlay the same + # container path with different files (dsv4's mla_indexer and glm53's + # glm53_tail_slot_persistent both own mla/indexer.py), and only the modules + # composed together ever meet at runtime. A global "last manifest wins" map + # checked glm53's attention.py against dsv4's indexer. Modules listed in + # no profile are not checked here. + profiles = {} + for env in sorted(glob.glob(os.path.join(REPO, "profiles", "*.env"))): + m = re.search(r'^MODULES="([^"]*)"', open(env, encoding="utf-8").read(), re.M) + if m: + profiles[os.path.basename(env)] = m.group(1).split() + + def _owners(mods): + owners = {} # dotted module path -> source path + for mod in mods: + manifest = os.path.join(REPO, "overlay", "modules", mod, "manifest.tsv") + moddir = os.path.dirname(manifest) + for raw in open(manifest, encoding="utf-8"): + line = raw.rstrip("\n") + if not line or line.startswith("#"): + continue + source, target = line.split("\t")[:2] + if not source.endswith(".py"): + continue + # Targets are absolute for image-bound overlays and relative to + # the package root for portable ones, so anchoring on "/vllm/" + # silently skipped every portable module -- including + # moe_gate_sm121, the one this check exists for. + if target.startswith("vllm/"): + rel = target + elif "/vllm/" in target: + rel = "vllm/" + target.split("/vllm/", 1)[1] + else: + continue + dotted = rel[:-3].replace("/", ".") + prev = owners.get(dotted) + check(prev is None, f"{mod} and {prev} both own {dotted}") + owners[dotted] = os.path.join(moddir, source) + return owners checked = 0 - for dotted, srcpath in sorted(owners.items()): - provided = set() - tree = ast.parse(open(srcpath, encoding="utf-8").read()) - for node in tree.body: - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): - provided.add(node.name) - elif isinstance(node, ast.Assign): - provided.update(t.id for t in node.targets if isinstance(t, ast.Name)) - elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): - provided.add(node.target.id) - elif isinstance(node, (ast.Import, ast.ImportFrom)): - provided.update((a.asname or a.name).split(".")[0] for a in node.names) - - for other in sorted(set(owners.values())): - if other == srcpath: - continue - for node in ast.walk(ast.parse(open(other, encoding="utf-8").read())): - if not isinstance(node, ast.ImportFrom) or node.module != dotted: + seen = set() + for pname, mods in sorted(profiles.items()): + owners = _owners(mods) + for dotted, srcpath in sorted(owners.items()): + provided = set() + tree = ast.parse(open(srcpath, encoding="utf-8").read()) + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + provided.add(node.name) + elif isinstance(node, ast.Assign): + provided.update(t.id for t in node.targets if isinstance(t, ast.Name)) + elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): + provided.add(node.target.id) + elif isinstance(node, (ast.Import, ast.ImportFrom)): + provided.update((a.asname or a.name).split(".")[0] for a in node.names) + + for other in sorted(set(owners.values())): + if other == srcpath: continue - for alias in node.names: - checked += 1 - check(alias.name in provided, - f"{os.path.basename(other)} imports {alias.name} from " - f"{dotted}, which {os.path.basename(srcpath)} does not define") - print(f" overlay symbol contracts ({checked}) ... OK") + for node in ast.walk(ast.parse(open(other, encoding="utf-8").read())): + if not isinstance(node, ast.ImportFrom) or node.module != dotted: + continue + for alias in node.names: + key = (other, dotted, alias.name, srcpath) + if key in seen: + continue + seen.add(key) + checked += 1 + check(alias.name in provided, + f"[{pname}] {os.path.basename(other)} imports {alias.name} from " + f"{dotted}, which {os.path.basename(srcpath)} does not define") + print(f" overlay symbol contracts ({checked}, per profile) ... OK") def test_profile_env_carried() -> None: @@ -3969,6 +3992,7 @@ def original(self, *args): "_glm53_cache_only_indexer_contract", "_glm53_cache_only_indexer_forward", "_glm53_fused_indexer_forward", + "_glm53_head_gate", "_HEAD_GATE", "install_glm53_prefill_fastpath", "prepare_glm53_prefill_fastpath", }, @@ -7038,6 +7062,72 @@ def test_glm53_async_dflash_contracts() -> None: print(" glm53 async dflash contracts .. OK") +def test_glm53_indexer_gate_splitk_contracts() -> None: + """glm53_indexer_gate_splitk: opt-in split-K head gate, small-M only, stock default.""" + mod_dir = os.path.join(REPO, "overlay", "modules", "glm53_indexer_gate_splitk") + kern = open(os.path.join(mod_dir, "glm53_indexer_gate.py"), encoding="utf-8").read() + attn = open(os.path.join(mod_dir, "glm5next_attention.py"), encoding="utf-8").read() + fast = open(os.path.join(REPO, "overlay", "modules", "glm53_model_wiring", + "glm53_prefill_fastpath.py"), encoding="utf-8").read() + profile = open(os.path.join(REPO, "profiles", "glm53.env"), encoding="utf-8").read() + modules = re.search(r'^MODULES="([^"]+)"', profile, re.M).group(1).split() + check("glm53_indexer_gate_splitk" in modules, "glm53 profile must mount glm53_indexer_gate_splitk") + check(re.search(r"^VLLM_GLM53_INDEXER_GATE_SPLITK=0$", profile, re.M) is not None, + "profile must ship VLLM_GLM53_INDEXER_GATE_SPLITK=0 (stock torch.mm by default)") + rows = [l.split("\t") for l in open(os.path.join(mod_dir, "manifest.tsv"), encoding="utf-8") + .read().splitlines() if l and not l.startswith("#")] + by_target = {r[1]: r[2] for r in rows} + check(len(rows) == 2 and re.fullmatch( + r"[0-9a-f]{64}", by_target.get("vllm/models/glm5next/nvidia/attention.py", "")) is not None + and by_target.get("vllm/models/glm5next/nvidia/glm53_indexer_gate.py") == "absent", + f"manifest must overlay attention.py (pinned) and add the kernel file (absent): {rows}") + stock = "torch.mm(hidden_states.float(), self._wp_fp32)" + helper = "_glm53_head_gate(hidden_states, self._wp_fp32)" + check(stock not in attn and attn.count(helper) == 1, + "attention overlay routes the head gate through the helper exactly once") + check("from vllm.models.glm5next.nvidia.glm53_indexer_gate import head_gate as _glm53_head_gate" in attn, + "attention overlay imports the helper from the module's kernel file") + check(stock not in fast and fast.count(helper) == 1, + "the fused-indexer forward (VLLM_GLM53_FUSED_K_GATE) routes through the same helper") + check("except ImportError:" in fast and "return torch.mm(x.float(), w)" in fast, + "the fastpath falls back to stock torch.mm when the module is not mounted") + check("MAX_M = 16" in kern and "x.shape[0] <= MAX_M" in kern, + "split-K only for the small-M decode shape (cuBLAS is already fast at M=32)") + check("return torch.mm(x.float(), w)" in kern, "the helper itself falls back to the stock product") + check("torch.zeros(M, N" in kern and "tl.atomic_add" in kern, + "atomic split-K must reduce into a zeroed output") + check("w.shape[0] % (_SPLIT * _BLOCK_K) == 0" in kern + and "w.dtype == torch.float32" in kern and "w.is_contiguous()" in kern, + "applicability: K tiles the split, fp32 contiguous weight") + check(kern.count("do_not_specialize") == 0 or "K" in kern, + "integer args do not need specialization pins here (K, N are shape constants per layer)") + tree = ast.parse(kern) + nodes = [n for n in tree.body + if (isinstance(n, ast.FunctionDef) and n.name == "gate_splitk_enabled") + or (isinstance(n, ast.Assign) and any(isinstance(t, ast.Name) and t.id == "ENV" for t in n.targets))] + check(len(nodes) == 2, "ENV constant and gate_splitk_enabled must be top-level definitions") + ns: dict = {"os": os} + exec(compile(ast.Module(body=nodes, type_ignores=[]), "glm53_indexer_gate", "exec"), ns) + fn = ns["gate_splitk_enabled"] + saved = os.environ.pop("VLLM_GLM53_INDEXER_GATE_SPLITK", None) + try: + check(fn() is False, "unset knob keeps stock torch.mm") + for v, want in (("0", False), ("1", True), (" 1 ", True), ("on", False), ("true", False), + ("2", False), ("shadow", False), ("", False)): + os.environ["VLLM_GLM53_INDEXER_GATE_SPLITK"] = v + check(fn() is want, f"knob {v!r} must map to {want} (only the exact string 1 arms)") + finally: + os.environ.pop("VLLM_GLM53_INDEXER_GATE_SPLITK", None) + if saved is not None: + os.environ["VLLM_GLM53_INDEXER_GATE_SPLITK"] = saved + readme = open(os.path.join(mod_dir, "README.md"), encoding="utf-8").read() + check("not bit-exact" in readme and "VLLM_GLM53_INDEXER_GATE_SPLITK=1 bash launchers/" in readme, + "README must state the numerics caveat and the caller-env arming form") + req = open(os.path.join(mod_dir, "requires"), encoding="utf-8").read().split() + check("glm53_model_wiring" in req, "the fastpath hunk lives in glm53_model_wiring") + print(" glm53 indexer gate split-K contracts .. OK") + + if __name__ == "__main__": test_skip_topk() test_prefill_chunker() @@ -7113,4 +7203,5 @@ def test_glm53_async_dflash_contracts() -> None: test_glm53_prep_fused_contracts() test_profile_keys_not_passed_via_extra_env() test_glm53_async_dflash_contracts() + test_glm53_indexer_gate_splitk_contracts() print(f"all OK ({PASS} checks)")